# TODO: use `glasgow.arch.qspi.nor` for decoding? (This was too much to refactor at once.)

from collections.abc import Buffer
from typing import BinaryIO
from abc import ABCMeta, abstractmethod
import enum
import struct
import logging
import argparse

from glasgow.database.jedec import jedec_mfg_name_from_bytes
from glasgow.support.logging import dump_hex
from glasgow.protocol.sfdp import SFDPCollection, SFDPJEDECFlashParametersTable
from glasgow.applet import GlasgowAppletTool
from . import Memory25QApplet


__all__ = ["MemoryImage", "Memory25QDecoder"]


class AddressingMode(enum.Enum):
    ThreeByte = enum.auto()
    FourByte  = enum.auto()


class MemoryImage:
    def __init__(self, size=None, *, wrap=None):
        self._wrap = wrap
        if wrap is not None:
            assert size is None
            size = wrap
        elif size is None:
            size = 0
        self._data = bytearray(size)
        self._mask = bytearray(size)

    def __len__(self):
        return len(self._data)

    def __bool__(self):
        try:
            self._mask.index(b"\xff")
        except ValueError:
            return False
        else:
            return True

    @property
    def data(self) -> memoryview:
        return memoryview(self._data)

    @property
    def mask(self) -> memoryview:
        return memoryview(self._mask)

    def read(self, addr: int, size: int, *, if_present=False) -> memoryview | None:
        if addr < 0:
            raise IndexError(f"start address {addr:#x} is out of bounds")
        if len(self._data) < addr + size:
            raise IndexError(f"end address {addr + size:#x} is out of bounds")
        if if_present:
            if self.mask[addr:addr + size] != b"\xff" * size:
                return None
        return self.data[addr:addr + size]

    def write(self, addr: int, chunk: Buffer):
        if self._wrap is None:
            extra_len = addr + len(chunk) - len(self.data)
            if extra_len > 0:
                self._data.extend(bytearray(extra_len))
                self._mask.extend(bytearray(extra_len))
            self._data[addr:addr + len(chunk)] = chunk
            self._mask[addr:addr + len(chunk)] = b"\xff" * len(chunk)
        else:
            # This is very slow, but currenly is only used for very small images.
            for offset, byte in enumerate(chunk):
                byte_addr = (addr + offset) % self._wrap
                self._data[byte_addr] = byte
                self._mask[byte_addr] = 0xff

    def save(self, data_file: BinaryIO | None, mask_file: BinaryIO | None = None):
        if data_file is not None:
            data_file.write(self._data)
        if mask_file is not None:
            mask_file.write(self._mask)


class Memory25QTraceTypeError(Exception):
    pass


class Memory25QAbstractTrace(metaclass=ABCMeta):
    @abstractmethod
    def read_copi(self, count: int | None = None) -> memoryview:
        """Reads ``count`` bytes (or until the end of trace) by sampling COPI line."""

    @abstractmethod
    def read_cipo(self, count: int | None = None) -> memoryview:
        """Reads ``count`` bytes (or until the end of trace) by sampling CIPO line."""

    @abstractmethod
    def read_dual(self, count: int | None = None) -> memoryview:
        """Reads ``count`` bytes (or until the end of trace) by sampling IO0/IO1
        (COPI/CIPO) lines.
        """

    @abstractmethod
    def read_quad(self, count: int | None = None) -> memoryview:
        """Reads ``count`` bytes (or until the end of trace) by sampling IO0/IO1/IO2/IO3
        (COPI/CIPO/WP#/HOLD#) lines.
        """


class Memory25QSPITrace(Memory25QAbstractTrace):
    def __init__(self, copi: bytes, cipo: bytes):
        self._at   = 0
        self._copi = memoryview(copi)
        self._cipo = memoryview(cipo)

    def read_copi(self, count: int | None = None) -> memoryview:
        if count is None:
            count = len(self._copi) - self._at
        if self._at + count > len(self._copi):
            raise IndexError("read past end of trace")
        data = self._copi[self._at:self._at + count]
        self._at += count
        return data

    def read_cipo(self, count: int | None = None) -> memoryview:
        if count is None:
            count = len(self._cipo) - self._at
        if self._at + count > len(self._cipo):
            raise IndexError("read past end of trace")
        data = self._cipo[self._at:self._at + count]
        self._at += count
        return data

    def read_dual(self, count: int | None = None) -> memoryview:
        raise Memory25QTraceTypeError("dual I/O read requires a QSPI trace")

    def read_quad(self, count: int | None = None) -> memoryview:
        raise Memory25QTraceTypeError("quad I/O read requires a QSPI trace")


class Memory25QQSPITrace(Memory25QAbstractTrace):
    def __init__(self, data: bytes):
        self._at   = 0
        self._data = memoryview(data)

    def _get(self, count: int):
        data = self._data[self._at:self._at + count]
        self._at += count
        return data

    def read_copi(self, count: int | None = None) -> memoryview:
        if count is None:
            count = (len(self._data) - self._at) // 4
        if self._at + 4 * count > len(self._data):
            raise IndexError("read past end of trace")
        def get_copi():
            byte = 0
            for code in self._get(4):
                byte = (byte << 1) | (code & 0x10) >> 4
                byte = (byte << 1) | (code & 0x01)
            return byte
        return memoryview(bytes(get_copi() for _ in range(count)))

    def read_cipo(self, count: int | None = None) -> memoryview:
        if count is None:
            count = (len(self._data) - self._at) // 4
        if self._at + 4 * count > len(self._data):
            raise IndexError("read past end of trace")
        def get_cipo():
            byte = 0
            for code in self._get(4):
                byte = (byte << 1) | (code & 0x20) >> 5
                byte = (byte << 1) | (code & 0x02) >> 1
            return byte
        return memoryview(bytes(get_cipo() for _ in range(count)))

    def read_dual(self, count: int | None = None) -> memoryview:
        if count is None:
            count = (len(self._data) - self._at) // 2
        if self._at + 2 * count > len(self._data):
            raise IndexError("read past end of trace")
        def get_dual():
            byte = 0
            for code in self._get(2):
                byte = (byte << 2) | (code & 0x30) >> 4
                byte = (byte << 2) | (code & 0x03)
            return byte
        return memoryview(bytes(get_dual() for _ in range(count)))

    def read_quad(self, count: int | None = None) -> memoryview:
        if count is None:
            count = len(self._data) - self._at
        return self._get(count)


class Memory25QDecoder:
    def __init__(self, logger: logging.Logger | None = None):
        self._logger   = logger

        self._jedec_id = None
        self._data     = MemoryImage()
        self._sfdp     = MemoryImage(256)
        self._uid      = MemoryImage(wrap=16)
        self._unknown  = set()

        self._index    = 0
        self._write    = False
        self._mode     = AddressingMode.ThreeByte

    def _log(self, message, *args, level=logging.DEBUG):
        if self._logger:
            self._logger.log(level, f"25x decode: [{self._index}] " + message, *args)

    @property
    def jedec_id(self) -> tuple[int, int] | None:
        return self._jedec_id

    @property
    def data(self) -> MemoryImage:
        return self._data

    @property
    def sfdp(self) -> MemoryImage:
        return self._sfdp

    @property
    def uid(self) -> MemoryImage:
        return self._uid

    @property
    def unknown(self) -> set[int]:
        return self._unknown

    def _decode_addr(self, trace, mode, *, x=1) -> int:
        match mode:
            case AddressingMode.ThreeByte:
                match x:
                    case 1: data = trace.read_copi(3)
                    case 2: data = trace.read_dual(3)
                    case 4: data = trace.read_quad(3)
                addr = data[0] << 16 | data[1] << 8  | data[2] << 0
                self._log(f"  {addr=:06x}")
            case AddressingMode.FourByte:
                match x:
                    case 1: data = trace.read_copi(4)
                    case 2: data = trace.read_dual(4)
                    case 4: data = trace.read_quad(4)
                addr = data[0] << 24 | data[1] << 16 | data[2] << 8 | data[3] << 0
                self._log(f"  {addr=:08x}")
        return addr

    def decode_trace(self, trace):
        try:
            cmd = None
            cmd = trace.read_copi(1)[0] # this can raise
            match cmd:
                case 0x03:
                    self._log(f"{cmd=:02X} (Read Data)")
                    addr = self._decode_addr(trace, self._mode)
                    data = trace.read_cipo()
                    self._log("  data=%s", dump_hex(data))
                    self._data.write(addr, data)

                case 0x04:
                    self._log(f"{cmd=:02X} (Write Disable)")
                    self._write = False

                case 0x05:
                    self._log(f"{cmd=:02X} (Read Status Register)")
                    # We don't do anything with the value; this could be used to verify whether
                    # preceding writes have completed or not.

                case 0x06:
                    self._log(f"{cmd=:02X} (Write Enable)")
                    self._write = True

                case 0x0B:
                    self._log(f"{cmd=:02X} (Fast Read)")
                    addr = self._decode_addr(trace, self._mode)
                    _    = trace.read_cipo(1)
                    data = trace.read_cipo()
                    self._log("  data=%s", dump_hex(data))
                    self._data.write(addr, data)

                case 0x4B:
                    self._log(f"{cmd=:02X} (Read Unique ID)")
                    addr = self._decode_addr(trace, AddressingMode.ThreeByte)
                    _    = trace.read_cipo(1)
                    data = trace.read_cipo()
                    self._log("  data=%s", dump_hex(data))
                    self._uid.write(addr, data)

                case 0x5A:
                    self._log(f"{cmd=:02X} (Read SFDP)")
                    addr = self._decode_addr(trace, AddressingMode.ThreeByte)
                    _    = trace.read_cipo(1)
                    data = trace.read_cipo()
                    self._log("  data=%s", dump_hex(data))
                    self._sfdp.write(addr, data)

                case 0x9F:
                    self._log(f"{cmd=:02X} (Read JEDEC ID)")
                    mfg_id, device_id = struct.unpack(">BH", trace.read_cipo(3))
                    self._jedec_id = (mfg_id, device_id)

                case 0xB7:
                    self._log(f"{cmd=:02X} (Enter 4-Byte Address Mode)")
                    self._mode = AddressingMode.FourByte

                case 0xE9:
                    self._log(f"{cmd=:02X} (Exit 4-Byte Address Mode)")
                    self._mode = AddressingMode.ThreeByte

                case 0x3B:
                    self._log(f"{cmd=:02X} (Dual Output Fast Read)")
                    addr = self._decode_addr(trace, self._mode, x=1)
                    _    = trace.read_copi(2)
                    data = trace.read_dual()
                    self._log("  data=%s", dump_hex(data))
                    self._data.write(addr, data)

                case 0x6B:
                    self._log(f"{cmd=:02X} (Quad Output Fast Read)")
                    addr = self._decode_addr(trace, self._mode, x=1)
                    _    = trace.read_copi(2)
                    data = trace.read_quad()
                    self._log("  data=%s", dump_hex(data))
                    self._data.write(addr, data)

                case 0xBB:
                    self._log(f"{cmd=:02X} (Dual I/O Fast Read)")
                    addr = self._decode_addr(trace, self._mode, x=2)
                    mode = trace.read_dual(1)[0]
                    _    = trace.read_dual(2)
                    data = trace.read_dual()
                    assert mode == 0, f"mode bits ({mode:08b}) not implemented"
                    self._log("  data=%s", dump_hex(data))
                    self._data.write(addr, data)

                case 0xEB:
                    self._log(f"{cmd=:02X} (Quad I/O Fast Read)")
                    addr = self._decode_addr(trace, self._mode, x=4)
                    mode = trace.read_quad(1)[0]
                    _    = trace.read_quad(2)
                    data = trace.read_quad()
                    assert mode == 0, f"mode bits ({mode:08b}) not implemented"
                    self._log("  data=%s", dump_hex(data))
                    self._data.write(addr, data)

                case _:
                    self._log(f"{cmd=:02X} (unknown)", level=logging.WARNING)
                    self._unknown.add(cmd)

        except IndexError:
            if cmd is None:
                self._log("(truncated)", level=logging.WARNING)
            else:
                self._log(f"{cmd=:02X} (truncated)", level=logging.ERROR)

        except Memory25QTraceTypeError as exn:
            self._log(str(exn), level=logging.ERROR)

        finally:
            self._index += 1

    def decode_spi(self, copi, cipo):
        self.decode_trace(Memory25QSPITrace(copi, cipo))

    def decode_qspi(self, data):
        self.decode_trace(Memory25QQSPITrace(data))


class Memory25QAppletTool(GlasgowAppletTool, applet=Memory25QApplet):
    logger = logging.getLogger(__name__)
    help = "decode communications with 25-series SPI NOR Flash memories and extract data"
    description = """
    Dissect captured SPI/QSPI transactions and extract data into linear memory image files.

    The expected capture file format is the same as ones used by `spi-analyzer` and `qspi-analyzer`
    applets. Specifically, one of the following Comma-Separated Value line formats is expected:

    * ``<COPI>,<CIPO>``, where <COPI> and <CIPO> are hexadecimal byte sequences with each eight
      bits corresponding to samples of COPI and CIPO, respectively (from MSB to LSB).
    * ``<DATA>``, where <DATA> is a hexadecimal nibble sequence with each four bits corresponding
      to samples of HOLD#, WP#, CIPO, COPI (from MSB to LSB).

    The extracted data can be saved as pairs of data files and mask files, where the mask file
    contains a 1 bit for every bit in the data file that has been observed in a transaction,
    and a 0 bit otherwise.

    The list below details every command that is recognized by this tool. If your capture includes
    commands not currently recognized, please open an issue with a capture file attached.

    * 03h (Read Data)
    * 04h (Write Disable)
    * 05h (Read Status Register)
    * 06h (Write Enable)
    * 0Bh (Fast Read)
    * 4Bh (Read Unique ID)
    * 5Ah (Read SFDP)
    * 9Fh (Read JEDEC ID)
    * B7h (Enter 4-Byte Address Mode)
    * E9h (Exit 4-Byte Address Mode)
    * 3Bh (Dual Output Fast Read)
    * 6Bh (Quad Output Fast Read)
    * BBh (Dual I/O Fast Read)
    * EBh (Quad I/O Fast Read)
    """

    @classmethod
    def add_arguments(cls, parser):
        parser.add_argument(
            "capture_file", metavar="CAPTURE-FILE", type=argparse.FileType("r"),
            help="read captured SPI transactions from CAPTURE-FILE")

        parser.add_argument(
            "--data", dest="data_file", metavar="DATA-FILE", type=argparse.FileType("wb"),
            help="write extracted data to DATA-FILE")
        parser.add_argument(
            "--data-mask", dest="data_mask_file", metavar="MASK-FILE", type=argparse.FileType("wb"),
            help="write data presence mask to MASK-FILE")

        parser.add_argument(
            "--sfdp", dest="sfdp_file", metavar="DATA-FILE", type=argparse.FileType("wb"),
            help="write extracted SFDP data to DATA-FILE")
        parser.add_argument(
            "--sfdp-mask", dest="sfdp_mask_file", metavar="MASK-FILE", type=argparse.FileType("wb"),
            help="write SFDP data presence mask to MASK-FILE")

        parser.add_argument(
            "--uid", dest="uid_file", metavar="DATA-FILE", type=argparse.FileType("wb"),
            help="write extracted UID data to DATA-FILE")
        parser.add_argument(
            "--uid-mask", dest="uid_mask_file", metavar="MASK-FILE", type=argparse.FileType("wb"),
            help="write UID data presence mask to MASK-FILE")

    async def run(self, args):
        decoder = Memory25QDecoder(self.logger)

        for index, line in enumerate(args.capture_file):
            try:
                match line.split(","):
                    case (copi, cipo):
                        decoder.decode_spi(bytes.fromhex(copi), bytes.fromhex(cipo))
                    case (data,):
                        decoder.decode_qspi(bytes.fromhex(data))
                    case _:
                        self.logger.error(f"line {index + 1}: unrecognized data")
            except ValueError:
                self.logger.error(f"line {index + 1}: invalid hex digit")

        if decoder.unknown:
            self.logger.warning("unknown commands encountered: %s",
                ", ".join(f"{cmd:02X}" for cmd in decoder.unknown))

        if decoder.jedec_id is not None:
            mfg_id, device_id = decoder.jedec_id
            mfg_name = jedec_mfg_name_from_bytes([mfg_id]) or "unknown"
            self.logger.info(
                f"JEDEC manufacturer {mfg_id:#04x} ({mfg_name}) device {device_id:#06x}")
        else:
            self.logger.info("capture does not have JEDEC device ID")

        try:
            async def read_sfdp(offset, length):
                return decoder.sfdp.read(offset, length)
            sfdp = await SFDPCollection.parse(read_sfdp)
            self.logger.info(f"capture contains valid {sfdp} descriptor")
            for line in sfdp.description():
                self.logger.info(f"  {line}")

            for table in sfdp:
                if isinstance(table, SFDPJEDECFlashParametersTable):
                    decoder.data.write(table.density_bits >> 3, b"") # stretch to real size

        except ValueError as exn:
            self.logger.info(f"capture does not have valid SFDP data: {exn}")

        if decoder.uid:
            self.logger.info("capture contains UID")

        decoder.data.save(args.data_file, args.data_mask_file)
        decoder.sfdp.save(args.sfdp_file, args.sfdp_mask_file)
        decoder.uid.save(args.uid_file, args.uid_mask_file)
