from core.crypto import CryptoContext
from core.device import Device
from core.registry import DeviceRegistry
from log.manager import LogManager
from utils.display import Display

from proto.c2_pb2 import Command, AgentMessage, AgentMsgType

# Forward declaration for type hinting to avoid circular import
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from cli.cli import CLI


class Transport:
    def __init__(self, registry: DeviceRegistry, logger: LogManager, cli_instance: 'CLI' = None):
        self.crypto = CryptoContext()
        self.registry = registry
        self.logger = logger
        self.cli = cli_instance # CLI instance for callback
        self.command_responses = {} # To track command responses

    def set_cli(self, cli_instance: 'CLI'):
        self.cli = cli_instance

    # ==================================================
    # RX  (ESP → C2)
    # ==================================================
    def handle_incoming(self, sock, addr, raw_data: bytes):
        """
        raw_data = BASE64( ChaCha20( Protobuf AgentMessage ) )
        """
        # Removed verbose transport debug prints

        # 1) base64 decode
        try:
            cipher = self.crypto.b64_decode(raw_data)
        except Exception as e:
            Display.error(f"Base64 decode failed from {addr}: {e}")
            return

        # 2) chacha decrypt
        try:
            protobuf_bytes = self.crypto.decrypt(cipher)
        except Exception as e:
            Display.error(f"Decrypt failed from {addr}: {e}")
            return

        # 3) protobuf decode → AgentMessage
        try:
            msg = AgentMessage.FromString(protobuf_bytes)
        except Exception as e:
            Display.error(f"Protobuf decode failed from {addr}: {e}")
            return

        if not msg.device_id:
            Display.error("AgentMessage received without device_id")
            return

        self._dispatch(sock, addr, msg)

    # ==================================================
    # DISPATCH
    # ==================================================
    def _dispatch(self, sock, addr, msg: AgentMessage):
        device = self.registry.get(msg.device_id)

        if not device:
            device = Device(
                id=msg.device_id,
                sock=sock,
                address=addr
            )
            self.registry.add(device)
            Display.device_event(device.id, f"Connected from {addr[0]}")
        else:
            device.touch()

        self._handle_agent_message(device, msg)

    # ==================================================
    # AGENT MESSAGE HANDLER
    # ==================================================
    def _handle_agent_message(self, device: Device, msg: AgentMessage):
        payload_str = ""
        if msg.payload:
            try:
                payload_str = msg.payload.decode(errors="ignore")
            except Exception:
                payload_str = repr(msg.payload)

        if msg.type == AgentMsgType.AGENT_CMD_RESULT:
            if msg.request_id and self.cli:
                self.cli.handle_command_response(msg.request_id, device.id, payload_str, msg.eof)
            else:
                Display.device_event(device.id, f"Command result (no request_id or CLI not set): {payload_str}")
        elif msg.type == AgentMsgType.AGENT_INFO:
            # Check for MLAT data (format: MLAT:x;y;rssi)
            if payload_str.startswith("MLAT:") and self.cli:
                mlat_data = payload_str[5:]  # Remove "MLAT:" prefix
                if self.cli.mlat_engine.parse_mlat_message(device.id, mlat_data):
                    # Recalculate position if we have enough scanners
                    state = self.cli.mlat_engine.get_state()
                    if state["scanners_count"] >= 3:
                        self.cli.mlat_engine.calculate_position()
                else:
                    Display.device_event(device.id, f"MLAT: Invalid data format: {mlat_data}")
            else:
                Display.device_event(device.id, f"INFO: {payload_str}")
        elif msg.type == AgentMsgType.AGENT_ERROR:
            Display.device_event(device.id, f"ERROR: {payload_str}")
        elif msg.type == AgentMsgType.AGENT_LOG:
            Display.device_event(device.id, f"LOG: {payload_str}")
        elif msg.type == AgentMsgType.AGENT_DATA:
            Display.device_event(device.id, f"DATA: {payload_str}")
        else:
            Display.device_event(device.id, f"UNKNOWN Message Type ({AgentMsgType.Name(msg.type)}): {payload_str}")

    # ==================================================
    # TX  (C2 → ESP)
    # ==================================================
    def send_command(self, sock, cmd: Command):
        """
        Command → Protobuf → ChaCha20 → Base64 → \\n
        """
        try:
            proto = cmd.SerializeToString()
            # Removed verbose transport debug prints

            # Encrypt
            cipher = self.crypto.encrypt(proto)

            # Base64
            b64 = self.crypto.b64_encode(cipher)

            sock.sendall(b64 + b"\n")

        except Exception as e:
            Display.error(f"Failed to send command to {cmd.device_id}: {e}")