"""
SOCKS5 Tunnel — servidor SOCKS5 local con multiplexing sobre WebSocket.

Implementa RFC 1928 mínimo (solo CONNECT, sin BIND ni UDP).
El tráfico se encapsula en frames JSON sobre el WebSocket existente
del AgentConnection, permitiendo que el sandbox cloud acceda a la red
del cliente a través del agente remoto.

Protocolo de frames:
  tunnel_open  → abre canal hacia target_host:target_port
  tunnel_data  → envía datos por un canal abierto (base64)
  tunnel_close → cierra un canal
  tunnel_error → notifica error en un canal
"""
import asyncio
import base64
import struct
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable, Coroutine

from jama_agent.logger import log

# --- Constantes SOCKS5 (RFC 1928) ---
SOCKS5_VERSION = 0x05
SOCKS5_NO_AUTH = 0x00
SOCKS5_CMD_CONNECT = 0x01
SOCKS5_ATYP_IPV4 = 0x01
SOCKS5_ATYP_DOMAIN = 0x03
SOCKS5_ATYP_IPV6 = 0x04
SOCKS5_REP_SUCCESS = 0x00
SOCKS5_REP_GENERAL_FAILURE = 0x01
SOCKS5_REP_NOT_ALLOWED = 0x02
SOCKS5_REP_NET_UNREACHABLE = 0x03
SOCKS5_REP_HOST_UNREACHABLE = 0x04
SOCKS5_REP_REFUSED = 0x05
SOCKS5_REP_CMD_NOT_SUPPORTED = 0x07

# Tamaño máximo de chunk para tunnel_data
MAX_CHUNK_SIZE = 32 * 1024  # 32 KB
# Timeout de inactividad por canal (segundos)
IDLE_TIMEOUT = 300  # 5 minutos


def parse_socks5_greeting(data: bytes) -> tuple[int, list[int]]:
    """Parsea el greeting inicial SOCKS5 (VER, NMETHODS, METHODS).

    Retorna (version, lista de métodos soportados).
    Lanza ValueError si el formato es inválido.
    """
    if len(data) < 3:
        raise ValueError("Greeting demasiado corto")
    version = data[0]
    nmethods = data[1]
    if len(data) < 2 + nmethods:
        raise ValueError(f"Greeting incompleto: esperaba {2 + nmethods} bytes, tiene {len(data)}")
    methods = list(data[2:2 + nmethods])
    return version, methods


def parse_socks5_request(data: bytes) -> tuple[int, int, str, int]:
    """Parsea un request SOCKS5 (VER, CMD, RSV, ATYP, DST.ADDR, DST.PORT).

    Retorna (cmd, atyp, target_host, target_port).
    Lanza ValueError si el formato es inválido.
    """
    if len(data) < 7:
        raise ValueError("Request demasiado corto")
    version = data[0]
    if version != SOCKS5_VERSION:
        raise ValueError(f"Versión SOCKS incorrecta: {version}")
    cmd = data[1]
    # data[2] es RSV (reservado)
    atyp = data[3]

    if atyp == SOCKS5_ATYP_IPV4:
        if len(data) < 10:
            raise ValueError("Request IPv4 incompleto")
        host = ".".join(str(b) for b in data[4:8])
        port = struct.unpack("!H", data[8:10])[0]
    elif atyp == SOCKS5_ATYP_DOMAIN:
        domain_len = data[4]
        if len(data) < 5 + domain_len + 2:
            raise ValueError("Request dominio incompleto")
        host = data[5:5 + domain_len].decode("ascii", errors="replace")
        port = struct.unpack("!H", data[5 + domain_len:7 + domain_len])[0]
    elif atyp == SOCKS5_ATYP_IPV6:
        if len(data) < 22:
            raise ValueError("Request IPv6 incompleto")
        addr_bytes = data[4:20]
        # Formato IPv6 estándar
        parts = []
        for i in range(0, 16, 2):
            parts.append(f"{addr_bytes[i]:02x}{addr_bytes[i+1]:02x}")
        host = ":".join(parts)
        port = struct.unpack("!H", data[20:22])[0]
    else:
        raise ValueError(f"ATYP no soportado: {atyp}")

    return cmd, atyp, host, port


def build_socks5_reply(rep: int, bind_addr: str = "0.0.0.0", bind_port: int = 0) -> bytes:
    """Construye un reply SOCKS5."""
    addr_parts = bind_addr.split(".")
    if len(addr_parts) == 4:
        addr_bytes = bytes(int(p) for p in addr_parts)
    else:
        addr_bytes = b"\x00\x00\x00\x00"
    return struct.pack(
        "!BBBB4sH",
        SOCKS5_VERSION,
        rep,
        0x00,  # RSV
        SOCKS5_ATYP_IPV4,
        addr_bytes,
        bind_port,
    )


@dataclass
class TunnelChannel:
    """Representa un canal de túnel bidireccional."""

    channel_id: str
    target_host: str
    target_port: int
    reader: asyncio.StreamReader | None = None
    writer: asyncio.StreamWriter | None = None
    # Cola para datos entrantes desde WS → SOCKS5 client
    incoming: asyncio.Queue[bytes | None] = field(default_factory=lambda: asyncio.Queue(maxsize=100))
    created_at: float = field(default_factory=lambda: asyncio.get_event_loop().time())
    last_activity: float = field(default_factory=lambda: asyncio.get_event_loop().time())
    closed: bool = False

    def touch(self) -> None:
        """Actualiza timestamp de última actividad."""
        try:
            self.last_activity = asyncio.get_event_loop().time()
        except RuntimeError:
            pass


class TunnelManager:
    """Gestiona canales de túnel multiplexados sobre WebSocket.

    Coordina entre conexiones SOCKS5 locales y el transporte WS.
    """

    def __init__(
        self,
        send_ws: Callable[[dict[str, Any]], Coroutine[Any, Any, None]],
        max_channels: int = 10,
    ) -> None:
        self._send_ws = send_ws
        self._max_channels = max_channels
        self._channels: dict[str, TunnelChannel] = {}
        self._lock = asyncio.Lock()

    @property
    def channel_count(self) -> int:
        return len(self._channels)

    @property
    def channels(self) -> dict[str, TunnelChannel]:
        return dict(self._channels)

    async def open_channel(
        self, target_host: str, target_port: int
    ) -> TunnelChannel | None:
        """Abre un nuevo canal de túnel. Retorna None si se excede el límite."""
        async with self._lock:
            if len(self._channels) >= self._max_channels:
                log.warning(
                    f"Límite de canales alcanzado ({self._max_channels})"
                )
                return None

            channel_id = str(uuid.uuid4())
            channel = TunnelChannel(
                channel_id=channel_id,
                target_host=target_host,
                target_port=target_port,
            )
            self._channels[channel_id] = channel

        # Notificar al backend que se abre un canal
        await self._send_ws({
            "type": "tunnel_open",
            "channel_id": channel_id,
            "target_host": target_host,
            "target_port": target_port,
        })

        log.info(
            f"Canal abierto: {channel_id[:8]}... → {target_host}:{target_port}"
        )
        return channel

    async def close_channel(self, channel_id: str, reason: str = "EOF") -> None:
        """Cierra un canal y notifica al peer."""
        async with self._lock:
            channel = self._channels.pop(channel_id, None)

        if not channel or channel.closed:
            return

        channel.closed = True

        # Señalar fin al queue
        try:
            channel.incoming.put_nowait(None)
        except asyncio.QueueFull:
            pass

        # Cerrar el stream local si existe
        if channel.writer and not channel.writer.is_closing():
            try:
                channel.writer.close()
                await channel.writer.wait_closed()
            except Exception:
                pass

        # Notificar al peer
        try:
            await self._send_ws({
                "type": "tunnel_close",
                "channel_id": channel_id,
                "reason": reason,
            })
        except Exception:
            pass

        log.info(f"Canal cerrado: {channel_id[:8]}... ({reason})")

    async def send_data(self, channel_id: str, data: bytes) -> bool:
        """Envía datos por un canal hacia el peer WS (base64)."""
        channel = self._channels.get(channel_id)
        if not channel or channel.closed:
            return False

        channel.touch()

        # Enviar en chunks si es necesario
        offset = 0
        while offset < len(data):
            chunk = data[offset:offset + MAX_CHUNK_SIZE]
            await self._send_ws({
                "type": "tunnel_data",
                "channel_id": channel_id,
                "data": base64.b64encode(chunk).decode("ascii"),
            })
            offset += MAX_CHUNK_SIZE

        return True

    async def receive_data(self, channel_id: str, data: bytes) -> None:
        """Recibe datos del peer WS y los encola para el SOCKS5 client."""
        channel = self._channels.get(channel_id)
        if not channel or channel.closed:
            return

        channel.touch()

        try:
            channel.incoming.put_nowait(data)
        except asyncio.QueueFull:
            log.warning(f"Cola llena para canal {channel_id[:8]}...")
            await self.close_channel(channel_id, "queue_full")

    async def handle_ws_message(self, message: dict[str, Any]) -> None:
        """Procesa un mensaje tunnel_* recibido del WS."""
        msg_type = message.get("type", "")
        channel_id = message.get("channel_id", "")

        if msg_type == "tunnel_data":
            raw_data = message.get("data", "")
            try:
                decoded = base64.b64decode(raw_data)
            except Exception:
                log.warning(f"Datos base64 inválidos en canal {channel_id[:8]}...")
                return
            await self.receive_data(channel_id, decoded)

        elif msg_type == "tunnel_close":
            reason = message.get("reason", "peer_close")
            await self.close_channel(channel_id, reason)

        elif msg_type == "tunnel_error":
            error = message.get("error", "Unknown error")
            log.warning(f"Error en canal {channel_id[:8]}...: {error}")
            await self.close_channel(channel_id, f"error: {error}")

        elif msg_type == "tunnel_open_ack":
            # Backend confirma apertura — el canal ya está creado localmente
            log.debug(f"Canal {channel_id[:8]}... confirmado por backend")

        else:
            log.warning(f"Mensaje tunnel desconocido: {msg_type}")

    async def close_all(self, reason: str = "shutdown") -> None:
        """Cierra todos los canales activos."""
        channel_ids = list(self._channels.keys())
        for cid in channel_ids:
            await self.close_channel(cid, reason)

    async def cleanup_idle(self, timeout: float = IDLE_TIMEOUT) -> int:
        """Cierra canales inactivos. Retorna cantidad cerrados."""
        try:
            now = asyncio.get_event_loop().time()
        except RuntimeError:
            return 0

        idle_ids = []
        for cid, channel in list(self._channels.items()):
            if now - channel.last_activity > timeout:
                idle_ids.append(cid)

        for cid in idle_ids:
            await self.close_channel(cid, "idle_timeout")

        if idle_ids:
            log.info(f"Cerrados {len(idle_ids)} canales por inactividad")
        return len(idle_ids)


class Socks5Server:
    """Servidor SOCKS5 local que encapsula tráfico en canales de túnel.

    Escucha conexiones SOCKS5 en localhost, y por cada CONNECT:
    1. Abre un canal en TunnelManager
    2. Reenvía datos bidireccionalmente entre el socket local y el canal WS
    """

    def __init__(
        self,
        tunnel_manager: TunnelManager,
        host: str = "127.0.0.1",
        port: int = 1080,
    ) -> None:
        self._tunnel = tunnel_manager
        self._host = host
        self._port = port
        self._server: asyncio.Server | None = None
        self._cleanup_task: asyncio.Task[Any] | None = None

    @property
    def is_running(self) -> bool:
        return self._server is not None and self._server.is_serving()

    @property
    def address(self) -> tuple[str, int]:
        return self._host, self._port

    async def start(self) -> None:
        """Inicia el servidor SOCKS5."""
        self._server = await asyncio.start_server(
            self._handle_client,
            self._host,
            self._port,
        )
        # Tarea de limpieza periódica de canales idle
        self._cleanup_task = asyncio.create_task(self._idle_cleanup_loop())
        log.info(f"SOCKS5 server escuchando en {self._host}:{self._port}")

    async def stop(self) -> None:
        """Detiene el servidor SOCKS5 y cierra todos los canales."""
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
            self._cleanup_task = None

        if self._server:
            self._server.close()
            await self._server.wait_closed()
            self._server = None

        await self._tunnel.close_all("server_stop")
        log.info("SOCKS5 server detenido")

    async def _idle_cleanup_loop(self) -> None:
        """Loop periódico para limpiar canales inactivos."""
        while True:
            await asyncio.sleep(60)  # Verificar cada minuto
            await self._tunnel.cleanup_idle()

    async def _handle_client(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        """Maneja una conexión SOCKS5 entrante."""
        peer = writer.get_extra_info("peername", ("?", 0))
        log.debug(f"SOCKS5 cliente conectado desde {peer}")

        try:
            # 1. Greeting (negociación de método auth)
            greeting_data = await asyncio.wait_for(reader.read(257), timeout=10)
            if not greeting_data:
                return

            version, methods = parse_socks5_greeting(greeting_data)
            if version != SOCKS5_VERSION:
                log.warning(f"Versión SOCKS no soportada: {version}")
                return

            if SOCKS5_NO_AUTH not in methods:
                # No soportamos auth — rechazar
                writer.write(bytes([SOCKS5_VERSION, 0xFF]))
                await writer.drain()
                return

            # Responder: NO_AUTH
            writer.write(bytes([SOCKS5_VERSION, SOCKS5_NO_AUTH]))
            await writer.drain()

            # 2. Request (CONNECT)
            request_data = await asyncio.wait_for(reader.read(262), timeout=10)
            if not request_data:
                return

            cmd, _atyp, target_host, target_port = parse_socks5_request(request_data)

            if cmd != SOCKS5_CMD_CONNECT:
                writer.write(build_socks5_reply(SOCKS5_REP_CMD_NOT_SUPPORTED))
                await writer.drain()
                return

            # 3. Abrir canal de túnel
            channel = await self._tunnel.open_channel(target_host, target_port)
            if not channel:
                writer.write(build_socks5_reply(SOCKS5_REP_NOT_ALLOWED))
                await writer.drain()
                return

            # Asociar streams al canal
            channel.reader = reader
            channel.writer = writer

            # Responder éxito
            writer.write(build_socks5_reply(SOCKS5_REP_SUCCESS))
            await writer.drain()

            log.info(
                f"SOCKS5 CONNECT {target_host}:{target_port} "
                f"→ canal {channel.channel_id[:8]}..."
            )

            # 4. Relay bidireccional
            await self._relay(channel, reader, writer)

        except asyncio.TimeoutError:
            log.warning(f"Timeout en handshake SOCKS5 desde {peer}")
        except ValueError as exc:
            log.warning(f"Error parsing SOCKS5: {exc}")
        except Exception as exc:
            log.error(f"Error en cliente SOCKS5: {exc}")
        finally:
            if not writer.is_closing():
                try:
                    writer.close()
                    await writer.wait_closed()
                except Exception:
                    pass

    async def _relay(
        self,
        channel: TunnelChannel,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        """Relay bidireccional: SOCKS5 socket ↔ tunnel channel."""

        async def local_to_ws() -> None:
            """Lee del socket local y envía por WS."""
            try:
                while not channel.closed:
                    data = await asyncio.wait_for(
                        reader.read(MAX_CHUNK_SIZE),
                        timeout=IDLE_TIMEOUT,
                    )
                    if not data:
                        break
                    await self._tunnel.send_data(channel.channel_id, data)
            except asyncio.TimeoutError:
                pass
            except (ConnectionError, asyncio.CancelledError):
                pass
            finally:
                await self._tunnel.close_channel(channel.channel_id, "client_eof")

        async def ws_to_local() -> None:
            """Lee del canal WS y escribe al socket local."""
            try:
                while not channel.closed:
                    data = await asyncio.wait_for(
                        channel.incoming.get(),
                        timeout=IDLE_TIMEOUT,
                    )
                    if data is None:
                        break
                    if not writer.is_closing():
                        writer.write(data)
                        await writer.drain()
            except asyncio.TimeoutError:
                pass
            except (ConnectionError, asyncio.CancelledError):
                pass

        # Ejecutar ambas direcciones en paralelo
        tasks = [
            asyncio.create_task(local_to_ws()),
            asyncio.create_task(ws_to_local()),
        ]
        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for task in pending:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
