"""USB Serial/JTAG SPI NAND file transfer client (Python 3 + pyserial)."""
import argparse
import os
import secrets
import stat
import struct
import sys
import time
import zlib
from pathlib import Path

try:
    import serial
except ImportError:
    serial = None

MAGIC = b"NXFR"
HEADER = struct.Struct("<4sBBHII")
INFO, LIST, MKDIR, BEGIN, BLOCK, END, VERIFY, CANCEL, PREPARE, CONFIRM, HELLO = range(1, 12)
ERRORS = ["OK", "bad frame", "bad path", "not mounted", "I/O error",
          "bad state", "CRC/size mismatch", "denied"]
ROOT = "/extflash"


def path_bytes(path, allow_root=False):
    if not isinstance(path, str) or not path.startswith(ROOT):
        raise ValueError("destination must start with /extflash")
    if path == ROOT:
        if not allow_root:
            raise ValueError("the volume root is not a file/directory name")
    elif not path.startswith(ROOT + "/"):
        raise ValueError("path outside /extflash")
    else:
        parts = path[len(ROOT) + 1:].split("/")
        if any(not p or p in (".", "..") or p[-1] in " ." or
               any(ord(c) < 32 or c in '\\:*?"<>|' for c in p) for p in parts):
            raise ValueError("invalid FAT path")
    value = path.encode("utf-8")
    if len(value) >= 512:
        raise ValueError("UTF-8 path too long")
    return value


def encode(cmd, seq, payload=b""):
    if len(payload) > 4352:
        raise ValueError("frame too large")
    body = HEADER.pack(MAGIC, 1, cmd, 0, seq, len(payload)) + payload
    return body + struct.pack("<I", zlib.crc32(body) & 0xffffffff)


class Client:
    def __init__(self, port):
        self.port = port
        self.seq = secrets.randbits(32)

    def wait_ready(self, timeout=0.4):
        deadline = time.monotonic() + timeout
        buf = b""
        old_timeout = getattr(self.port, "timeout", None)
        try:
            if old_timeout is not None:
                self.port.timeout = min(0.05, timeout)
            while time.monotonic() < deadline:
                chunk = self.port.read(1024)
                if not chunk:
                    continue
                buf = (buf + chunk)[-96:]
                if b"NAND-XFER READY" in buf:
                    return True
        finally:
            if old_timeout is not None:
                self.port.timeout = old_timeout
        return False

    def drain_rx(self):
        old_timeout = getattr(self.port, "timeout", None)
        try:
            if old_timeout is not None:
                self.port.timeout = 0
            while self.port.read(4096):
                pass
        finally:
            if old_timeout is not None:
                self.port.timeout = old_timeout

    def connect(self):
        self.wait_ready()
        self.drain_rx()
        nonce = secrets.token_bytes(4)
        if self.call(HELLO, nonce) != nonce:
            raise ValueError("device handshake failed")

    def _frame(self, deadline):
        match = bytearray()
        while time.monotonic() < deadline:
            byte = self.port.read(1)
            if not byte:
                continue
            match = (match + byte)[-4:]
            if match == MAGIC:
                break
        else:
            raise TimeoutError("no response magic")
        header = MAGIC + self._exact(12, deadline)
        magic, version, cmd, flags, seq, size = HEADER.unpack(header)
        if version != 1 or flags or size > 4352:
            raise ValueError("invalid response header")
        payload = self._exact(size, deadline)
        checksum, = struct.unpack("<I", self._exact(4, deadline))
        if (zlib.crc32(header + payload) & 0xffffffff) != checksum:
            raise ValueError("response CRC mismatch")
        return cmd, seq, payload

    def _exact(self, size, deadline):
        data = bytearray()
        while len(data) < size and time.monotonic() < deadline:
            data.extend(self.port.read(size - len(data)))
        if len(data) != size:
            raise TimeoutError("incomplete response")
        return bytes(data)

    def call(self, cmd, payload=b"", timeout=12, retries=3):
        self.seq = (self.seq + 1) & 0xffffffff
        seq = self.seq
        packet = encode(cmd, seq, payload)
        for attempt in range(retries):
            self.port.write(packet)
            deadline = time.monotonic() + timeout
            while time.monotonic() < deadline:
                try:
                    reply_cmd, reply_seq, body = self._frame(deadline)
                except (ValueError, TimeoutError):
                    break
                if reply_seq != seq or reply_cmd != (cmd | 0x80):
                    continue
                if not body:
                    raise RuntimeError("empty reply")
                if body[0]:
                    code = body[0]
                    raise RuntimeError(ERRORS[code] if code < len(ERRORS) else f"error {code}")
                return body[1:]
        raise TimeoutError(f"command {cmd} unacknowledged after {retries} attempts")

    def info(self):
        page, pages, mounted, total, free = struct.unpack("<IIBQQ", self.call(INFO))
        return page, pages, bool(mounted), total, free

    def verify(self, path):
        return struct.unpack("<QI", self.call(VERIFY, path_bytes(path), timeout=60))

    def listdir(self, path):
        encoded = path_bytes(path, allow_root=True)
        cursor = 0
        while True:
            body = self.call(LIST, struct.pack("<H", cursor) + encoded)
            if len(body) < 3:
                raise ValueError("invalid directory response")
            next_cursor, done = struct.unpack("<HB", body[:3])
            i = 3
            while i < len(body):
                if i + 10 > len(body):
                    raise ValueError("truncated directory entry")
                is_dir, name_len, size = struct.unpack_from("<BBQ", body, i)
                i += 10
                if i + name_len > len(body):
                    raise ValueError("truncated directory name")
                name = body[i:i + name_len].decode("utf-8")
                i += name_len
                yield (is_dir, name, size)
            if done:
                return
            if next_cursor <= cursor:
                raise ValueError("directory cursor did not advance")
            cursor = next_cursor

    def upload(self, source, dest, log=print):
        dest = path_bytes(dest).decode("utf-8")
        try:
            self.call(BEGIN, path_bytes(dest))
            size = 0
            crc = 0
            with open(source, "rb") as stream:
                while True:
                    data = stream.read(4096)
                    if not data:
                        break
                    if size > 0xffffffff:
                        raise ValueError("file exceeds 4 GiB protocol offset limit")
                    ack = self.call(BLOCK, struct.pack("<I", size) + data, timeout=30)
                    size += len(data)
                    crc = zlib.crc32(data, crc) & 0xffffffff
                    if len(ack) != 8 or struct.unpack("<Q", ack)[0] != size:
                        raise ValueError("unexpected device offset")
            result = self.call(END, struct.pack("<QI", size, crc), timeout=90)
            if len(result) != 12 or struct.unpack("<QI", result) != (size, crc):
                raise ValueError("device readback differs from source")
            if self.verify(dest) != (size, crc):
                raise ValueError("post-close verification differs")
            log(f"OK {source} -> {dest} ({size} bytes, CRC32 {crc:08x})")
        except Exception:
            try:
                self.call(CANCEL, timeout=1, retries=1)
            except Exception:
                pass
            raise


def directory_entries(source):
    root = Path(source)
    for base, dirs, files in os.walk(root, followlinks=False):
        dirs.sort()
        files.sort()
        for name in sorted(dirs + files):
            path = Path(base) / name
            mode = path.lstat().st_mode
            if stat.S_ISLNK(mode) or (hasattr(path, "is_junction") and path.is_junction()):
                raise ValueError(f"link/junction not allowed: {path}")
            if not (stat.S_ISDIR(mode) or stat.S_ISREG(mode)):
                raise ValueError(f"not a normal file or directory: {path}")
            yield path, path.relative_to(root).parts, stat.S_ISDIR(mode)


def open_port(name):
    port = serial.Serial()
    port.port = name
    port.baudrate = 115200
    port.timeout = 0.2
    port.write_timeout = 15
    port.dsrdtr = False
    port.rtscts = False
    port.open()
    try:
        port.setRTS(False)
        port.setDTR(True)
    except (AttributeError, OSError):
        pass
    return port


def format_volume(client):
    token = client.call(PREPARE)
    if len(token) != 4:
        raise ValueError("invalid format challenge")
    # Confirmation is acquired before PREPARE so the 30-second device token cannot expire at the prompt.
    client.call(CONFIRM, token + struct.pack("<I", client.seq), timeout=180, retries=1)


def format_with_confirmation(client, input_fn=input):
    print("WARNING: this erases ALL files on the external NAND volume.")
    if input_fn("Type FORMAT /extflash to continue: ") != "FORMAT /extflash":
        raise ValueError("format cancelled")
    format_volume(client)
    print("NAND formatted and remounted")


def main(argv=None):
    parser = argparse.ArgumentParser(description="Transfer files to external SPI NAND over USB Serial/JTAG")
    parser.add_argument("port")
    parser.add_argument("source", nargs="?")
    parser.add_argument("destination", nargs="?")
    parser.add_argument("--recursive", action="store_true")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--info", action="store_true")
    group.add_argument("--list", metavar="DEVICE_DIR")
    group.add_argument("--mkdir", metavar="DEVICE_DIR")
    group.add_argument("--verify", metavar="DEVICE_FILE")
    group.add_argument("--format", action="store_true")
    args = parser.parse_args(argv)
    if args.recursive and not (args.source and args.destination):
        parser.error("--recursive requires source and destination")
    if args.source or args.destination:
        if not args.source or not args.destination or any((args.info, args.list, args.mkdir, args.verify, args.format)):
            parser.error("specify both source and destination, with no command option")
    elif not any((args.info, args.list, args.mkdir, args.verify, args.format)):
        parser.error("specify a file/directory transfer or a command")
    if serial is None:
        parser.error("pyserial missing: py -3 -m pip install pyserial")
    try:
        with open_port(args.port) as port:
            client = Client(port)
            client.connect()
            if args.format:
                format_with_confirmation(client)
            elif args.info:
                page, pages, mounted, total, free = client.info()
                print(f"NAND: {page * pages} bytes ({pages} pages x {page}); mounted={mounted}; FAT={total} bytes, free={free}")
            elif args.list:
                for is_dir, name, size in client.listdir(args.list):
                    print(f"{'DIR ' if is_dir else 'FILE'} {name} {size if not is_dir else ''}")
            elif args.mkdir:
                client.call(MKDIR, path_bytes(args.mkdir))
            elif args.verify:
                size, crc = client.verify(args.verify)
                print(f"{args.verify}: {size} bytes, CRC32 {crc:08x}")
            else:
                source = Path(args.source)
                if source.is_symlink() or (hasattr(source, "is_junction") and source.is_junction()):
                    raise ValueError("links and junctions are not allowed")
                if source.is_file() and not args.recursive:
                    client.upload(source, args.destination)
                elif source.is_dir() and args.recursive:
                    target = path_bytes(args.destination).decode("utf-8")
                    entries = list(directory_entries(source))
                    for _, parts, _ in entries:
                        path_bytes(target + "/" + "/".join(parts))
                    client.call(MKDIR, path_bytes(target))
                    for path, parts, is_dir in entries:
                        remote = target + "/" + "/".join(parts)
                        if is_dir:
                            client.call(MKDIR, path_bytes(remote))
                        else:
                            client.upload(path, remote)
                elif not source.exists():
                    raise ValueError(f"source not found: {source}")
                else:
                    raise ValueError("use --recursive for a directory; omit it for a single file")
    except (OSError, ValueError, RuntimeError, TimeoutError, serial.SerialException) as exc:
        parser.exit(1, f"error: {exc}\n")


if __name__ == "__main__":
    main()