"""SSH connection manager with per-node caching and transparent reconnect."""

from __future__ import annotations

import logging
import time

import asyncssh

from vmcluster_mcp.models import ExecResult

logger = logging.getLogger("vmcluster_mcp.exec.ssh")


class SSHConnectionManager:
    """Async SSH connection manager with per-node caching and transparent reconnect.

    Connections are keyed by (ip, port, user, key_path). On a stale connection
    (e.g. VM rebooted), the next execute() call detects the failure, invalidates
    the cache entry, and retries once with a fresh connection — transparently.
    """

    def __init__(self) -> None:
        self._connections: dict[tuple[str, int, str, str], asyncssh.SSHClientConnection] = {}

    async def execute(
        self,
        ip: str,
        port: int,
        user: str,
        key_path: str,
        command: str,
        timeout_s: int = 30,
    ) -> ExecResult:
        """Run a shell command on a remote host, using a cached connection when available.

        On stale connection (ConnectionLost, DisconnectError, OSError), invalidates
        the cache entry and retries once with a fresh connection.

        Args:
            ip: Host IP address.
            port: SSH port.
            user: SSH username.
            key_path: Path to the SSH private key file.
            command: Shell command to run.
            timeout_s: Command execution timeout in seconds.

        Returns:
            ExecResult with exit_code, stdout, stderr, duration_s, timed_out.
        """
        key = (ip, port, user, key_path)
        start = time.monotonic()

        # Try with cached or new connection
        for attempt in range(2):
            conn = self._connections.get(key)
            try:
                if conn is None:
                    conn = await asyncssh.connect(
                        ip,
                        port=port,
                        username=user,
                        client_keys=[key_path],
                        known_hosts=None,
                        login_timeout=10,
                    )
                    self._connections[key] = conn

                result = await conn.run(command, timeout=timeout_s)
                elapsed = time.monotonic() - start
                return ExecResult(
                    exit_code=result.exit_status if result.exit_status is not None else -1,
                    stdout=result.stdout or "",
                    stderr=result.stderr or "",
                    duration_s=elapsed,
                    timed_out=False,
                )

            except asyncssh.TimeoutError:
                elapsed = time.monotonic() - start
                return ExecResult(
                    exit_code=-1,
                    stdout="",
                    stderr="Command timed out",
                    duration_s=elapsed,
                    timed_out=True,
                )

            except (asyncssh.ConnectionLost, asyncssh.DisconnectError, OSError) as exc:
                # Stale connection — invalidate and retry once
                logger.warning("Stale SSH connection to %s:%d — %s", ip, port, exc)
                if key in self._connections:
                    try:
                        self._connections[key].close()
                    except Exception:
                        pass
                    del self._connections[key]
                conn = None
                if attempt == 1:
                    # Second attempt also failed
                    elapsed = time.monotonic() - start
                    return ExecResult(
                        exit_code=-1,
                        stdout="",
                        stderr=str(exc),
                        duration_s=elapsed,
                        timed_out=False,
                    )
                # Continue to retry

            except Exception as exc:
                elapsed = time.monotonic() - start
                return ExecResult(
                    exit_code=-1,
                    stdout="",
                    stderr=str(exc),
                    duration_s=elapsed,
                    timed_out=False,
                )

        # Should not be reached
        elapsed = time.monotonic() - start
        return ExecResult(
            exit_code=-1,
            stdout="",
            stderr="Unknown error",
            duration_s=elapsed,
            timed_out=False,
        )

    async def invalidate(self, ip: str) -> None:
        """Remove all cached connections for the given IP address.

        Args:
            ip: Host IP — all (port, user, key) variants for this IP are removed.
        """
        keys_to_remove = [k for k in self._connections if k[0] == ip]
        for k in keys_to_remove:
            try:
                self._connections[k].close()
            except Exception:
                pass
            del self._connections[k]

    async def close_all(self) -> None:
        """Close all cached connections and clear the cache."""
        for conn in self._connections.values():
            try:
                conn.close()
            except Exception:
                pass
        self._connections.clear()
