"""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)
    """
    
    def __init__(self, flash_image_path: Path, flash_size: int = 16 * 1024 * 1024) -> None:
        """Initialize SRIO flash controller simulator."""
        self.flash_image_path = flash_image_path
        self.flash_size = flash_size
        
        # Initialize flash image file
        if not flash_image_path.exists():
            flash_image_path.parent.mkdir(parents=True, exist_ok=True)
            with open(flash_image_path, "wb") as f:
                f.write(b"\xFF" * flash_size)
        print(f"Flash controller initialized: {flash_image_path} ({flash_size} bytes)")
        
        # Register state (256-byte aligned blocks)
        self.memory: Dict[int, bytearray] = {}  # Stores 256-byte blocks
        
        # Flash operation state
        self.flash_busy = False
        self.tx_fifo_buffer: bytes = b""
        self.mode_reg = 0
        self.cmd_reg = 0
        self.addr_reg = 0
        self.num_bytes_reg = 0
    
    def _get_block_key(self, address: int) -> int:
        """Get 256-byte aligned block key."""
        return address & 0xFFFFFF00
    
    def _ensure_block(self, address: int) -> bytearray:
        """Ensure 256-byte block exists in memory."""
        key = self._get_block_key(address)
        if key not in self.memory:
            self.memory[key] = bytearray(256)
        return self.memory[key]
    
    def write_data(self, address: int, data: bytes) -> None:
        """Write data to SRIO address (256-byte aligned)."""
        if len(data) != 256:
            print(f"  WARNING: Write size {len(data)} != 256 bytes")
            # Pad or truncate
            if len(data) < 256:
                data = data + b"\x00" * (256 - len(data))
            else:
                data = data[:256]
        
        block = self._ensure_block(address)
        block[:] = data
        
        # Check if this is a register write
        if self._is_register_address(address):
            self._handle_register_write(address, data)
        elif address == TX_FIFO_REG:
            # TX FIFO data for flash write
            self.tx_fifo_buffer = data
            print(f"  TX_FIFO: Stored 256 bytes for flash write")
    
    def read_data(self, address: int, length: int) -> bytes:
        """Read data from SRIO address."""
        if length != 256:
            print(f"  WARNING: Read size {length} != 256 bytes")
        
        block = self._ensure_block(address)
        
        # Check if this is a register read
        if self._is_register_address(address):
            self._handle_register_read(address, block)
        
        return bytes(block[:length])
    
    def _is_register_address(self, address: int) -> bool:
        """Check if address is a register."""
        base = address & 0xFFFFFF00
        offset = address & 0xFF
        reg_bases = [
            MODE_REG & 0xFFFFFF00,
            CMD_REG & 0xFFFFFF00,
            CTRL_REG & 0xFFFFFF00,
            STATUS_REG & 0xFFFFFF00,
            RX_FIFO_REG & 0xFFFFFF00,
        ]
        return base in reg_bases or address == TX_FIFO_REG
    
    def _handle_register_write(self, address: int, data: bytes) -> None:
        """Process register write and update state."""
        base = address & 0xFFFFFF00
        
        # Extract register values from 256-byte block
        if base == (MODE_REG & 0xFFFFFF00):
            offset = MODE_REG & 0xFF
            self.mode_reg = struct.unpack('<I', data[offset:offset+4])[0]
            print(f"  MODE_REG = 0x{self.mode_reg:02X}")
        
        if base == (CMD_REG & 0xFFFFFF00):
            offset = CMD_REG & 0xFF
            self.cmd_reg = struct.unpack('<I', data[offset:offset+4])[0]
            print(f"  CMD_REG = 0x{self.cmd_reg:02X}")
        
        if base == (ADDR_REG & 0xFFFFFF00):
            offset = ADDR_REG & 0xFF
            self.addr_reg = struct.unpack('<I', data[offset:offset+4])[0]
            print(f"  ADDR_REG = 0x{self.addr_reg:08X}")
        
        if base == (NUM_BYTE_REG & 0xFFFFFF00):
            offset = NUM_BYTE_REG & 0xFF
            self.num_bytes_reg = struct.unpack('<I', data[offset:offset+4])[0]
            print(f"  NUM_BYTE_REG = {self.num_bytes_reg}")
        
        if base == (CTRL_REG & 0xFFFFFF00):
            offset = CTRL_REG & 0xFF
            ctrl_value = struct.unpack('<I', data[offset:offset+4])[0]
            self._handle_ctrl_write(ctrl_value)
    
    def _handle_ctrl_write(self, value: int) -> None:
        """Handle CTRL_REG write (triggers state machine)."""
        print(f"  CTRL_REG = 0x{value:02X}", end="")
        
        status_block = self._ensure_block(STATUS_REG)
        status_offset = STATUS_REG & 0xFF
        
        if value == 0x01:
            # START_REQUEST or part of WRITE_DATA sequence
            # Set STATUS[11] = 1 (busy/ack)
            status = struct.unpack('<I', status_block[status_offset:status_offset+4])[0]
            status |= (1 << 11)
            struct.pack_into('<I', status_block, status_offset, status)
            print(" → START (STATUS[11]=1)")
        elif value == 0x00:
            # END_REQUEST
            # Clear STATUS[11]
            status = struct.unpack('<I', status_block[status_offset:status_offset+4])[0]
            status &= ~(1 << 11)
            # Clear error bits [1,4,5,6,10]
            status &= ~((1 << 1) | (1 << 4) | (1 << 5) | (1 << 6) | (1 << 10))
            struct.pack_into('<I', status_block, status_offset, status)
            print(" → END (STATUS[11]=0, errors cleared)")
        elif value == 0x03:
            # WRITE_DATA / READ_DATA sequence
            # Set STATUS[12] = 1 (tx_fifo_empty / ready)
            status = struct.unpack('<I', status_block[status_offset:status_offset+4])[0]
            status |= (1 << 12)
            struct.pack_into('<I', status_block, status_offset, status)
            print(" → WRITE/READ_DATA (STATUS[12]=1)")
            
            # Execute flash operation if CMD is set
            if self.cmd_reg != 0:
                self._execute_flash_operation()
        else:
            print()
    
    def _handle_register_read(self, address: int, block: bytearray) -> None:
        """Prepare register read response."""
        base = address & 0xFFFFFF00
        
        # Read STATUS_REG
        if base == (STATUS_REG & 0xFFFFFF00):
            offset = STATUS_REG & 0xFF
            status = struct.unpack('<I', block[offset:offset+4])[0]
            print(f"  STATUS_REG = 0x{status:08X} (bit11={status>>11&1}, bit12={status>>12&1})")
        
        # Read RX_FIFO_REG (flash status or read data)
        if base == (RX_FIFO_REG & 0xFFFFFF00):
            offset = RX_FIFO_REG & 0xFF
            rx_data = struct.unpack('<I', block[offset:offset+4])[0]
            print(f"  RX_FIFO_REG = 0x{rx_data:08X}")
    
    def _execute_flash_operation(self) -> None:
        """Execute flash operation based on CMD_REG."""
        cmd = self.cmd_reg
        addr = self.addr_reg
        num_bytes = self.num_bytes_reg
        
        # Prepare RX_FIFO response
        rx_block = self._ensure_block(RX_FIFO_REG)
        rx_offset = RX_FIFO_REG & 0xFF
        
        if cmd == 0x06:
            print(f"    → Flash: Write Enable")
            # No action needed
        elif cmd == 0xD8:
            print(f"    → Flash: ERASE sector at 0x{addr:08X}")
            self._flash_erase(addr, 65536)
            # Set RX_FIFO[0] = 1 (busy), will clear on next wait_flash
            struct.pack_into('<I', rx_block, rx_offset, 0x01)
            self.flash_busy = True
        elif cmd == 0x02:
            print(f"    → Flash: WRITE 256 bytes to 0x{addr:08X}")
            self._flash_write(addr, self.tx_fifo_buffer)
            struct.pack_into('<I', rx_block, rx_offset, 0x01)
            self.flash_busy = True
        elif cmd == 0x0B:
            print(f"    → Flash: READ {num_bytes} bytes from 0x{addr:08X}")
            data = self._flash_read(addr, num_bytes)
            # Store first 4 bytes in RX_FIFO
            if len(data) >= 4:
                value = struct.unpack('<I', data[:4])[0]
                struct.pack_into('<I', rx_block, rx_offset, value)
            print(f"      RX_FIFO[0] = 0x{struct.unpack('<I', rx_block[rx_offset:rx_offset+4])[0]:08X}")
        elif cmd == 0x05:
            # Read flash status register
            if self.flash_busy:
                print(f"    → Flash: READ STATUS (busy=1)")
                struct.pack_into('<I', rx_block, rx_offset, 0x01)
                self.flash_busy = False  # Next read will show ready
            else:
                print(f"    → Flash: READ STATUS (ready=0)")
                struct.pack_into('<I', rx_block, rx_offset, 0x00)
        elif cmd == 0x61:
            print(f"    → Flash: Write Volatile Config Register")
        elif cmd == 0xB7:
            print(f"    → Flash: Enter 4-byte Address Mode")
        else:
            print(f"    → Flash: Unknown CMD 0x{cmd:02X}")
    
    def _flash_erase(self, address: int, length: int) -> None:
        """Erase flash sector (fill with 0xFF)."""
        try:
            with open(self.flash_image_path, "r+b") as f:
                f.seek(address)
                f.write(b"\xFF" * length)
        except Exception as e:
            print(f"      ERROR: {e}")
    
    def _flash_write(self, address: int, data: bytes) -> None:
        """Write to flash."""
        try:
            with open(self.flash_image_path, "r+b") as f:
                f.seek(address)
                f.write(data)
        except Exception as e:
            print(f"      ERROR: {e}")
    
    def _flash_read(self, address: int, length: int) -> bytes:
        """Read from flash."""
        try:
            with open(self.flash_image_path, "rb") as f:
                f.seek(address)
                return f.read(length)
        except Exception as e:
            print(f"      ERROR: {e}")
            return b"\xFF" * length


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:
    """TFTP handler that processes SRIO commands via flash controller."""

    def __init__(self, workdir: Path, flash_controller: SRIOFlashController) -> None:
        self.workdir = workdir
        self.flash_controller = flash_controller
        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_controller.read_data(address, length)
                return data
            except Exception as e:
                print(f"  ERROR: {e}")
                return None
        else:
            # Normal file read
            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_controller.write_data(address, data)
                return True
            except Exception as e:
                print(f"  ERROR: {e}")
                return False
        else:
            # Normal file write
            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."""

    # 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 handle_rrq(self, filename: str, addr: tuple) -> None:
        """Handle Read Request (GET)."""
        print(f"\nRRQ from {addr[0]}:{addr[1]}: {filename}")
        
        transfer_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        transfer_sock.bind(('0.0.0.0', 0))
        
        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)

                transfer_sock.settimeout(5.0)
                try:
                    ack_packet, _ = transfer_sock.recvfrom(4096)
                    opcode = struct.unpack("!H", ack_packet[:2])[0]
                    if opcode != self.OP_ACK:
                        return
                except socket.timeout:
                    print(f"  Timeout waiting for ACK {block_num}")
                    return

                offset += len(chunk)
                block_num += 1

                if len(chunk) < self.block_size:
                    break

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

    def handle_wrq(self, filename: str, addr: tuple) -> None:
        """Handle Write Request (PUT)."""
        print(f"\nWRQ from {addr[0]}:{addr[1]}: {filename}")

        transfer_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        transfer_sock.bind(('0.0.0.0', 0))

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

            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

                            if len(data) < self.block_size:
                                break

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

            success = self.handler.handle_write(filename, bytes(received_data))
            if success:
                print(f"  ✓ WRQ completed: {len(received_data)} bytes received\n")
        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:
            parts = packet[2:].split(b"\x00")
            if len(parts) >= 2:
                filename = parts[0].decode('ascii', errors='ignore')
                self.handle_rrq(filename, addr)

        elif opcode == self.OP_WRQ:
            parts = packet[2:].split(b"\x00")
            if len(parts) >= 2:
                filename = parts[0].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)
                thread = threading.Thread(target=self.handle_request, args=(packet, addr), daemon=True)
                thread.start()
            except KeyboardInterrupt:
                print("\nShutting down TFTP server.")
                break


def main() -> None:
    parser = argparse.ArgumentParser(description="TFTP receiver emulator for SRIO flash 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_controller = SRIOFlashController(args.flash_image, args.flash_size)
    handler = SRIOTFTPHandler(workdir, flash_controller)

    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("\nServer stopped.")
        sys.exit(0)


if __name__ == "__main__":
    main()