# target_simulator/core/sfp_transport.py

"""
Provides a reusable transport layer for the Simple Fragmentation Protocol (SFP).

This module handles UDP socket communication, SFP header parsing, fragment
reassembly, ACK generation, and payload sending. It is application-agnostic
and uses a callback/handler system to pass fully reassembled payloads to the
application layer based on the SFP_FLOW identifier.
"""

import socket
import logging
import threading
import time
import ctypes
from typing import Dict, Callable, Optional

from target_simulator.utils.network import create_udp_socket, close_udp_socket
from target_simulator.core.sfp_structures import SFPHeader

PayloadHandler = Callable[[bytearray], None]

logger = logging.getLogger(__name__)


class SfpTransport:
    """Manages SFP communication, payload reassembly, and sending."""

    # Max size for a script payload, conservative to fit in server buffers.
    MAX_SCRIPT_PAYLOAD_SIZE = 1020
    # Size of the fixed part of the script message (the DataTag).
    SCRIPT_TAG_SIZE = 8

    def __init__(
        self,
        host: str,
        port: int,
        payload_handlers: Dict[int, PayloadHandler],
        ack_config: Optional[Dict[int, int]] = None,
        raw_packet_callback: Optional[Callable[[bytes, tuple], None]] = None,
    ):
        self._log_prefix = "[SfpTransport]"
        logger.info(f"{self._log_prefix} Initializing for {host}:{port}...")

        self._host = host
        self._port = port
        self._payload_handlers = payload_handlers
        self._ack_config = ack_config if ack_config is not None else {}
        self._raw_packet_callback = raw_packet_callback

        self._socket: Optional[socket.socket] = None
        self._receiver_thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()
        self._tid_counter = 0  # Simple transaction ID counter for sending
        self._send_lock = threading.Lock()  # Protects TID counter and socket sending

        # transaction state: key=(flow, tid) -> {frag_index: total_frags}
        self._fragments: Dict[tuple, Dict[int, int]] = {}
        # buffers for reassembly: key=(flow, tid) -> bytearray(total_size)
        self._buffers: Dict[tuple, bytearray] = {}

        logger.debug(
            f"{self._log_prefix} Registered handlers for flows: "
            f"{[chr(k) if 32 <= k <= 126 else k for k in self._payload_handlers.keys()]}"
        )
        logger.debug(f"{self._log_prefix} ACK window config: {self._ack_config}")

    def start(self) -> bool:
        """Starts the receiving thread."""
        if self._receiver_thread is not None and self._receiver_thread.is_alive():
            logger.warning(
                f"{self._log_prefix} Start called, but receiver is already running."
            )
            return True

        self._socket = create_udp_socket(self._host, self._port)
        if not self._socket:
            logger.critical(
                f"{self._log_prefix} Failed to create and bind socket. Cannot start."
            )
            return False

        self._stop_event.clear()
        self._receiver_thread = threading.Thread(
            target=self._receive_loop, name="SfpTransportThread", daemon=True
        )
        self._receiver_thread.start()
        return True

    def shutdown(self):
        """Stops the receiving thread and closes the socket."""
        self._stop_event.set()

        if self._socket:
            close_udp_socket(self._socket)
            self._socket = None

        if self._receiver_thread and self._receiver_thread.is_alive():
            logger.debug(f"{self._log_prefix} Waiting for receiver thread to join...")
            self._receiver_thread.join(timeout=2.0)
            if self._receiver_thread.is_alive():
                logger.warning(
                    f"{self._log_prefix} Receiver thread did not join cleanly."
                )

        logger.info(f"{self._log_prefix} Shutdown complete.")

    def send_script_command(
        self, command_string: str, destination: tuple, flow_id: int = ord("R")
    ) -> bool:
        """
        Encapsulates and sends a text command as a single-fragment SFP packet.

        This method constructs a payload similar to the server's `script_message_t`,
        wraps it in an SFP header, and sends it to the specified destination.

        Args:
            command_string: The text command to send (e.g., "tgtinit ...").
            destination: A tuple (ip, port) for the destination.
            flow_id: The SFP flow ID to use for this message.

        Returns:
            True if the packet was sent successfully, False otherwise.
        """
        log_prefix = f"{self._log_prefix} Send"
        if not self._socket:
            logger.error(f"{log_prefix} Cannot send: socket is not open.")
            return False

        # Define the expected payload structures using ctypes
        # The server expects script_payload_t which contains a ctrl_tag and ctrl[32]
        # before the script_tag, so replicate that layout here.
        class LocalDataTag(ctypes.Structure):
            _pack_ = 1
            _fields_ = [
                ("ID", ctypes.c_uint8 * 2),
                ("VALID", ctypes.c_uint8),
                ("VERSION", ctypes.c_uint8),
                ("SIZE", ctypes.c_uint32),
            ]

        class ScriptPayload(ctypes.Structure):
            _pack_ = 1
            _fields_ = [
                ("ctrl_tag", LocalDataTag),
                ("ctrl", ctypes.c_uint32 * 32),
                ("script_tag", LocalDataTag),
                ("script", ctypes.c_uint8 * self.MAX_SCRIPT_PAYLOAD_SIZE),
            ]
        try:
            # Normalize command: ensure it starts with '$' (no blank after $) and ends with newline
            cs = command_string or ""
            cs = cs.strip()
            #if cs and not cs.startswith("$"):
            #    cs = "$" + cs.lstrip()
            if cs and not cs.endswith("\n"):
                cs = cs + "\n"

            # (no transformation) send cs as-is; server-side now accepts spaces
            command_bytes = cs.encode("utf-8")
            if len(command_bytes) > self.MAX_SCRIPT_PAYLOAD_SIZE:
                logger.error(
                    f"{log_prefix} Command is too long ({len(command_bytes)} bytes). "
                    f"Max is {self.MAX_SCRIPT_PAYLOAD_SIZE}."
                )
                return False

            # Create and populate the script payload structure
            payload_struct = ScriptPayload()
            payload_struct.script_tag.ID[0] = ord("C")
            payload_struct.script_tag.ID[1] = ord("S")
            payload_struct.script_tag.VALID = 1
            payload_struct.script_tag.VERSION = 1
            payload_struct.script_tag.SIZE = len(command_bytes)
            ctypes.memmove(
                payload_struct.script, command_bytes, len(command_bytes)
            )

            payload_bytes = bytes(payload_struct)
            # Compute the offset of the script buffer within the payload structure
            try:
                script_offset = ScriptPayload.script.offset
            except Exception:
                # Fallback: assume script tag only (legacy)
                script_offset = self.SCRIPT_TAG_SIZE

            actual_payload_size = script_offset + len(command_bytes)
            payload_bytes = payload_bytes[:actual_payload_size]

            # Create and populate the SFP header
            header = SFPHeader()
            with self._send_lock:
                self._tid_counter = (self._tid_counter + 1) % 256
                header.SFP_TID = self._tid_counter

            header.SFP_DIRECTION = ord(">")
            header.SFP_FLOW = flow_id
            header.SFP_TOTFRGAS = 1  # Single fragment message
            header.SFP_FRAG = 0
            header.SFP_PLSIZE = len(payload_bytes)
            header.SFP_TOTSIZE = len(payload_bytes)
            header.SFP_PLOFFSET = 0

            full_packet = bytes(header) + payload_bytes

            # Send the packet
            self._socket.sendto(full_packet, destination)
            # Log the actual normalized command that was placed into the payload
            try:
                sent_preview = cs if isinstance(cs, str) else cs.decode('utf-8', errors='replace')
            except Exception:
                sent_preview = repr(cs)
            logger.info(
                f"{log_prefix} Sent command to {destination} (TID: {header.SFP_TID}): {sent_preview!r}"
            )
            return True

        except Exception as e:
            logger.exception(f"{log_prefix} Failed to send script command.")
            return False

    def _receive_loop(self):
        log_prefix = f"{self._log_prefix} Loop"
        logger.info(f"{log_prefix} Starting receive loop.")

        while not self._stop_event.is_set():
            if not self._socket:
                logger.error(f"{log_prefix} Socket is not available. Stopping loop.")
                break

            try:
                data, addr = self._socket.recvfrom(65535)
                if not data:
                    continue
                if self._raw_packet_callback:
                    self._raw_packet_callback(data, addr)
            except socket.timeout:
                continue
            except OSError as e:
                if not self._stop_event.is_set():
                    logger.error(f"{log_prefix} Socket error: {e}")
                break
            except Exception:
                logger.exception(f"{log_prefix} Unexpected error in recvfrom.")
                time.sleep(0.01)
                continue

            self._process_packet(data, addr)

        logger.info(f"{log_prefix} Receive loop terminated.")

    def _process_packet(self, data: bytes, addr: tuple):
        header_size = SFPHeader.size()
        if len(data) < header_size:
            logger.warning(
                f"Packet from {addr} is too small for SFP header. Ignoring."
            )
            return

        try:
            header = SFPHeader.from_buffer_copy(data)
        except (ValueError, TypeError):
            logger.error(f"Failed to parse SFP header from {addr}. Ignoring.")
            return

        flow = header.SFP_FLOW
        tid = header.SFP_TID
        frag = header.SFP_FRAG
        total_frags = header.SFP_TOTFRGAS
        pl_size = header.SFP_PLSIZE
        pl_offset = header.SFP_PLOFFSET
        total_size = header.SFP_TOTSIZE
        flags = header.SFP_FLAGS

        key = (flow, tid)

        if header.SFP_FLAGS & 0x01:
            self._send_ack(addr, data[:header_size])

        if total_frags == 0 or total_frags > 60000 or total_size <= 0:
            logger.warning(
                f"Invalid metadata for {key}: total_frags={total_frags}, "
                f"total_size={total_size}. Ignoring."
            )
            return

        if frag == 0:
            self._cleanup_lingering_transactions(flow, tid)
            logger.debug(
                f"New transaction started for key={key}. Total size: {total_size} bytes."
            )
            self._fragments[key] = {}
            try:
                self._buffers[key] = bytearray(total_size)
            except (MemoryError, ValueError):
                logger.error(
                    f"Failed to allocate {total_size} bytes for key={key}. Ignoring transaction."
                )
                self._fragments.pop(key, None)
                return

        if key not in self._buffers or key not in self._fragments:
            logger.debug(
                f"Ignoring fragment {frag} for untracked transaction key={key}."
            )
            return

        self._fragments[key][frag] = total_frags
        payload = data[header_size:]
        bytes_to_copy = min(pl_size, len(payload))

        if (pl_offset + bytes_to_copy) > len(self._buffers[key]):
            logger.error(
                f"Payload for key={key}, frag={frag} would overflow buffer. Ignoring."
            )
            return

        self._buffers[key][pl_offset : pl_offset + bytes_to_copy] = payload[
            :bytes_to_copy
        ]

        if len(self._fragments[key]) == total_frags:
            logger.debug(
                f"Transaction complete for key={key}. Handing off to application layer."
            )

            completed_payload = self._buffers.pop(key)
            self._fragments.pop(key)

            handler = self._payload_handlers.get(flow)
            if handler:
                try:
                    handler(completed_payload)
                except Exception:
                    logger.exception(
                        f"Error executing payload handler for flow {flow}."
                    )
            else:
                logger.warning(f"No payload handler registered for flow ID {flow}.")

    def _send_ack(self, dest_addr: tuple, original_header_bytes: bytes):
        log_prefix = f"{self._log_prefix} ACK"
        if not self._socket:
            return

        try:
            ack_header = bytearray(original_header_bytes)
            flow = ack_header[SFPHeader.get_field_offset("SFP_FLOW")]
            window_size = self._ack_config.get(flow, 0)

            ack_header[SFPHeader.get_field_offset("SFP_DIRECTION")] = ord("<")
            ack_header[SFPHeader.get_field_offset("SFP_WIN")] = window_size
            original_flags = ack_header[SFPHeader.get_field_offset("SFP_FLAGS")]
            ack_header[SFPHeader.get_field_offset("SFP_FLAGS")] = (
                original_flags | 0x01
            ) & ~0x02

            self._socket.sendto(ack_header, dest_addr)
            flow_char = chr(flow) if 32 <= flow <= 126 else flow
            logger.debug(
                f"{log_prefix} Sent to {dest_addr} for flow {flow_char} with WIN={window_size}."
            )
        except Exception:
            logger.exception(f"{log_prefix} Failed to send to {dest_addr}.")

    def _cleanup_lingering_transactions(self, current_flow: int, current_tid: int):
        """Remove old, incomplete transactions for the same flow."""
        keys_to_remove = [
            key
            for key in list(self._fragments.keys())
            if key[0] == current_flow and key[1] != current_tid
        ]
        for key in keys_to_remove:
            logger.warning(
                f"Cleaning up lingering/incomplete transaction for key={key}."
            )
            self._fragments.pop(key, None)
            self._buffers.pop(key, None)