"""SSH backend: connection pooling, process execution, cancellation."""

import asyncio
import time
from dataclasses import dataclass
from typing import Dict, Optional

import asyncssh


class SSHError(Exception):
    """Base exception for SSH operations."""
    pass


class SSHConnectionError(SSHError):
    """Failed to establish SSH connection."""
    pass


class SSHExecutionError(SSHError):
    """Process execution failed."""
    pass


@dataclass
class SSHConnEntry:
    """SSH connection pool entry."""

    conn: asyncssh.SSHClientConnection
    host: str
    last_used: float


class SSHPool:
    """
    SSH connection pool: one connection per unique host, reused across runs.

    Guarantees:
    - pool.get(host) returns a usable asyncssh.SSHClientConnection
    - Subsequent calls to get(same_host) return the same connection (if still open)
    - Stale connections are detected and replaced

    Contingencies:
    - Connection keeps open via keepalive; if remote SSH server dies, detection timing varies
    - No cleanup of stale hosts; pool grows with unique hosts

    Assumptions:
    - Host uniqueness is sufficient (ignores port, username, auth method)
    - All requests to same host use same credentials
    """

    def __init__(
        self,
        known_hosts: Optional[str] = None,
        connect_timeout: float = 10.0,
    ):
        """
        Initialize SSH connection pool.

        Args:
            known_hosts: Path to known_hosts file for host key verification.
                        If None, host keys are not verified (INSECURE).
                        CRITICAL: For v1, this must be a real path.
            connect_timeout: Timeout in seconds for SSH connection attempt.
        """
        self._known_hosts = known_hosts
        self._connect_timeout = connect_timeout
        self._lock = asyncio.Lock()
        self._by_host: Dict[str, SSHConnEntry] = {}

    async def get(
        self,
        host: str,
        *,
        username: Optional[str] = None,
    ) -> asyncssh.SSHClientConnection:
        """
        Get or create SSH connection to host.

        Args:
            host: Hostname or IP address
            username: SSH username (optional; uses default SSH config if not provided)

        Returns:
            asyncssh.SSHClientConnection ready for use

        Raises:
            SSHConnectionError: Connection failed

        Guarantee:
        - Returned connection is open and ready for create_process()
        """
        async with self._lock:
            entry = self._by_host.get(host)

            # Reuse existing connection if still open
            if entry and not entry.conn.is_closing():
                entry.last_used = time.time()
                return entry.conn

            # Connection doesn't exist or is stale; create new one
            try:
                conn = await asyncio.wait_for(
                    asyncssh.connect(
                        host,
                        username=username,
                        known_hosts=self._known_hosts,
                        # CRITICAL: known_hosts must NOT be None for v1
                        # If None, connection is vulnerable to MITM
                        keepalive_interval=30,
                        keepalive_count_max=3,
                    ),
                    timeout=self._connect_timeout,
                )
            except asyncio.TimeoutError as e:
                raise SSHConnectionError(
                    f"SSH connection to {host} timed out after {self._connect_timeout}s"
                ) from e
            except asyncssh.Error as e:
                raise SSHConnectionError(f"SSH connection to {host} failed: {e}") from e

            entry = SSHConnEntry(conn=conn, host=host, last_used=time.time())
            self._by_host[host] = entry
            return conn

    async def close_all(self) -> None:
        """Close all pooled connections."""
        async with self._lock:
            for entry in self._by_host.values():
                entry.conn.close()
            self._by_host.clear()


@dataclass
class RunHandle:
    """Handle to a remote process: wait, cancel, timeout."""

    run_id: int
    host: str
    proc: asyncssh.SSHClientProcess
    started_at: float

    async def wait(self) -> int:
        """
        Wait for process to complete.

        Guarantees:
        - Waits until process has exited
        - Returns exit code

        Returns:
            Process exit code (0 = success, non-zero = failure)
        """
        await self.proc.wait()
        return self.proc.exit_status

    def cancel(self, hard: bool = False) -> None:
        """
        Terminate process.

        Args:
            hard: If False, sends SIGTERM (soft). If True, sends SIGKILL (hard).

        Guarantees:
        - Signal is sent to remote process

        Contingencies:
        - SIGTERM may be ignored by process; SIGKILL enforces termination
        - Process death is asynchronous; may take a few seconds
        """
        if hard:
            self.proc.kill()
        else:
            self.proc.terminate()

    @property
    def elapsed(self) -> float:
        """Seconds since process started."""
        return time.time() - self.started_at


class AsyncSSHBackend:
    """
    Execute processes via SSH channels on pooled connections.

    Guarantees:
    - One connection per unique host (reused across runs)
    - Each run is a distinct SSH channel
    - Timeout is enforced: process is killed after timeout_s

    Contingencies:
    - Remote process parallelism depends on remote OS scheduler and CPU
    - Multiple channels may serialize if remote SSH server implements it that way
    """

    def __init__(self, pool: SSHPool):
        """
        Initialize backend.

        Args:
            pool: SSHPool instance for connection management
        """
        self.pool = pool

    async def run(
        self,
        run_id: int,
        host: str,
        cmd: str,
        *,
        username: Optional[str] = None,
        cwd: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
        timeout_s: int = 60,
    ) -> RunHandle:
        """
        Execute command on remote host via SSH channel.

        The command is executed as: bash -lc 'cd <cwd> && export <env> && <cmd>'

        CRITICAL: cwd and env must have been validated by execution.validate_command_execution()
        before calling this function. If they contain shell metacharacters, injection is possible.

        Args:
            run_id: Unique run identifier
            host: Remote hostname
            cmd: Command to execute
            username: SSH username
            cwd: Working directory (must be validated for shell safety)
            env: Environment variables dict (must be validated for shell safety)
            timeout_s: Timeout in seconds; process is killed if it exceeds this

        Returns:
            RunHandle: Handle to wait, cancel, check elapsed time

        Raises:
            SSHConnectionError: Connection to host failed
            SSHExecutionError: Process execution failed
        """
        # Get SSH connection (pooled)
        try:
            conn = await self.pool.get(host, username=username)
        except SSHConnectionError as e:
            raise SSHExecutionError(f"Cannot connect to {host}: {e}") from e

        # Build command with cwd and env injection
        # NOTE: cwd and env MUST have been validated for metacharacters
        env_prefix = ""
        if env:
            # Build export statements: export VAR1=value1 VAR2=value2 ...
            # Note: values are already validated to not contain shell metacharacters
            env_prefix = "export " + " ".join(
                f"{k}={self._sh_quote(v)}" for k, v in env.items()
            )
            env_prefix += " && "

        cd_prefix = ""
        if cwd:
            # cd to working directory
            # Note: cwd is already validated to not contain shell metacharacters
            cd_prefix = f"cd {self._sh_quote(cwd)} && "

        remote_cmd = f"bash -lc 'set -euo pipefail; {cd_prefix}{env_prefix}{cmd}'"

        # Execute process on remote
        try:
            proc = await conn.create_process(remote_cmd)
        except asyncssh.Error as e:
            raise SSHExecutionError(f"Failed to execute on {host}: {e}") from e

        return RunHandle(
            run_id=run_id,
            host=host,
            proc=proc,
            started_at=time.time(),
        )

    @staticmethod
    def _sh_quote(s: str) -> str:
        """
        Shell-quote a string for use in bash -lc 'command' context.

        Uses single-quote wrapping with internal quote escaping.

        Args:
            s: String to quote

        Returns:
            Safely quoted string that can be interpolated into bash -lc template

        Examples:
            _sh_quote("hello") → "'hello'"
            _sh_quote("it's") → "'it'\"'\"'s'"
        """
        return "'" + s.replace("'", "'\"'\"'") + "'"
