"""
Async SSH session manager.

Merges nterm's full auth stack with VelociTerm's WebSocket streaming model.
Uses asyncio.to_thread for paramiko's blocking I/O.

Changes from v1:
- Stream stats tracking (bytes, messages, duration) for sniffer diagnostics
- WebSocket send wrapped in try/except (fixes crash on client disconnect during send)
- Structured stream events for frontend diagnostics
- on_output callback hook for sniffer integration

Changes from v2:
- Legacy algorithms applied per-Transport via transport_factory (no global mutation)
- Fresh SSHClient on SHA2 RSA retry (fixes stale transport "unknown cipher")
- Algorithm lists validated against paramiko 3.4 registry
"""

import asyncio
import base64
import socket
import time
import logging
from typing import Optional, Callable, Awaitable
from dataclasses import dataclass, field

import paramiko
from fastapi import WebSocket

from .config import SSHConfig, LegacySSHSupport, load_private_key, filter_ansi_sequences

logger = logging.getLogger(__name__)


@dataclass
class StreamStats:
    """Per-session output stream statistics for diagnostics."""
    bytes_received: int = 0
    messages_sent: int = 0
    started_at: float = 0.0
    last_recv_at: float = 0.0
    recv_count: int = 0          # channel.recv calls that returned data
    empty_count: int = 0         # timeout/empty recv calls
    errors: int = 0

    @property
    def duration(self) -> float:
        if not self.started_at:
            return 0.0
        return time.time() - self.started_at

    @property
    def bytes_per_sec(self) -> float:
        d = self.duration
        return self.bytes_received / d if d > 0 else 0.0

    def to_dict(self) -> dict:
        return {
            "bytes_received": self.bytes_received,
            "messages_sent": self.messages_sent,
            "recv_count": self.recv_count,
            "empty_count": self.empty_count,
            "errors": self.errors,
            "duration_sec": round(self.duration, 1),
            "bytes_per_sec": round(self.bytes_per_sec, 1),
        }


class SSHSession:
    """
    Single SSH session with full auth support.

    Wraps paramiko client + channel with nterm's connection logic.
    All connect/disconnect methods are synchronous — the manager
    calls them via asyncio.to_thread.
    """

    def __init__(self, config: SSHConfig):
        self.config = config
        self.client: Optional[paramiko.SSHClient] = None
        self.channel: Optional[paramiko.Channel] = None
        self.connected: bool = False
        self.stats: StreamStats = StreamStats()

    def _new_client(self) -> paramiko.SSHClient:
        """Create a fresh SSHClient instance."""
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        return client

    def _build_connect_params(self, disable_sha2: bool = True) -> dict:
        """
        Build paramiko connect kwargs from config.

        Args:
            disable_sha2: If True, disable SHA2 RSA algorithms for
                          compatibility with older devices. Set False
                          on retry for modern device fallback.
        """
        params = {
            "hostname": self.config.hostname,
            "port": self.config.port,
            "username": self.config.username,
            "timeout": self.config.timeout,
            "allow_agent": False,
            "look_for_keys": False,
        }

        if disable_sha2:
            params["disabled_algorithms"] = {
                "pubkeys": ["rsa-sha2-512", "rsa-sha2-256"]
            }

        # Legacy mode: per-transport factory (instance-level, not global)
        if self.config.legacy_mode:
            params["transport_factory"] = LegacySSHSupport.make_transport

        # Auth: key takes priority, with optional password fallback
        if self.config.key_content or self.config.key_file:
            params["pkey"] = load_private_key(self.config)
            if self.config.password:
                params["password"] = self.config.password
        elif self.config.password:
            params["password"] = self.config.password

        return params

    def connect(self) -> None:
        """
        Establish SSH connection and open interactive shell.

        Connection strategy:
            1. Build connect params (with legacy transport_factory if needed)
            2. Attempt connection with SHA2 RSA disabled (compatibility)
            3. On failure, fresh client + retry with SHA2 RSA enabled (modern)
            4. Set transport keepalive
            5. Open interactive shell with xterm PTY
        """
        self.client = self._new_client()
        connect_params = self._build_connect_params(disable_sha2=True)

        try:
            self.client.connect(**connect_params)
        except Exception as first_error:
            logger.debug(
                f"First connect attempt failed for {self.config.hostname}: "
                f"{first_error}. Retrying with SHA2 RSA algorithms enabled."
            )
            # Fresh client — stale transport can't renegotiate
            self.client = self._new_client()
            connect_params = self._build_connect_params(disable_sha2=False)
            self.client.connect(**connect_params)

        # Transport keepalive
        transport = self.client.get_transport()
        if transport:
            transport.set_keepalive(self.config.keepalive_interval)

        # Open interactive shell
        self.channel = self.client.invoke_shell(term="xterm", width=200, height=24)
        self.channel.settimeout(0.5)
        self.connected = True

        legacy_tag = " [legacy]" if self.config.legacy_mode else ""
        logger.info(
            f"Connected to {self.config.hostname}:{self.config.port} "
            f"via {self.config.auth_method}{legacy_tag}"
        )

    def disconnect(self) -> None:
        """Close channel and client cleanly."""
        self.connected = False

        if self.channel:
            try:
                self.channel.close()
            except Exception as e:
                logger.debug(f"Channel close error: {e}")
            self.channel = None

        if self.client:
            try:
                self.client.close()
            except Exception as e:
                logger.debug(f"Client close error: {e}")
            self.client = None

        logger.debug(f"Disconnected from {self.config.hostname}")

    def send(self, data: str) -> None:
        """Send input to the SSH channel."""
        if self.channel and not self.channel.closed:
            self.channel.send(data)

    def resize(self, cols: int, rows: int) -> None:
        """Resize the PTY."""
        if self.channel and not self.channel.closed:
            self.channel.resize_pty(width=cols, height=rows)


class SSHSessionManager:
    """
    Manages multiple concurrent SSH sessions keyed by tab_id.

    WebSocket message protocol:
        Client → Server:
            { type: "connect", hostname, port, username, password?, key_file?,
              key_content?, key_passphrase?, legacy_mode? }
            { type: "input", data: "<keystroke(s)>" }
            { type: "resize", cols: N, rows: N }

        Server → Client:
            { type: "ssh_output", data: "<base64>", tabId: "<uuid>" }
            { type: "connected", tabId, hostname, auth_method }
            { type: "disconnected", tabId, stats }
            { type: "error", message, tabId }
    """

    def __init__(self):
        self.sessions: dict[str, SSHSession] = {}

    async def connect(
        self, tab_id: str, config: SSHConfig, websocket: WebSocket
    ) -> bool:
        """
        Create and connect an SSH session.
        Returns True on success, False on failure.
        """
        config.validate()

        session = SSHSession(config)
        self.sessions[tab_id] = session

        try:
            await asyncio.to_thread(session.connect)
            await self._ws_send(websocket, {
                "type": "connected",
                "tabId": tab_id,
                "hostname": config.hostname,
                "auth_method": config.auth_method,
            })
            return True

        except Exception as e:
            error_msg = str(e)
            logger.error(f"SSH connection failed for {tab_id}: {error_msg}")

            error_display = f"\r\n\x1b[31mConnection failed: {error_msg}\x1b[0m\r\n"
            encoded = base64.b64encode(error_display.encode()).decode()
            await self._ws_send(websocket, {
                "type": "ssh_output", "data": encoded, "tabId": tab_id,
            })
            await self._ws_send(websocket, {
                "type": "error", "message": error_msg, "tabId": tab_id,
            })

            self.sessions.pop(tab_id, None)
            return False

    async def stream_output(self, tab_id: str, websocket: WebSocket) -> None:
        """
        Stream SSH output to the WebSocket.

        SNIFFER TAP POINT:
            The filtered text variable is where Parsing Fire will
            receive its input. Currently only computed for diagnostics.
        """
        session = self.sessions.get(tab_id)
        if not session or not session.channel:
            return

        session.stats.started_at = time.time()

        try:
            while session.connected and not session.channel.closed:
                try:
                    data = await asyncio.to_thread(session.channel.recv, 4096)
                    if not data:
                        break

                    now = time.time()
                    session.stats.bytes_received += len(data)
                    session.stats.recv_count += 1
                    session.stats.last_recv_at = now

                    # Raw bytes to terminal — xterm.js renders ANSI natively
                    encoded = base64.b64encode(data).decode("utf-8")
                    await self._ws_send(websocket, {
                        "type": "ssh_output", "data": encoded, "tabId": tab_id,
                    })
                    session.stats.messages_sent += 1

                    # ── SNIFFER TAP POINT ───────────────────────
                    # ANSI-filtered text: what Parsing Fire will see
                    filtered = filter_ansi_sequences(
                        data.decode("utf-8", errors="replace")
                    )
                    # logger.debug(f"[sniffer:{tab_id[:8]}] {filtered!r}")
                    # ────────────────────────────────────────────

                except socket.timeout:
                    session.stats.empty_count += 1
                    continue
                except OSError:
                    break

        except asyncio.CancelledError:
            logger.debug(f"Stream cancelled for {tab_id}")

        finally:
            stats = session.stats.to_dict()
            logger.info(f"[{tab_id[:8]}] Stream ended: {stats}")
            await self._ws_send(websocket, {
                "type": "disconnected", "tabId": tab_id, "stats": stats,
            })

    async def send_input(self, tab_id: str, data: str) -> None:
        """Send keystroke input to an SSH session."""
        session = self.sessions.get(tab_id)
        if session:
            session.send(data)

    async def resize(self, tab_id: str, cols: int, rows: int) -> None:
        """Resize a session's PTY dimensions."""
        session = self.sessions.get(tab_id)
        if session:
            session.resize(cols, rows)

    async def disconnect(self, tab_id: str) -> None:
        """Disconnect and remove a session."""
        session = self.sessions.pop(tab_id, None)
        if session:
            await asyncio.to_thread(session.disconnect)

    async def disconnect_all(self) -> None:
        """Disconnect all sessions — clean shutdown."""
        tab_ids = list(self.sessions.keys())
        for tab_id in tab_ids:
            await self.disconnect(tab_id)

    def get_session_info(self) -> list[dict]:
        """Return info about all active sessions (for /api/diag endpoint)."""
        info = []
        for tab_id, session in self.sessions.items():
            info.append({
                "tab_id": tab_id[:8],
                "hostname": session.config.hostname,
                "port": session.config.port,
                "auth_method": session.config.auth_method,
                "connected": session.connected,
                "stats": session.stats.to_dict(),
            })
        return info

    @staticmethod
    async def _ws_send(websocket: WebSocket, data: dict) -> None:
        """
        Send JSON to WebSocket, swallowing errors if client already gone.

        Prevents server crash when browser tab closes while output streams.
        """
        try:
            await websocket.send_json(data)
        except Exception:
            pass