#
# A class which knows how to write and parse 'Quetzal' files, which is
# the standard save-file format for modern Z-machine implementations.
# This allows ZVM's saved games to load in other interpreters, and
# vice versa.
#
# The Quetzal format is documented at:
#    http://www.ifarchive.org/if-archive/infocom/\
#           interpreters/specification/savefile_14.txt
#
# For the license of this file, please consult the LICENSE file in the
# root directory of this distribution.
#

import os
import struct

from . import bitfield, zstackmanager
from .zlogging import log


def _read_iff_chunk(f):
    """Read one IFF chunk from a file-like object. Returns (name, data).

    Raises EOFError when there are no more chunks to read.
    """
    header = f.read(8)
    if len(header) < 8:
        raise EOFError
    name = header[:4]
    (size,) = struct.unpack(">I", header[4:8])
    data = f.read(size)
    # IFF chunks are padded to even byte boundaries
    if size % 2 == 1:
        f.read(1)
    return name, data


# The general format of Queztal is that of a "FORM" IFF file, which is
# a container class for 'chunks'.
#
#   "FORM", 4 bytes of container total-length, "IFZS",
#        4-byte chunkname, 4-byte length, length bytes of data
#        4-byte chunkname, 4-byte length, length bytes of data
#        4-byte chunkname, 4-byte length, length bytes of data
#        ...


class QuetzalError(Exception):
    "General exception for Quetzal classes."

    pass


class QuetzalMalformedChunk(QuetzalError):
    "Malformed chunk detected."


class QuetzalNoSuchSavefile(QuetzalError):
    "Cannot locate save-game file."


class QuetzalUnrecognizedFileFormat(QuetzalError):
    "Not a valid Quetzal file."


class QuetzalIllegalChunkOrder(QuetzalError):
    "IFhd chunk came after Umem/Cmem/Stks chunks (see section 5.4)."


class QuetzalMismatchedFile(QuetzalError):
    "Quetzal file dosen't match current game."


class QuetzalMemoryOutOfBounds(QuetzalError):
    "Decompressed dynamic memory has gone out of bounds."


class QuetzalMemoryMismatch(QuetzalError):
    "Savefile's dynamic memory image is incorrectly sized."


class QuetzalStackFrameOverflow(QuetzalError):
    "Stack frame parsing went beyond bounds of 'Stks' chunk."


class QuetzalParser:
    """A class to read a Quetzal save-file and modify a z-machine."""

    def __init__(self, zmachine):
        log("Creating new instance of QuetzalParser")
        self._zmachine = zmachine
        self._seen_mem_or_stks = False
        self._last_loaded_metadata = {}  # metadata for tests & debugging

    def _parse_ifhd(self, data):
        """Parse a chunk of type IFhd, and check that the quetzal file
        really belongs to the current story (by comparing release number,
        serial number, and checksum.)"""

        # Spec says that this chunk *must* come before memory or stack chunks.
        if self._seen_mem_or_stks:
            raise QuetzalIllegalChunkOrder

        bytes = data
        if len(bytes) != 13:
            raise QuetzalMalformedChunk

        chunk_release = (data[0] << 8) + data[1]
        chunk_serial = data[2:8]
        chunk_checksum = (data[8] << 8) + data[9]
        chunk_pc = (data[10] << 16) + (data[11] << 8) + data[12]
        self._zmachine._opdecoder.program_counter = chunk_pc

        log(f"  Found release number {chunk_release}")
        log(f"  Found serial number {int(chunk_serial)}")
        log(f"  Found checksum {chunk_checksum}")
        log(f"  Initial program counter value is {chunk_pc}")
        self._last_loaded_metadata["release number"] = chunk_release
        self._last_loaded_metadata["serial number"] = chunk_serial
        self._last_loaded_metadata["checksum"] = chunk_checksum
        self._last_loaded_metadata["program counter"] = chunk_pc

        # Verify the save-file params against the current z-story header
        mem = self._zmachine._mem
        if mem.read_word(2) != chunk_release:
            raise QuetzalMismatchedFile
        serial_bytes = chunk_serial
        if serial_bytes != mem[0x12:0x18]:
            raise QuetzalMismatchedFile
        mem_checksum = mem.read_word(0x1C)
        if mem_checksum != 0 and (mem_checksum != chunk_checksum):
            raise QuetzalMismatchedFile
        log("  Quetzal file correctly verifies against original story.")

    def _parse_cmem(self, data):
        """Parse a chunk of type Cmem.  Decompress an image of dynamic
        memory, and place it into the ZMachine."""

        log("  Decompressing dynamic memory image")
        self._seen_mem_or_stks = True

        # Just duplicate the dynamic memory block of the pristine story image,
        # and then make tweaks to it as we decode the runlength-encoding.
        pmem = self._zmachine._pristine_mem
        cmem = self._zmachine._mem
        savegame_mem = list(pmem[pmem._dynamic_start : (pmem._dynamic_end + 1)])
        memlen = len(savegame_mem)
        memcounter = 0
        log(f"  Dynamic memory length is {memlen}")
        self._last_loaded_metadata["memory length"] = memlen

        runlength_bytes = data
        bytelen = len(runlength_bytes)
        bytecounter = 0

        log("  Decompressing dynamic memory image")
        while bytecounter < bytelen:
            byte = runlength_bytes[bytecounter]
            if byte != 0:
                savegame_mem[memcounter] = byte ^ pmem[memcounter]
                memcounter += 1
                bytecounter += 1
                log(f"   Set byte {memcounter}:{savegame_mem[memcounter]}")
            else:
                bytecounter += 1
                num_extra_zeros = runlength_bytes[bytecounter]
                memcounter += 1 + num_extra_zeros
                bytecounter += 1
                log(f"   Skipped {1 + num_extra_zeros} unchanged bytes")
            if memcounter >= memlen:
                raise QuetzalMemoryOutOfBounds

        # If memcounter finishes less then memlen, that's totally fine, it
        # just means there are no more diffs to apply.

        cmem[cmem._dynamic_start : (cmem._dynamic_end + 1)] = savegame_mem
        log("  Successfully installed new dynamic memory.")

    def _parse_umem(self, data):
        """Parse a chunk of type Umem.  Suck a raw image of dynamic memory
        and place it into the ZMachine."""

        ### TODO:  test this by either finding an interpreter that ouptuts
        ## this type of chunk, or by having own QuetzalWriter class
        ## (optionally) do it.
        log("  Loading uncompressed dynamic memory image")
        self._seen_mem_or_stks = True

        cmem = self._zmachine._mem
        dynamic_len = (cmem._dynamic_end - cmem._dynamic_start) + 1
        log(f"  Dynamic memory length is {dynamic_len}")
        self._last_loaded_metadata["dynamic memory length"] = dynamic_len

        savegame_mem = [ord(x) for x in data]
        if len(savegame_mem) != dynamic_len:
            raise QuetzalMemoryMismatch

        cmem[cmem._dynamic_start : (cmem._dynamic_end + 1)] = savegame_mem
        log("  Successfully installed new dynamic memory.")

    def _parse_stks(self, data):
        """Parse a chunk of type Stks."""

        log("  Begin parsing of stack frames")

        # Our strategy here is simply to create an entirely new
        # ZStackManager object and populate it with a series of ZRoutine
        # stack-frames parses from the quetzal file.  We then attach this
        # new ZStackManager to our z-machine, and allow the old one to be
        # garbage collected.
        stackmanager = zstackmanager.ZStackManager(self._zmachine._mem)

        self._seen_mem_or_stks = True
        bytes = data
        total_len = len(bytes)
        ptr = 0

        # Read successive stack frames:
        while ptr < total_len:
            log("  Parsing stack frame...")
            return_pc = (bytes[ptr] << 16) + (bytes[ptr + 1] << 8) + bytes[ptr + 2]
            ptr += 3
            flags_bitfield = bitfield.BitField(bytes[ptr])
            ptr += 1
            varnum = bytes[ptr]
            ptr += 1
            _argflag = bytes[ptr]
            ptr += 1
            evalstack_size = (bytes[ptr] << 8) + bytes[ptr + 1]
            ptr += 2

            # Quetzal flags bit 4: if set, routine discards its return value
            discard_result = flags_bitfield[4]
            store_var = None if discard_result else varnum

            # read anywhere from 0 to 15 local vars, pad to 15
            num_locals = flags_bitfield[0:4]
            local_vars = []
            for _i in range(num_locals):
                var = (bytes[ptr] << 8) + bytes[ptr + 1]
                ptr += 2
                local_vars.append(var)
            # runtime expects 15 slots, save only stores declared count
            local_vars.extend([0] * (15 - num_locals))
            log(f"    Found {num_locals} local vars (padded to 15)")

            # least recent to most recent stack values:
            stack_values = []
            for _i in range(evalstack_size):
                val = (bytes[ptr] << 8) + bytes[ptr + 1]
                ptr += 2
                stack_values.append(val)
            log(f"    Found {len(stack_values)} local stack values")

            # return_pc belongs on the CALLER frame (previous in the stack).
            # When this routine finishes, finish_routine() returns
            # caller.program_counter as the resume address.
            prev_frame = stackmanager._call_stack[-1]
            prev_frame.program_counter = return_pc

            # store_var (varnum) is the variable that receives the return
            # value when this routine finishes — NOT the return PC.
            routine = zstackmanager.ZRoutine(
                0, store_var, self._zmachine._mem, [], local_vars, stack_values
            )
            stackmanager.push_routine(routine)
            log("    Added new frame to stack.")

            if ptr > total_len:
                raise QuetzalStackFrameOverflow

        self._zmachine._stackmanager = stackmanager
        # Update cached references in subsystems that store the stack manager
        # (they cache the stack at init time and won't see the replacement)
        if hasattr(self._zmachine, "_cpu"):
            self._zmachine._cpu._stackmanager = stackmanager
        if hasattr(self._zmachine, "_opdecoder"):
            self._zmachine._opdecoder._stack = stackmanager
        log("  Successfully installed all stack frames.")

    def _parse_intd(self, data):
        """Parse a chunk of type IntD, which is interpreter-dependent info."""

        log("  Begin parsing of interpreter-dependent metadata")
        bytes = [ord(x) for x in data]

        _os_id = bytes[0:3]
        _flags = bytes[4]
        _contents_id = bytes[5]
        _reserved = bytes[6:8]
        _interpreter_id = bytes[8:12]
        _private_data = bytes[12:]
        ### TODO:  finish this

    # The following 3 chunks are totally optional metadata, and are
    # artifacts of the larger IFF standard.  We're not required to do
    # anything when we see them, though maybe it would be nice to print
    # them to the user?

    def _parse_auth(self, data):
        """Parse a chunk of type AUTH.  Display the author."""

        log(f"Author of file: {data}")
        self._last_loaded_metadata["author"] = data

    def _parse_copyright(self, data):
        """Parse a chunk of type (c) .  Display the copyright."""

        log(f"Copyright: (C) {data}")
        self._last_loaded_metadata["copyright"] = data

    def _parse_anno(self, data):
        """Parse a chunk of type ANNO.  Display any annotation"""

        log(f"Annotation: {data}")
        self._last_loaded_metadata["annotation"] = data

    # --------- Public APIs -----------

    def get_last_loaded(self):
        """Return a list of metadata about the last loaded Quetzal file, for
        debugging and test verification."""
        return self._last_loaded_metadata

    def load_from_bytes(self, data):
        """Parse Quetzal data from raw bytes (instead of a file).

        Used by op_restore when filesystem.restore_game() returns raw bytes.
        """
        import io

        self._last_loaded_metadata = {}

        if len(data) < 12:
            raise QuetzalUnrecognizedFileFormat

        # Validate FORM header
        if data[0:4] != b"FORM":
            raise QuetzalUnrecognizedFileFormat

        # Read total length
        self._len = (data[4] << 24) + (data[5] << 16) + (data[6] << 8) + data[7]
        log(f"Total length of FORM data is {self._len}")
        self._last_loaded_metadata["total length"] = self._len

        # Validate IFZS type
        if data[8:12] != b"IFZS":
            raise QuetzalUnrecognizedFileFormat

        # Create a BytesIO object to use with chunk module
        self._file = io.BytesIO(data)
        self._file.seek(12)  # Skip FORM header, length, and IFZS type

        log("Parsing chunks from byte data")
        try:
            while 1:
                chunkname, chunk_data = _read_iff_chunk(self._file)
                chunksize = len(chunk_data)
                log(f"** Found chunk ID {chunkname}: length {chunksize}")
                self._last_loaded_metadata[chunkname] = chunksize

                if chunkname == b"IFhd":
                    self._parse_ifhd(chunk_data)
                elif chunkname == b"CMem":
                    self._parse_cmem(chunk_data)
                elif chunkname == b"UMem":
                    self._parse_umem(chunk_data)
                elif chunkname == b"Stks":
                    self._parse_stks(chunk_data)
                elif chunkname == b"IntD":
                    self._parse_intd(chunk_data)
                elif chunkname == b"AUTH":
                    self._parse_auth(chunk_data)
                elif chunkname == b"(c) ":
                    self._parse_copyright(chunk_data)
                elif chunkname == b"ANNO":
                    self._parse_anno(chunk_data)
                else:
                    # spec says to ignore and skip past unrecognized chunks
                    pass

        except EOFError:
            pass

        self._file.close()
        log("Finished parsing Quetzal data.")

    def load(self, savefile_path):
        """Parse each chunk of the Quetzal file at SAVEFILE_PATH,
        initializing associated zmachine subsystems as needed."""

        self._last_loaded_metadata = {}

        if not os.path.isfile(savefile_path):
            raise QuetzalNoSuchSavefile

        log(f"Attempting to load saved game from '{savefile_path}'")
        self._file = open(savefile_path, "rb")  # noqa: SIM115

        # The python 'chunk' module is pretty dumb; it doesn't understand
        # the FORM chunk and the way it contains nested chunks.
        # Therefore, we deliberately seek 12 bytes into the file so that
        # we can start sucking out chunks.  This also allows us to
        # validate that the FORM type is "IFZS".
        header = self._file.read(4)
        if header != b"FORM":
            raise QuetzalUnrecognizedFileFormat
        bytestring = self._file.read(4)
        self._len = bytestring[0] << 24
        self._len += bytestring[1] << 16
        self._len += bytestring[2] << 8
        self._len += bytestring[3]
        log(f"Total length of FORM data is {self._len}")
        self._last_loaded_metadata["total length"] = self._len

        type = self._file.read(4)
        if type != b"IFZS":
            raise QuetzalUnrecognizedFileFormat

        try:
            while 1:
                chunkname, data = _read_iff_chunk(self._file)
                chunksize = len(data)
                log(f"** Found chunk ID {chunkname}: length {chunksize}")
                self._last_loaded_metadata[chunkname] = chunksize

                if chunkname == b"IFhd":
                    self._parse_ifhd(data)
                elif chunkname == b"CMem":
                    self._parse_cmem(data)
                elif chunkname == b"UMem":
                    self._parse_umem(data)
                elif chunkname == b"Stks":
                    self._parse_stks(data)
                elif chunkname == b"IntD":
                    self._parse_intd(data)
                elif chunkname == b"AUTH":
                    self._parse_auth(data)
                elif chunkname == b"(c) ":
                    self._parse_copyright(data)
                elif chunkname == b"ANNO":
                    self._parse_anno(data)
                else:
                    # spec says to ignore and skip past unrecognized chunks
                    pass

        except EOFError:
            pass

        self._file.close()
        log("Finished parsing Quetzal file.")


# ------------------------------------------------------------------------------


class QuetzalWriter:
    """A class to write the current state of a z-machine into a
    Quetzal-format file."""

    def __init__(self, zmachine):
        log("Creating new instance of QuetzalWriter")
        self._zmachine = zmachine

    def _generate_ifhd_chunk(self):
        """Return a chunk of type IFhd, containing metadata about the
        zmachine and story being played."""

        mem = self._zmachine._mem

        # Release number (2 bytes, big-endian) from header bytes 2-3
        release = mem.read_word(2)

        # Serial number (6 bytes) from header bytes 0x12-0x17
        serial = bytes(mem[0x12:0x18])

        # Checksum (2 bytes, big-endian) from header bytes 0x1C-0x1D
        checksum = mem.read_word(0x1C)

        # Program counter (3 bytes, big-endian) - current PC
        pc = self._zmachine._cpu._program_counter

        # Build the 13-byte chunk
        chunk_data = bytearray(13)

        # Bytes 0-1: Release number
        chunk_data[0] = (release >> 8) & 0xFF
        chunk_data[1] = release & 0xFF

        # Bytes 2-7: Serial number
        chunk_data[2:8] = serial

        # Bytes 8-9: Checksum
        chunk_data[8] = (checksum >> 8) & 0xFF
        chunk_data[9] = checksum & 0xFF

        # Bytes 10-12: Program counter (24-bit)
        chunk_data[10] = (pc >> 16) & 0xFF
        chunk_data[11] = (pc >> 8) & 0xFF
        chunk_data[12] = pc & 0xFF

        return bytes(chunk_data)

    def _generate_cmem_chunk(self):
        """Return a compressed chunk of data representing the compressed
        image of the zmachine's main memory."""

        pmem = self._zmachine._pristine_mem
        cmem = self._zmachine._mem

        # XOR current dynamic memory with pristine dynamic memory
        dynamic_start = pmem._dynamic_start
        dynamic_end = pmem._dynamic_end
        diffarray = []

        for index in range(dynamic_start, dynamic_end + 1):
            xor_value = pmem[index] ^ cmem[index]
            diffarray.append(xor_value)

        log(f"Generated XOR array of {len(diffarray)} bytes")

        # Run-length encode the XOR result
        result = []
        zerocounter = 0

        for byte in diffarray:
            if byte == 0:
                zerocounter += 1
            else:
                # Flush any pending zeros
                while zerocounter > 0:
                    # Encode: 0x00 followed by count of ADDITIONAL zeros
                    # Maximum count in one byte is 255, meaning 256 zeros total (1+255)
                    if zerocounter > 256:
                        result.append(0x00)
                        result.append(0xFF)  # 1 + 255 = 256 zeros
                        zerocounter -= 256
                    else:
                        result.append(0x00)
                        result.append(zerocounter - 1)  # count of additional zeros
                        zerocounter = 0

                # Output non-zero byte
                result.append(byte)

        # Don't encode trailing zeros - parser can leave them as-is
        # (per spec: "If memcounter finishes less than memlen, that's totally fine")

        log(f"Compressed to {len(result)} bytes")
        return bytes(result)

    def _generate_stks_chunk(self):
        """Return a stacks chunk, describing the stack state of the
        zmachine at this moment."""

        result = bytearray()
        stackmanager = self._zmachine._stackmanager
        call_stack = stackmanager._call_stack

        # Skip the ZStackBottom sentinel (first element)
        for i, frame in enumerate(call_stack[1:], start=1):
            num_local_vars = len(frame.local_vars)

            # Quetzal return_pc = caller's saved program counter.
            # The previous frame (caller) stores the resume PC that
            # finish_routine() returns when this frame exits.
            prev_frame = call_stack[i - 1]
            return_pc = prev_frame.program_counter or 0
            result.append((return_pc >> 16) & 0xFF)
            result.append((return_pc >> 8) & 0xFF)
            result.append(return_pc & 0xFF)

            # Write flags byte (bits 0-3 = num local vars,
            # bit 4 = discard return value)
            flags = num_local_vars & 0x0F
            if frame.return_addr is None:
                flags |= 0x10
            result.append(flags)

            # Write varnum (which variable gets the return value)
            varnum = frame.return_addr if frame.return_addr is not None else 0
            result.append(varnum & 0xFF)

            # Write argflag (bitmask of supplied arguments)
            # TODO: track this properly, for now use 0
            result.append(0)

            # Write eval_stack_size as 16-bit big-endian
            eval_stack_size = len(frame.stack)
            result.append((eval_stack_size >> 8) & 0xFF)
            result.append(eval_stack_size & 0xFF)

            # Write local variables (16-bit big-endian each)
            for local_var in frame.local_vars:
                result.append((local_var >> 8) & 0xFF)
                result.append(local_var & 0xFF)

            # Write evaluation stack values (16-bit big-endian each)
            for stack_val in frame.stack:
                result.append((stack_val >> 8) & 0xFF)
                result.append(stack_val & 0xFF)

        return bytes(result)

    def _generate_anno_chunk(self):
        """Return an annotation chunk, containing metadata about the ZVM
        interpreter which created the savefile."""

        ### TODO:  write this
        return b"0"

    # --------- Public APIs -----------

    def generate_save_data(self):
        """Generate complete Quetzal save data as bytes (IFF/FORM/IFZS container).

        Returns bytes representing the complete save file in IFF format.
        """
        log("Generating Quetzal save data")

        # Generate all chunks
        ifhd_chunk = self._generate_ifhd_chunk()
        cmem_chunk = self._generate_cmem_chunk()
        stks_chunk = self._generate_stks_chunk()
        anno_chunk = self._generate_anno_chunk()

        # Build IFF container with proper chunk headers
        result = bytearray()

        # Helper to write a chunk with its header
        def write_chunk(chunk_id, chunk_data):
            result.extend(chunk_id.encode("ascii"))
            size = len(chunk_data)
            result.append((size >> 24) & 0xFF)
            result.append((size >> 16) & 0xFF)
            result.append((size >> 8) & 0xFF)
            result.append(size & 0xFF)
            result.extend(chunk_data)
            # IFF chunks must be padded to even byte boundaries
            if size % 2 == 1:
                result.append(0)  # padding byte
            log(f"  Added {chunk_id} chunk ({size} bytes)")

        # Write nested chunks
        write_chunk("IFhd", ifhd_chunk)
        write_chunk("CMem", cmem_chunk)
        write_chunk("Stks", stks_chunk)
        write_chunk("ANNO", anno_chunk)

        # Calculate total size (everything after FORM header + size field)
        total_size = len(result) + 4  # +4 for "IFZS"

        # Build final FORM container
        container = bytearray()
        container.extend(b"FORM")
        container.append((total_size >> 24) & 0xFF)
        container.append((total_size >> 16) & 0xFF)
        container.append((total_size >> 8) & 0xFF)
        container.append(total_size & 0xFF)
        container.extend(b"IFZS")
        container.extend(result)

        log(f"Generated {len(container)} bytes of save data")
        return bytes(container)

    def write(self, savefile_path):
        """Write the current zmachine state to a new Quetzal-file at
        SAVEFILE_PATH."""

        log(f"Attempting to write game-state to '{savefile_path}'")
        data = self.generate_save_data()

        with open(savefile_path, "wb") as f:
            f.write(data)

        log("Done writing game-state to savefile.")
