"""
Conexión WebSocket persistente del Agent CLI al servidor cloud.

Mantiene un canal bidireccional con heartbeats, reconexión automática
con exponential backoff y procesamiento de tareas remotas.
"""
import asyncio
import json
import platform
import time
from typing import Any, Callable

import psutil  # type: ignore
import websockets
from websockets.exceptions import (
    ConnectionClosed,
    ConnectionClosedError,
    InvalidStatus,
)

from jama_agent.logger import log
from jama_agent.result_collector import ResultCollector
from jama_agent.task_executor import execute_task
from jama_agent.tunnel import TunnelManager


class AgentConnection:
    """Conexión WebSocket persistente al servidor Jama.

    Gestiona heartbeats, reconexión automática con exponential backoff,
    ejecución de tareas remotas, envío de resultados y túneles SOCKS5.
    """

    def __init__(
        self,
        server: str,
        agent_id: str,
        token: str,
        heartbeat_interval: int = 30,
        reconnect_max_delay: int = 60,
        on_disconnect: Callable[[], None] | None = None,
    ) -> None:
        self.server = server
        self.agent_id = agent_id
        self.token = token
        self.heartbeat_interval = heartbeat_interval
        self.reconnect_max_delay = reconnect_max_delay
        self.on_disconnect = on_disconnect

        self._ws: websockets.WebSocketClientProtocol | None = None
        self._running = False
        self._connected = False
        self._start_time: float = 0.0
        self._reconnect_attempt = 0
        self._collector = ResultCollector()
        self._tasks: dict[str, asyncio.Task[Any]] = {}
        self._tunnel_manager: TunnelManager | None = None
        self._allow_insecure: bool = False

    @property
    def is_connected(self) -> bool:
        """Indica si la conexión está activa."""
        return self._connected

    @property
    def uptime(self) -> float:
        """Segundos desde que se inició la conexión."""
        if self._start_time == 0:
            return 0.0
        return time.monotonic() - self._start_time

    def _build_ws_url(self) -> str:
        """Construye la URL del WebSocket.

        En producción (no-localhost), requiere wss:// a menos que
        se haya habilitado --insecure.
        """
        base = self.server.rstrip("/")
        # Convertir http→ws, https→wss
        if base.startswith("https://"):
            ws_base = "wss://" + base[8:]
        elif base.startswith("http://"):
            ws_base = "ws://" + base[7:]
        else:
            ws_base = "ws://" + base

        # TLS enforcement: rechazar ws:// en producción
        if ws_base.startswith("ws://") and not self._allow_insecure:
            host = ws_base.split("//", 1)[1].split("/")[0].split(":")[0]
            if host not in ("localhost", "127.0.0.1", "::1"):
                raise ConnectionError(
                    f"Conexión insegura (ws://) a host remoto '{host}' no permitida. "
                    "Usa https:// o --insecure para override."
                )

        return f"{ws_base}/ws/agent/{self.agent_id}?token={self.token}"

    def _get_reconnect_delay(self) -> float:
        """Calcula el delay de reconexión con exponential backoff.

        Secuencia: 1, 2, 4, 8, 16, 32, 60, 60, 60...
        """
        delay = min(2**self._reconnect_attempt, self.reconnect_max_delay)
        return float(max(1, delay))

    async def connect(self) -> None:
        """Inicia la conexión WebSocket y los loops de heartbeat y mensajes."""
        self._running = True
        self._start_time = time.monotonic()

        while self._running:
            try:
                url = self._build_ws_url()
                log.info(f"Conectando a {self.server}...")

                extra_headers = {}
                if self.token:
                    extra_headers["Authorization"] = f"Bearer {self.token}"

                async with websockets.connect(
                    url,
                    additional_headers=extra_headers,
                    ping_interval=None,  # Gestionamos pings manualmente
                    close_timeout=5,
                ) as ws:
                    self._ws = ws
                    self._connected = True
                    self._reconnect_attempt = 0
                    log.info(
                        f"Conectado al servidor (agent_id={self.agent_id})"
                    )

                    # Enviar resultados pendientes del buffer
                    await self._flush_pending_results()

                    # Correr heartbeat y message handler en paralelo
                    heartbeat_task = asyncio.create_task(self._heartbeat_loop())
                    handler_task = asyncio.create_task(self._message_handler())

                    done, pending = await asyncio.wait(
                        [heartbeat_task, handler_task],
                        return_when=asyncio.FIRST_COMPLETED,
                    )

                    # Cancelar tareas pendientes
                    for task in pending:
                        task.cancel()
                        try:
                            await task
                        except asyncio.CancelledError:
                            pass

            except (ConnectionClosed, ConnectionClosedError) as exc:
                log.warning(f"Conexión cerrada: {exc}")
            except (OSError, InvalidStatus) as exc:
                log.warning(f"Error de conexión: {exc}")
            except asyncio.CancelledError:
                log.info("Conexión cancelada")
                break
            except Exception as exc:
                log.error(f"Error inesperado: {exc}")

            self._connected = False
            self._ws = None

            if not self._running:
                break

            # Reconexión con exponential backoff
            await self._reconnect()

    async def _reconnect(self) -> None:
        """Espera con exponential backoff antes de reconectar."""
        delay = self._get_reconnect_delay()
        self._reconnect_attempt += 1
        log.info(
            f"Reconectando en {delay:.0f}s "
            f"(intento {self._reconnect_attempt})..."
        )
        await asyncio.sleep(delay)

    async def _heartbeat_loop(self) -> None:
        """Envía heartbeats periódicos al servidor."""
        while self._running and self._connected:
            try:
                await self._send_heartbeat()
                await asyncio.sleep(self.heartbeat_interval)
            except (ConnectionClosed, asyncio.CancelledError):
                break
            except Exception as exc:
                log.warning(f"Error en heartbeat: {exc}")
                break

    async def _send_heartbeat(self) -> None:
        """Envía un mensaje pong con métricas del sistema."""
        if not self._ws:
            return

        # Recoger métricas del sistema
        try:
            cpu = psutil.cpu_percent(interval=0)
            memory = psutil.virtual_memory().percent
        except Exception:
            cpu = 0.0
            memory = 0.0

        message = {
            "type": "pong",
            "uptime": int(self.uptime),
            "cpu": cpu,
            "memory": memory,
            "os": platform.system().lower(),
            "hostname": platform.node(),
        }

        await self._ws.send(json.dumps(message))

    async def _message_handler(self) -> None:
        """Procesa mensajes recibidos del servidor."""
        if not self._ws:
            return

        async for raw_message in self._ws:
            if not self._running:
                break

            try:
                message = json.loads(raw_message)
            except json.JSONDecodeError:
                log.warning(f"Mensaje no válido recibido: {raw_message!r}")
                continue

            await self._handle_message(message)

    async def _handle_message(self, message: dict[str, Any]) -> None:
        """Despacha un mensaje según su tipo."""
        msg_type = message.get("type", "unknown")

        if msg_type == "execute":
            await self._handle_execute(message)
        elif msg_type == "ping":
            await self._send_heartbeat()
        elif msg_type == "kill":
            log.warning("Recibido comando KILL del servidor")
            self._running = False
            await self.disconnect()
        elif msg_type.startswith("tunnel_"):
            # Delegar al TunnelManager
            if self._tunnel_manager:
                await self._tunnel_manager.handle_ws_message(message)
            else:
                log.warning("Mensaje tunnel recibido pero no hay tunnel_manager")
        else:
            log.warning(f"Mensaje desconocido: {msg_type}")

    async def _handle_execute(self, message: dict[str, Any]) -> None:
        """Ejecuta una tarea recibida del servidor."""
        task_id = message.get("task_id", "")
        tool = message.get("tool", "")
        args = message.get("args", {})
        timeout = message.get("timeout", 300)

        if not task_id or not tool:
            log.warning(f"Tarea incompleta: {message}")
            return

        log.info(f"Tarea recibida: {task_id} → {tool}")

        # Ejecutar en background para no bloquear el message handler
        task = asyncio.create_task(
            self._execute_and_report(task_id, tool, args, timeout)
        )
        self._tasks[task_id] = task

    async def _execute_and_report(
        self,
        task_id: str,
        tool: str,
        args: dict[str, Any],
        timeout: int,
    ) -> None:
        """Ejecuta una tarea y envía el resultado al servidor."""
        try:
            result = await execute_task(task_id, tool, args, timeout)

            response = {
                "type": "result",
                "task_id": result.task_id,
                "status": result.status,
                "output": result.output,
                "stderr": result.stderr,
                "exit_code": result.exit_code,
                "duration_ms": result.duration_ms,
            }

            if self._connected and self._ws:
                try:
                    await self._ws.send(json.dumps(response))
                    log.info(f"Resultado enviado: {task_id} ({result.status})")
                except Exception:
                    # Si falla, guardar en buffer local
                    self._collector.add(response)
                    log.warning(f"Resultado buffereado: {task_id}")
            else:
                self._collector.add(response)
                log.warning(f"Sin conexión, resultado buffereado: {task_id}")

        except Exception as exc:
            error_response = {
                "type": "error",
                "task_id": task_id,
                "message": str(exc),
            }
            if self._connected and self._ws:
                try:
                    await self._ws.send(json.dumps(error_response))
                except Exception:
                    self._collector.add(error_response)
            else:
                self._collector.add(error_response)
        finally:
            self._tasks.pop(task_id, None)

    async def _flush_pending_results(self) -> None:
        """Envía resultados pendientes del buffer al servidor."""
        pending = self._collector.get_pending()
        if not pending:
            return

        log.info(f"Enviando {len(pending)} resultados pendientes...")
        sent_ids: list[str] = []

        for result in pending:
            if self._ws:
                try:
                    await self._ws.send(json.dumps(result))
                    task_id = result.get("task_id", "unknown")
                    sent_ids.append(task_id)
                except Exception:
                    break  # Parar si falla — lo intentaremos en la próxima conexión

        # Limpiar los que se enviaron
        for task_id in sent_ids:
            self._collector.remove(task_id)

        if sent_ids:
            log.info(f"Enviados {len(sent_ids)} resultados pendientes")

    async def _send_ws_message(self, message: dict[str, Any]) -> None:
        """Envía un mensaje JSON por WebSocket."""
        if self._connected and self._ws:
            await self._ws.send(json.dumps(message))

    def get_tunnel_manager(self) -> TunnelManager:
        """Obtiene o crea el TunnelManager asociado a esta conexión."""
        if self._tunnel_manager is None:
            self._tunnel_manager = TunnelManager(
                send_ws=self._send_ws_message,
                max_channels=10,
            )
        return self._tunnel_manager

    async def send_discovery(self, network_map: dict[str, Any]) -> bool:
        """Envía un mapa de red descubierto al servidor."""
        message = {
            "type": "discovery",
            "network_map": network_map,
        }

        if self._connected and self._ws:
            try:
                await self._ws.send(json.dumps(message))
                return True
            except Exception as exc:
                log.warning(f"Error enviando discovery: {exc}")
                return False
        return False

    async def disconnect(self) -> None:
        """Cierra la conexión de forma limpia."""
        self._running = False

        # Cerrar túneles activos
        if self._tunnel_manager:
            await self._tunnel_manager.close_all("disconnect")

        # Cancelar tareas en ejecución
        for task_id, task in list(self._tasks.items()):
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
        self._tasks.clear()

        # Cerrar WebSocket
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None

        self._connected = False
        log.info("Desconectado del servidor")

        if self.on_disconnect:
            self.on_disconnect()
