"""TFTP receiver emulator for testing SRIO-over-TFTP flash protocol.

This server emulates the target hardware: it decodes incoming TFTP filenames
that encode SRIO commands ($SRIO:<slot>/<address>+<len>) and simulates
the complete SRIO flash controller with registers (MODE_REG, CMD_REG, STATUS_REG, etc.).

Usage:
    python tools/tftp_receiver.py --port 6969 --workdir ./tftp_root

Reference: dwl_fw.cpp, fpgaflashinterface.h, qgtftptargetsim.cpp

Features:
- Simulates SRIO flash controller registers (from dwl_fw.h)
- Handles op_init_qspi() initialization sequence
- Processes erase_section(), write_flash(), read_flash(), wait_flash()
- Logs all register operations for debugging

Note: Uses custom TFTP server implementation for Windows compatibility.
"""
from __future__ import annotations

import argparse
import re
import socket
import struct
import sys
import threading
import time
from pathlib import Path
from typing import Dict, Optional, Tuple


# SRIO Register addresses (from dwl_fw.h)
MODE_REG = 0x4700002C
CMD_REG = 0x47000030
ADDR_REG = 0x47000034
NUM_BYTE_REG = 0x47000038
CTRL_REG = 0x47000060
TX_FIFO_REG = 0x47000400
STATUS_REG = 0x47000864
RX_FIFO_REG = 0x47000C00


class SRIOFlashController:
    """Simulates SRIO flash controller with register state machine.
    
    Emulates hardware behavior including:
    - Register state (MODE_REG, CMD_REG, CTRL_REG, STATUS_REG, etc.)
    - Flash operations (erase, write, read)
    - Status bit transitions (busy, ack, error)
    """Simulates a flash memory using a file on disk."""

    def __init__(self, image_path: Path, size: int = 16 * 1024 * 1024) -> None:
        self.image_path = image_path
        self.size = size
        if not image_path.exists():
            # Create empty flash image
            with open(image_path, "wb") as f:
                f.write(b"\xFF" * size)
        print(f"Flash emulator initialized: {image_path} ({size} bytes)")

    def read(self, address: int, length: int) -> bytes:
        """Read from flash."""
        if address + length > self.size:
            raise ValueError(f"Read out of bounds: 0x{address:X}+{length} > {self.size}")
        with open(self.image_path, "rb") as f:
            f.seek(address)
            return f.read(length)

    def write(self, address: int, data: bytes) -> None:
        """Write to flash."""
        length = len(data)
        if address + length > self.size:
            raise ValueError(f"Write out of bounds: 0x{address:X}+{length} > {self.size}")
        with open(self.image_path, "r+b") as f:
            f.seek(address)
            f.write(data)
        print(f"  Flash write: 0x{address:08X} + {length} bytes")

    def erase(self, address: int, length: int = 65536) -> None:
        """Erase flash region (fill with 0xFF)."""
        if address + length > self.size:
            raise ValueError(f"Erase out of bounds: 0x{address:X}+{length} > {self.size}")
        with open(self.image_path, "r+b") as f:
            f.seek(address)
            f.write(b"\xFF" * length)
        print(f"  Flash erase: 0x{address:08X} + {length} bytes")


def parse_srio_filename(filename: str) -> Optional[Tuple[str, int, int]]:
    """Parse SRIO filename format: $SRIO:<slot>/<address>+<len>.

    Returns:
        (slot, address, length) or None if not a valid SRIO command.

    Examples:
        >>> parse_srio_filename("$SRIO:0x10/0x1000+256")
        ('0x10', 4096, 256)
    """
    match = re.match(r"\$SRIO:([^/]+)/(0x[0-9A-Fa-f]+)\+(\d+)", filename)
    if match:
        slot = match.group(1)
        address = int(match.group(2), 16)
        length = int(match.group(3))
        return (slot, address, length)
    return None


class SRIOTFTPHandler:
    """Custom TFTP handler that processes SRIO commands."""

    def __init__(self, workdir: Path, flash: FlashEmulator) -> None:
        self.workdir = workdir
        self.flash = flash
        self.workdir.mkdir(parents=True, exist_ok=True)

    def handle_read(self, filename: str) -> Optional[bytes]:
        """Handle TFTP GET (read) request."""
        parsed = parse_srio_filename(filename)
        if parsed:
            slot, address, length = parsed
            print(f"  SRIO READ: slot={slot}, addr=0x{address:08X}, len={length}")
            try:
                data = self.flash.read(address, length)
                return data
            except Exception as e:
                print(f"    ERROR: {e}")
                return None
        else:
            # Normal file read (not SRIO command)
            filepath = self.workdir / filename
            if filepath.exists():
                with open(filepath, "rb") as f:
                    return f.read()
            return None

    def handle_write(self, filename: str, data: bytes) -> bool:
        """Handle TFTP PUT (write) request."""
        parsed = parse_srio_filename(filename)
        if parsed:
            slot, address, length = parsed
            print(f"  SRIO WRITE: slot={slot}, addr=0x{address:08X}, len={len(data)}")
            try:
                self.flash.write(address, data)
                return True
            except Exception as e:
                print(f"    ERROR: {e}")
                return False
        else:
            # Normal file write (not SRIO command)
            filepath = self.workdir / filename
            with open(filepath, "wb") as f:
                f.write(data)
            print(f"  File written: {filepath}")
            return True


class SimpleTFTPServer:
    """Minimal TFTP server (RFC 1350) for Windows compatibility.
    
    Supports RRQ and WRQ operations with custom handler for SRIO commands.
    """

    # TFTP opcodes
    OP_RRQ = 1
    OP_WRQ = 2
    OP_DATA = 3
    OP_ACK = 4
    OP_ERROR = 5

    def __init__(self, host: str, port: int, handler: SRIOTFTPHandler) -> None:
        self.host = host
        self.port = port
        self.handler = handler
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind((host, port))
        self.block_size = 512

    def send_error(self, error_code: int, error_msg: str, addr: tuple) -> None:
        """Send TFTP ERROR packet."""
        packet = struct.pack("!HH", self.OP_ERROR, error_code) + error_msg.encode() + b"\x00"
        self.sock.sendto(packet, addr)

    def send_ack(self, block_num: int, addr: tuple) -> None:
        """Send TFTP ACK packet."""
        packet = struct.pack("!HH", self.OP_ACK, block_num)
        self.sock.sendto(packet, addr)

    def send_data(self, block_num: int, data: bytes, addr: tuple) -> None:
        """Send TFTP DATA packet."""
        packet = struct.pack("!HH", self.OP_DATA, block_num) + data
        self.sock.sendto(packet, addr)

    def handle_rrq(self, filename: str, addr: tuple) -> None:
        """Handle Read Request (GET)."""
        print(f"RRQ from {addr}: {filename}")
        
        # Create dedicated socket for this transfer (RFC 1350)
        transfer_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        transfer_sock.bind(('0.0.0.0', 0))  # Use ephemeral port
        
        try:
            data = self.handler.handle_read(filename)
            if data is None:
                packet = struct.pack("!HH", self.OP_ERROR, 1) + b"File not found\x00"
                transfer_sock.sendto(packet, addr)
                return

            # Send data in 512-byte blocks
            block_num = 1
            offset = 0
            while offset < len(data):
                chunk = data[offset : offset + self.block_size]
                packet = struct.pack("!HH", self.OP_DATA, block_num) + chunk
                transfer_sock.sendto(packet, addr)

                # Wait for ACK with timeout
                transfer_sock.settimeout(5.0)
                try:
                    ack_packet, ack_addr = transfer_sock.recvfrom(4096)
                    opcode = struct.unpack("!H", ack_packet[:2])[0]
                    if opcode != self.OP_ACK:
                        print(f"  Expected ACK, got opcode {opcode}")
                        return
                    ack_block = struct.unpack("!H", ack_packet[2:4])[0]
                    if ack_block != block_num:
                        print(f"  ACK mismatch: expected {block_num}, got {ack_block}")
                        return
                except socket.timeout:
                    print(f"  Timeout waiting for ACK {block_num}")
                    return

                offset += len(chunk)
                block_num += 1

                # Last block (less than 512 bytes)
                if len(chunk) < self.block_size:
                    break

            print(f"  RRQ completed: {len(data)} bytes sent")
        finally:
            transfer_sock.close()

    def handle_wrq(self, filename: str, addr: tuple) -> None:
        """Handle Write Request (PUT)."""
        print(f"WRQ from {addr}: {filename}")

        # Create dedicated socket for this transfer (RFC 1350)
        transfer_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        transfer_sock.bind(('0.0.0.0', 0))  # Use ephemeral port

        try:
            # Send initial ACK (block 0)
            packet = struct.pack("!HH", self.OP_ACK, 0)
            transfer_sock.sendto(packet, addr)

            # Receive data blocks
            received_data = bytearray()
            expected_block = 1

            while True:
                transfer_sock.settimeout(10.0)
                try:
                    packet, data_addr = transfer_sock.recvfrom(4096)
                    opcode = struct.unpack("!H", packet[:2])[0]

                    if opcode == self.OP_DATA:
                        block_num = struct.unpack("!H", packet[2:4])[0]
                        data = packet[4:]

                        if block_num == expected_block:
                            received_data.extend(data)
                            ack_packet = struct.pack("!HH", self.OP_ACK, block_num)
                            transfer_sock.sendto(ack_packet, data_addr)
                            expected_block += 1

                            # Last block (less than 512 bytes)
                            if len(data) < self.block_size:
                                break
                        else:
                            print(f"  Block mismatch: expected {expected_block}, got {block_num}")
                            # Re-send last ACK
                            ack_packet = struct.pack("!HH", self.OP_ACK, expected_block - 1)
                            transfer_sock.sendto(ack_packet, data_addr)

                    elif opcode == self.OP_ERROR:
                        error_code = struct.unpack("!H", packet[2:4])[0]
                        error_msg = packet[4:].decode('ascii', errors='ignore').rstrip('\x00')
                        print(f"  Client error {error_code}: {error_msg}")
                        return

                except socket.timeout:
                    print(f"  Timeout waiting for DATA block {expected_block}")
                    return

            # Process received data
            success = self.handler.handle_write(filename, bytes(received_data))
            if success:
                print(f"  WRQ completed: {len(received_data)} bytes received")
            else:
                print(f"  WRQ failed")
        finally:
            transfer_sock.close()

    def handle_request(self, packet: bytes, addr: tuple) -> None:
        """Handle incoming TFTP request."""
        opcode = struct.unpack("!H", packet[:2])[0]

        if opcode == self.OP_RRQ:
            # Parse RRQ: opcode, filename\0, mode\0
            parts = packet[2:].split(b"\x00")
            if len(parts) >= 2:
                filename = parts[0].decode('ascii', errors='ignore')
                mode = parts[1].decode('ascii', errors='ignore')
                self.handle_rrq(filename, addr)

        elif opcode == self.OP_WRQ:
            # Parse WRQ: opcode, filename\0, mode\0
            parts = packet[2:].split(b"\x00")
            if len(parts) >= 2:
                filename = parts[0].decode('ascii', errors='ignore')
                mode = parts[1].decode('ascii', errors='ignore')
                self.handle_wrq(filename, addr)

    def serve_forever(self) -> None:
        """Main server loop."""
        print(f"TFTP server listening on {self.host}:{self.port}")
        print("Press Ctrl+C to stop.\n")

        while True:
            try:
                packet, addr = self.sock.recvfrom(4096)
                # Handle each request in a separate thread
                thread = threading.Thread(target=self.handle_request, args=(packet, addr), daemon=True)
                thread.start()
            except KeyboardInterrupt:
                print("\nShutting down TFTP server.")
                break
            except Exception as e:
                print(f"Error: {e}")


def main() -> None:
    parser = argparse.ArgumentParser(description="TFTP receiver emulator for SRIO protocol")
    parser.add_argument("--port", type=int, default=6969, help="TFTP listen port (default: 6969)")
    parser.add_argument(
        "--workdir",
        type=Path,
        default=Path("./tftp_root"),
        help="Working directory for TFTP files",
    )
    parser.add_argument(
        "--flash-image",
        type=Path,
        default=Path("./tftp_root/target_flash.img"),
        help="Flash memory image file",
    )
    parser.add_argument(
        "--flash-size",
        type=int,
        default=16 * 1024 * 1024,
        help="Flash size in bytes (default: 16 MB)",
    )
    args = parser.parse_args()

    workdir = args.workdir
    workdir.mkdir(parents=True, exist_ok=True)

    flash = FlashEmulator(args.flash_image, args.flash_size)
    handler = SRIOTFTPHandler(workdir, flash)

    print(f"Working directory: {workdir.absolute()}")
    print(f"Flash image: {args.flash_image.absolute()}")
    print(f"Flash size: {args.flash_size // (1024*1024)} MB\n")

    server = SimpleTFTPServer("0.0.0.0", args.port, handler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down TFTP server.")
        sys.exit(0)


if __name__ == "__main__":
    main()