﻿import functools
import operator
import math
import string
import struct
import sys

__title__ = 'puzpy'
__version__ = '0.2.6'
__author__ = 'Alex DeJarnatt'
__author_email__ = 'adejarnatt@gmail.com'
__maintainer__ = 'Simeon Visser'
__maintainer_email__ = 'simeonvisser@gmail.com'
__license__ = 'MIT'
__copyright__ = 'Copyright 2009 Alex DeJarnatt'


PY3 = sys.version_info[0] >= 3

if PY3:
    str = str
    range = range
else:
    str = unicode  # noqa: F821
    range = xrange  # noqa: F821

HEADER_FORMAT = '''<
             H 11s        xH
             Q       4s  2sH
             12s         BBH
             H H '''

HEADER_CKSUM_FORMAT = '<BBH H H '

EXTENSION_HEADER_FORMAT = '< 4s  H H '

MASKSTRING = 'ICHEATED'

ENCODING = 'ISO-8859-1'
ENCODING_UTF8 = 'UTF-8'
ENCODING_ERRORS = 'strict'  # raises an exception for bad chars; change to 'replace' for laxer handling

ACROSSDOWN = b'ACROSS&DOWN'

BLACKSQUARE = '.'
BLACKSQUARE2 = ':'
BLANKSQUARE = '-'


def enum(**enums):
    return type('Enum', (), enums)


PuzzleType = enum(
    Normal=0x0001,
    Diagramless=0x0401
)

# the following diverges from the documentation
# but works for the files I've tested
SolutionState = enum(
    # solution is available in plaintext
    Unlocked=0x0000,
    # solution is not present in the file
    NotProvided=0x0002,
    # solution is locked (scrambled) with a key
    Locked=0x0004
)

GridMarkup = enum(
    # ordinary grid cell
    Default=0x00,
    # marked incorrect at some point
    PreviouslyIncorrect=0x10,
    # currently showing incorrect
    Incorrect=0x20,
    # user got a hint
    Revealed=0x40,
    # circled
    Circled=0x80
)


# refer to Extensions as Extensions.Rebus, Extensions.Markup
Extensions = enum(
    # grid of rebus indices: 0 for non-rebus;
    # i+1 for key i into RebusSolutions map
    # should be same size as the grid
    Rebus=b'GRBS',

    # map of rebus solution entries eg 0:HEART;1:DIAMOND;17:CLUB;23:SPADE;
    RebusSolutions=b'RTBL',

    # user's rebus entries, same format as RebusSolutions
    RebusFill=b'RUSR',

    # timer state: 'a,b' where a is the number of seconds elapsed and
    # b is a boolean (0,1) for whether the timer is running
    Timer=b'LTIM',

    # grid cell markup: previously incorrect: 0x10;
    # currently incorrect: 0x20,
    # hinted: 0x40,
    # circled: 0x80
    Markup=b'GEXT'
)


def read(filename):
    """
    Read a .puz file and return the Puzzle object.
    raises PuzzleFormatError if there's any problem with the file format.
    """
    with open(filename, 'rb') as f:
        return load(f.read())


def read_text(filename):
    """
    Read an Across Lite .txt text format file and return the Puzzle object.
    raises PuzzleFormatError if there's any problem with the file format.
    """
    with open(filename, 'r', encoding='utf-8', errors='replace') as f:
        return load_text(f.read())
    

def load(data):
    """
    Read .puz file data and return the Puzzle object.
    raises PuzzleFormatError if there's any problem with the file format.
    """
    puz = Puzzle()
    puz.load(data)
    return puz


def load_text(text):
    """
    Parse Across Lite Text format from a string or bytes and return a Puzzle object.
    raises PuzzleFormatError if there's any problem with the format.
    """
    return from_text_format(text)


class PuzzleFormatError(Exception):
    """
    Indicates a format error in the .puz file. May be thrown due to
    invalid headers, invalid checksum validation, or other format issues.
    """
    def __init__(self, message=''):
        self.message = message


class Puzzle:
    """Represents a puzzle
    """
    def __init__(self, version='1.3'):
        """Initializes a blank puzzle
        """
        self.preamble = b''
        self.postscript = b''
        self.title = ''
        self.author = ''
        self.copyright = ''
        self.width = 0
        self.height = 0
        if not isinstance(version, bytes):
            version = bytes(version.encode('utf8'))
        self.version = version
        self.fileversion = version + b'\0'
        self.encoding = ENCODING
        # these are bytes that might be unused
        self.unk1 = b'\0' * 2
        self.unk2 = b'\0' * 12
        self.scrambled_cksum = 0
        self.fill = ''
        self.solution = ''
        self.clues = []
        self.notes = ''
        self.extensions = {}
        # the folowing is so that we can round-trip values in order:
        self._extensions_order = []
        self.puzzletype = PuzzleType.Normal
        self.solution_state = SolutionState.Unlocked
        self.helpers = {}  # add-ons like Rebus and Markup

    def load(self, data):
        s = PuzzleBuffer(data)

        # advance to start - files may contain some data before the
        # start of the puzzle use the ACROSS&DOWN magic string as a waypoint
        # save the preamble for round-tripping
        if not s.seek_to(ACROSSDOWN, -2):
            raise PuzzleFormatError("Data does not appear to represent a "
                                    "puzzle. Are you sure you didn't intend "
                                    "to use read?")

        self.preamble = s.data[:s.pos]

        puzzle_data = s.unpack(HEADER_FORMAT)
        cksum_gbl = puzzle_data[0]
        # acrossDown = puzzle_data[1]
        cksum_hdr = puzzle_data[2]
        cksum_magic = puzzle_data[3]
        self.fileversion = puzzle_data[4]
        # since we don't know the role of these bytes, just round-trip them
        self.unk1 = puzzle_data[5]
        self.scrambled_cksum = puzzle_data[6]
        self.unk2 = puzzle_data[7]
        self.width = puzzle_data[8]
        self.height = puzzle_data[9]
        numclues = puzzle_data[10]
        self.puzzletype = puzzle_data[11]
        self.solution_state = puzzle_data[12]

        self.version = self.fileversion[:3]
        # Once we have fileversion we can guess the encoding
        self.encoding = ENCODING if self.version_tuple()[0] < 2 else ENCODING_UTF8
        s.encoding = self.encoding

        self.solution = s.read(self.width * self.height).decode(self.encoding)
        self.fill = s.read(self.width * self.height).decode(self.encoding)

        self.title = s.read_string()
        self.author = s.read_string()
        self.copyright = s.read_string()

        self.clues = [s.read_string() for i in range(0, numclues)]
        self.notes = s.read_string()

        ext_cksum = {}
        while s.can_unpack(EXTENSION_HEADER_FORMAT):
            code, length, cksum = s.unpack(EXTENSION_HEADER_FORMAT)
            ext_cksum[code] = cksum
            # extension data is represented as a null-terminated string,
            # but since the data can contain nulls we can't use read_string
            self.extensions[code] = s.read(length)
            s.read(1)  # extensions have a trailing byte
            # save the codes in order for round-tripping
            self._extensions_order.append(code)

        # sometimes there's some extra garbage at
        # the end of the file, usually \r\n
        if s.can_read():
            self.postscript = s.read_to_end()

        if cksum_gbl != self.global_cksum():
            raise PuzzleFormatError('global checksum does not match')
        if cksum_hdr != self.header_cksum():
            raise PuzzleFormatError('header checksum does not match')
        if cksum_magic != self.magic_cksum():
            raise PuzzleFormatError('magic checksum does not match')
        for code, cksum_ext in ext_cksum.items():
            if cksum_ext != data_cksum(self.extensions[code]):
                raise PuzzleFormatError(
                    'extension %s checksum does not match' % code
                )

    def save(self, filename):
        puzzle_bytes = self.tobytes()
        with open(filename, 'wb') as f:
            f.write(puzzle_bytes)

    def tobytes(self):
        s = PuzzleBuffer(encoding=self.encoding)
        # commit any changes from helpers
        for h in self.helpers.values():
            if 'save' in dir(h):
                h.save()

        # include any preamble text we might have found on read
        s.write(self.preamble)

        s.pack(HEADER_FORMAT,
               self.global_cksum(), ACROSSDOWN,
               self.header_cksum(), self.magic_cksum(),
               self.fileversion, self.unk1, self.scrambled_cksum,
               self.unk2, self.width, self.height,
               len(self.clues), self.puzzletype, self.solution_state)

        s.write(self.encode(self.solution))
        s.write(self.encode(self.fill))

        s.write_string(self.title)
        s.write_string(self.author)
        s.write_string(self.copyright)

        for clue in self.clues:
            s.write_string(clue)

        s.write_string(self.notes)

        # do a bit of extra work here to ensure extensions round-trip in the
        # order they were read. this makes verification easier. But allow
        # for the possibility that extensions were added or removed from
        # self.extensions
        ext = dict(self.extensions)
        for code in self._extensions_order:
            data = ext.pop(code, None)
            if data:
                s.pack(EXTENSION_HEADER_FORMAT, code,
                       len(data), data_cksum(data))
                s.write(data + b'\0')

        for code, data in ext.items():
            s.pack(EXTENSION_HEADER_FORMAT, code, len(data), data_cksum(data))
            s.write(data + b'\0')

        # postscript is initialized, read, and stored as bytes. In case it is
        # overwritten as a string, this try/except converts it back.
        try:
            postscript_bytes = self.encode(self.postscript)
        except AttributeError:
            postscript_bytes = self.postscript

        s.write(postscript_bytes)

        return s.tobytes()

    def encode(self, s):
        return s.encode(self.encoding, ENCODING_ERRORS)

    def encode_zstring(self, s):
        return self.encode(s) + b'\0'

    def version_tuple(self):
        return tuple(map(int, self.version.split(b'.')))

    def set_version(self, version):
        if not isinstance(version, bytes):
            version = bytes(version.encode('utf-8'))
        self.version = version
        self.fileversion = version + b'\0'

    def has_rebus(self):
        return self.rebus().has_rebus()

    def rebus(self):
        if 'rebus' not in self.helpers:
            self.helpers['rebus'] = Rebus(self)
        return self.helpers['rebus']

    def has_markup(self):
        return self.markup().has_markup()

    def markup(self):
        if 'markup' not in self.helpers:
            self.helpers['markup'] = Markup(self)
        return self.helpers['markup']

    def clue_numbering(self):
        if 'clues' not in self.helpers:
            self.helpers['clues'] = DefaultClueNumbering(self.fill, self.clues, self.width, self.height)
        return self.helpers['clues']

    def blacksquare(self):
        return BLACKSQUARE2 if self.puzzletype == PuzzleType.Diagramless else BLACKSQUARE

    def is_solution_locked(self):
        return bool(self.solution_state == SolutionState.Locked)

    def unlock_solution(self, key):
        if self.is_solution_locked():
            unscrambled = unscramble_solution(self.solution, self.width, self.height, key,
                                              ignore_chars=self.blacksquare())
            if not self.check_answers(unscrambled):
                return False

            # clear the scrambled bit and cksum
            self.solution = unscrambled
            self.scrambled_cksum = 0
            self.solution_state = SolutionState.Unlocked

        return True

    def lock_solution(self, key):
        if not self.is_solution_locked():
            # set the scrambled bit and cksum
            self.scrambled_cksum = scrambled_cksum(self.solution, self.width, self.height,
                                                   ignore_chars=self.blacksquare(), encoding=self.encoding)
            self.solution_state = SolutionState.Locked
            scrambled = scramble_solution(self.solution, self.width, self.height, key,
                                          ignore_chars=self.blacksquare())
            self.solution = scrambled

    def check_answers(self, fill):
        if self.is_solution_locked():
            scrambled = scrambled_cksum(fill, self.width, self.height,
                                        ignore_chars=self.blacksquare(), encoding=self.encoding)
            return scrambled == self.scrambled_cksum
        else:
            return fill == self.solution

    def header_cksum(self, cksum=0):
        return data_cksum(struct.pack(HEADER_CKSUM_FORMAT,
                          self.width, self.height, len(self.clues),
                          self.puzzletype, self.solution_state), cksum)

    def text_cksum(self, cksum=0):
        # for the checksum to work these fields must be added in order with
        # null termination, followed by all non-empty clues without null
        # termination, followed by notes (but only for version >= 1.3)
        if self.title:
            cksum = data_cksum(self.encode_zstring(self.title), cksum)
        if self.author:
            cksum = data_cksum(self.encode_zstring(self.author), cksum)
        if self.copyright:
            cksum = data_cksum(self.encode_zstring(self.copyright), cksum)

        for clue in self.clues:
            if clue:
                cksum = data_cksum(self.encode(clue), cksum)

        # notes included in global cksum starting v1.3 of format
        if self.version_tuple() >= (1, 3) and self.notes:
            cksum = data_cksum(self.encode_zstring(self.notes), cksum)

        return cksum

    def global_cksum(self):
        cksum = self.header_cksum()
        cksum = data_cksum(self.encode(self.solution), cksum)
        cksum = data_cksum(self.encode(self.fill), cksum)
        cksum = self.text_cksum(cksum)
        # extensions do not seem to be included in global cksum
        return cksum

    def magic_cksum(self):
        cksums = [
            self.header_cksum(),
            data_cksum(self.encode(self.solution)),
            data_cksum(self.encode(self.fill)),
            self.text_cksum()
        ]

        cksum_magic = 0
        for (i, cksum) in enumerate(reversed(cksums)):
            cksum_magic <<= 8
            cksum_magic |= (
                ord(MASKSTRING[len(cksums) - i - 1]) ^ (cksum & 0x00ff)
            )
            cksum_magic |= (
                (ord(MASKSTRING[len(cksums) - i - 1 + 4]) ^ (cksum >> 8)) << 32
            )

        return cksum_magic


class PuzzleBuffer:
    """PuzzleBuffer class
    wraps a data buffer ('' or []) and provides .puz-specific methods for
    reading and writing data
    """
    def __init__(self, data=None, encoding=ENCODING):
        self.data = data or []
        self.encoding = encoding
        self.pos = 0

    def can_read(self, n_bytes=1):
        return self.pos + n_bytes <= len(self.data)

    def length(self):
        return len(self.data)

    def read(self, n_bytes):
        start = self.pos
        self.pos += n_bytes
        return self.data[start:self.pos]

    def read_to_end(self):
        start = self.pos
        self.pos = self.length()
        return self.data[start:self.pos]

    def read_string(self):
        return self.read_until(b'\0')

    def read_until(self, c):
        start = self.pos
        self.seek_to(c, 1)  # read past
        return str(self.data[start:self.pos-1], self.encoding)

    def seek_to(self, s, offset=0):
        try:
            self.pos = self.data.index(s, self.pos) + offset
            return True
        except ValueError:
            # s not found, advance to end
            self.pos = self.length()
            return False

    def write(self, s):
        self.data.append(s)

    def write_string(self, s):
        s = s or ''
        self.data.append(s.encode(self.encoding, ENCODING_ERRORS) + b'\0')

    def pack(self, struct_format, *values):
        self.data.append(struct.pack(struct_format, *values))

    def can_unpack(self, struct_format):
        return self.can_read(struct.calcsize(struct_format))

    def unpack(self, struct_format):
        start = self.pos
        try:
            res = struct.unpack_from(struct_format, self.data, self.pos)
            self.pos += struct.calcsize(struct_format)
            return res
        except struct.error:
            message = 'could not unpack values at {} for format {}'.format(
                start, struct_format
            )
            raise PuzzleFormatError(message)

    def tobytes(self):
        return b''.join(self.data)


# clue numbering helper

def get_grid_numbering(grid, width, height):
    # Add numbers to the grid based on positions of black squares
    def col(index):
        return index % width

    def row(index):
        return int(math.floor(index / width))
    
    def len_across(index):
        for c in range(0, width - col(index)):
            if is_blacksquare(grid[index + c]):
                return c
        return c + 1

    def len_down(index):
        for c in range(0, height - row(index)):
            if is_blacksquare(grid[index + c*width]):
                return c
        return c + 1

    a = []
    d = []
    c = 0
    n = 1
    for i in range(0, len(grid)):
        if not is_blacksquare(grid[i]):
            lastc = c
            is_across = col(i) == 0 or is_blacksquare(grid[i - 1])
            if is_across and len_across(i) > 1:
                a.append({
                    'num': n,
                    'clue': None,  # filled in by caller
                    'clue_index': c,
                    'cell': i,
                    'row': row(i),
                    'col': col(i),
                    'len': len_across(i),
                    'dir': 'across',
                })
                c += 1
            is_down = row(i) == 0 or is_blacksquare(grid[i - width])
            if is_down and len_down(i) > 1:
                d.append({
                    'num': n,
                    'clue': None,  # filled in by caller
                    'clue_index': c,
                    'cell': i,
                    'row': row(i),
                    'col': col(i),
                    'len': len_down(i),
                    'dir': 'down'
                })
                c += 1
            if c > lastc:
                n += 1
    return a, d

class DefaultClueNumbering:
    def __init__(self, grid, clues, width, height):
        self.grid = grid
        self.clues = clues
        self.width = width
        self.height = height

        self.across, self.down = get_grid_numbering(grid, width, height)
        for entry in self.across:
            entry['clue'] = clues[entry['clue_index']]
        for entry in self.down:
            entry['clue'] = clues[entry['clue_index']]

    # The following methods are no longer in use, but left here in case
    # anyone was using them externally. They may be removed in a future release.
    def col(self, index):
        return index % self.width

    def row(self, index):
        return int(math.floor(index / self.width))

    def len_across(self, index):
        for c in range(0, self.width - self.col(index)):
            if is_blacksquare(self.grid[index + c]):
                return c
        return c + 1

    def len_down(self, index):
        for c in range(0, self.height - self.row(index)):
            if is_blacksquare(self.grid[index + c*self.width]):
                return c
        return c + 1


class Grid:
    def __init__(self, grid, width, height):
        self.grid = grid
        self.width = width
        self.height = height
        assert len(self.grid) == self.width * self.height

    def get_cell(self, row, col):
        return self.grid[self.get_cell_index(row, col)]

    def get_cell_index(self, row, col):
        return row * self.width + col

    def get_range(self, row, col, length, dir='across'):
        if dir == 'across':
            return self.get_range_across(row, col, length)
        elif dir == 'down':
            return self.get_range_down(row, col, length)
        else:
            assert False, "dir not one of 'across' or 'down'"

    def get_range_across(self, row, col, length):
        return [self.grid[self.get_cell_index(row, col + i)] for i in range(length)]

    def get_range_down(self, row, col, length):
        return [self.grid[self.get_cell_index(row + i, col)] for i in range(length)]

    def get_range_for_clue(self, clue):
        return self.get_range(clue['row'], clue['col'], clue['len'], clue['dir'])

    def get_row(self, row):
        return self.get_range_across(row, 0, self.width)

    def get_column(self, col):
        return self.get_range_down(0, col, self.height)

    def get_string(self, row, col, length, dir='across'):
        return ''.join(self.get_range(row, col, length, dir))

    def get_string_across(self, row, col, length):
        return ''.join(self.get_range_across(row, col, length))

    def get_string_down(self, row, col, length):
        return ''.join(self.get_range_down(row, col, length))

    def get_string_for_clue(self, clue):
        return ''.join(self.get_range_for_clue(clue))


class Rebus:
    def __init__(self, puzzle):
        self.puzzle = puzzle
        self.fill = {}
        self.solutions = {}
        self.table = [0] * (self.puzzle.width * self.puzzle.height)

        # parse rebus data
        if Extensions.Rebus in self.puzzle.extensions:
            rebus_data = self.puzzle.extensions[Extensions.Rebus]
            self.table = parse_bytes(rebus_data)
        
        if Extensions.RebusSolutions in self.puzzle.extensions:
            raw_solution_data = self.puzzle.extensions[Extensions.RebusSolutions]
            solutions_str = raw_solution_data.decode(puzzle.encoding)
            self.solutions = {
               int(item[0]): item[1]
                for item in parse_dict(solutions_str).items()
            }
        
        if Extensions.RebusFill in self.puzzle.extensions:
            raw_fill_data = self.puzzle.extensions[Extensions.RebusFill]
            fill_str = raw_fill_data.decode(puzzle.encoding)
            self.fill = {
                int(item[0]): item[1]
                for item in parse_dict(fill_str).items()
            }

    def has_rebus(self):
        return Extensions.Rebus in self.puzzle.extensions or any(self.table)

    def is_rebus_square(self, index):
        return bool(self.table[index])

    def get_rebus_squares(self):
        return [i for i, b in enumerate(self.table) if b]

    def get_rebus_solution(self, index):
        if self.is_rebus_square(index):
            return self.solutions[self.table[index] - 1]
        return None

    def get_rebus_fill(self, index):
        if self.is_rebus_square(index):
            return self.fill[self.table[index] - 1]
        return None

    def set_rebus_fill(self, index, value):
        if self.is_rebus_square(index):
            self.fill[self.table[index] - 1] = value

    def save(self):
        if self.has_rebus():
            # commit changes back to puzzle.extensions
            self.puzzle.extensions[Extensions.Rebus] = pack_bytes(self.table)
            rebus_solutions = self.puzzle.encode(dict_to_string(self.solutions))
            self.puzzle.extensions[Extensions.RebusSolutions] = rebus_solutions
            rebus_fill = self.puzzle.encode(dict_to_string(self.fill))
            self.puzzle.extensions[Extensions.RebusFill] = rebus_fill


class Markup:
    def __init__(self, puzzle):
        self.puzzle = puzzle
        # parse markup data
        markup_data = self.puzzle.extensions.get(Extensions.Markup, b'')
        self.markup = parse_bytes(markup_data)

    def has_markup(self):
        return any(bool(b) for b in self.markup)

    def get_markup_squares(self):
        return [i for i, b in enumerate(self.markup) if b]

    def is_markup_square(self, index):
        return bool(self.markup[index])

    def save(self):
        if self.has_markup():
            self.puzzle.extensions[Extensions.Markup] = pack_bytes(self.markup)


# helper functions for cksums and scrambling
def data_cksum(data, cksum=0):
    for b in data:
        if isinstance(b, bytes):
            b = ord(b)
        # right-shift one with wrap-around
        lowbit = (cksum & 0x0001)
        cksum = (cksum >> 1)
        if lowbit:
            cksum = (cksum | 0x8000)

        # then add in the data and clear any carried bit past 16
        cksum = (cksum + b) & 0xffff

    return cksum


def replace_chars(s, chars, replacement=''):
    for ch in chars:
        s = s.replace(ch, replacement)
    return s


def scramble_solution(solution, width, height, key, ignore_chars=BLACKSQUARE):
    sq = square(solution, width, height)
    data = restore(sq, scramble_string(replace_chars(sq, ignore_chars), key))
    return square(data, height, width)


def scramble_string(s, key):
    """
    s is the puzzle's solution in column-major order, omitting black squares:
    i.e. if the puzzle is:
        C A T
        # # A
        # # R
    solution is CATAR


    Key is a 4-digit number in the range 1000 <= key <= 9999

    """
    key = key_digits(key)
    for k in key:          # foreach digit in the key
        s = shift(s, key)  # for each char by each digit in the key in sequence
        s = s[k:] + s[:k]  # cut the sequence around the key digit
        s = shuffle(s)     # do a 1:1 shuffle of the 'deck'

    return s


def unscramble_solution(scrambled, width, height, key, ignore_chars=BLACKSQUARE):
    # width and height are reversed here
    sq = square(scrambled, width, height)
    data = restore(sq, unscramble_string(replace_chars(sq, ignore_chars), key))
    return square(data, height, width)


def unscramble_string(s, key):
    key = key_digits(key)
    l = len(s)  # noqa: E741
    for k in key[::-1]:
        s = unshuffle(s)
        s = s[l-k:] + s[:l-k]
        s = unshift(s, key)

    return s


def scrambled_cksum(scrambled, width, height, ignore_chars=BLACKSQUARE, encoding=ENCODING):
    data = replace_chars(square(scrambled, width, height), ignore_chars)
    return data_cksum(data.encode(encoding, ENCODING_ERRORS))


def key_digits(key):
    return [int(c) for c in str(key).zfill(4)]


def square(data, w, h):
    aa = [data[i:i+w] for i in range(0, len(data), w)]
    return ''.join(
        [''.join([aa[r][c] for r in range(0, h)]) for c in range(0, w)]
    )


def shift(s, key):
    atoz = string.ascii_uppercase
    return ''.join(
        atoz[(atoz.index(c) + key[i % len(key)]) % len(atoz)]
        for i, c in enumerate(s)
    )


def unshift(s, key):
    return shift(s, [-k for k in key])


def shuffle(s):
    mid = int(math.floor(len(s) / 2))
    items = functools.reduce(operator.add, zip(s[mid:], s[:mid]))
    return ''.join(items) + (s[-1] if len(s) % 2 else '')


def unshuffle(s):
    return s[1::2] + s[::2]


def restore(s, t):
    """
    s is the source string, it can contain '.'
    t is the target, it's smaller than s by the number of '.'s in s

    Each char in s is replaced by the corresponding
    char in t, jumping over '.'s in s.

    >>> restore('ABC.DEF', 'XYZABC')
    'XYZ.ABC'
    """
    t = (c for c in t)
    return ''.join(next(t) if not is_blacksquare(c) else c for c in s)


def is_blacksquare(c):
    if isinstance(c, int):
        c = chr(c)
    return c in [BLACKSQUARE, BLACKSQUARE2]


#
# functions for parsing / serializing primitives
#


def parse_bytes(s):
    return list(struct.unpack('B' * len(s), s))


def pack_bytes(a):
    return struct.pack('B' * len(a), *a)


# dict string format is k1:v1;k2:v2;...;kn:vn;
# (for whatever reason there's a trailing ';')
def parse_dict(s):
    return dict(p.split(':', 1) for p in s.split(';') if ':' in p)


def dict_to_string(d):
    return ';'.join(':'.join(map(str, [k, v])) for k, v in d.items()) + ';'


def from_text_format(s):
    d = text_file_as_dict(s)

    if 'ACROSS PUZZLE' in d:
        file_version = 'v1'
    elif 'ACROSS PUZZLE v2' in d:
        file_version = 'v2'
    else:
        raise PuzzleFormatError('Not a valid Across Lite text puzzle')

    p = Puzzle()
    across_clues = []
    down_clues = []
    if 'TITLE' in d:
        p.title = d['TITLE']
    if 'AUTHOR' in d:
        p.author = d['AUTHOR']
    if 'COPYRIGHT' in d:
        p.copyright = d['COPYRIGHT']
    if 'SIZE' in d:
        size = d['SIZE'].split('x')
        p.width = int(size[0])
        p.height = int(size[1])
    if 'GRID' in d:
        solution_lines = d['GRID'].splitlines()
        p.solution = ''.join(line.strip() for line in solution_lines if line.strip())
    if 'ACROSS' in d:
        across_clues.extend(line.strip() for line in d['ACROSS'].splitlines() if line.strip())
    if 'DOWN' in d:
        down_clues.extend(line.strip() for line in d['DOWN'].splitlines() if line.strip())
    if 'NOTEPAD' in d:
        p.notes = d['NOTEPAD']
    if 'REBUS' in d:
        pass # TODO: text file REBUS
    
    if p.solution:
        p.fill = ''.join(c if c == BLACKSQUARE else BLANKSQUARE for c in p.solution)
        across, down = get_grid_numbering(p.fill, p.width, p.height)
        # we have to match puzfile's expected clue ordering or we won't be able to 
        # write the puzzle out as a valid .puz file
        p.clues = [''] * (len(across) + len(down))
        for i in range(len(across)):
            clue = across_clues[i] if i < len(across_clues) else ''
            across[i]['clue'] = clue
            p.clues[across[i]['clue_index']] = clue
        for i in range(len(down)):
            clue = down_clues[i] if i < len(down_clues) else ''
            down[i]['clue'] = clue
            p.clues[down[i]['clue_index']] = clue

    return p

def text_file_as_dict(s):
    d = {}
    k = None
    v = []
    for line in s.splitlines():
        line = line.strip()
        if line.startswith('<') and line.endswith('>'):
            if k:
                d[k] = '\n'.join(v)
            k = line[1:-1]
            v = []
        else:
            v.append(line)

    if k:
        d[k] = '\n'.join(v)
    return d

def to_text_format(p, text_version='v1'):
    TAB = '\t'  # most lines begin indented with whitespace
    lines = []
    if text_version == 'v1':
        lines.append('<ACROSS PUZZLE>')
    elif text_version:
        lines.append(f'<ACROSS PUZZLE {text_version}>')
    else:
        raise ValueError("invalid text_version")

    lines.append('<TITLE>')
    lines.append(TAB + p.title)
    lines.append('<AUTHOR>')
    lines.append(TAB + p.author)
    lines.append('<COPYRIGHT>')
    lines.append(TAB + p.copyright)
    lines.append('<SIZE>')
    lines.append(TAB + f'{p.width}x{p.height}')
    lines.append('<GRID>')
    for r in range(p.height):
        row = p.solution[r*p.width:(r+1)*p.width]
        lines.append(TAB + row)

    # get clues in across/down order
    numbering = p.clue_numbering()
    lines.append('<ACROSS>')
    for clue in numbering.across:
        lines.append(TAB + clue['clue'])
    lines.append('<DOWN>')
    for clue in numbering.down:
        lines.append(TAB + clue['clue'])

    lines.append('<NOTEPAD>')
    lines.append(p.notes)  # no tab here, idk why

    return '\n'.join(lines)