# Copyright 2026 Celesto AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""SSH command execution for SmolVM guests.

Uses `paramiko <https://www.paramiko.org/>`_ for persistent SSH connections
to guest VMs.  A single TCP connection is established on first use (or during
:meth:`SSHClient.wait_for_ssh`) and reused for all subsequent commands,
eliminating the ~170ms overhead of forking a new ``ssh`` process per call.
"""

import logging
import shlex
import socket
import time
from typing import Literal

import paramiko

from smolvm.exceptions import OperationTimeoutError, SmolVMError
from smolvm.types import CommandResult

logger = logging.getLogger(__name__)

ShellMode = Literal["login", "raw"]


class SSHClient:
    """Execute commands on a microVM guest via SSH.

    Maintains a persistent paramiko SSH connection that is established
    lazily on first use and reused for all subsequent commands.

    Args:
        host: Guest IP address or hostname.
        user: SSH user (default ``root``).
        port: SSH port on the guest (default ``22``).
        key_path: Optional path to an SSH private key file.
        password: Optional password for authentication.
        connect_timeout: Seconds to wait for the TCP connection.
    """

    def __init__(
        self,
        host: str,
        user: str = "root",
        port: int = 22,
        key_path: str | None = None,
        password: str | None = None,
        connect_timeout: int = 10,
    ) -> None:
        if not host:
            raise ValueError("host cannot be empty")
        if not user:
            raise ValueError("user cannot be empty")
        if port < 1 or port > 65535:
            raise ValueError(f"port must be 1-65535, got {port}")
        if connect_timeout < 1:
            raise ValueError("connect_timeout must be >= 1")

        self.host = host
        self.user = user
        self.port = port
        self.key_path = key_path
        self.password = password
        self.connect_timeout = connect_timeout

        self._client: paramiko.SSHClient | None = None

    # ── Connection lifecycle ────────────────────────────────────

    def _connect(self) -> paramiko.SSHClient:
        """Establish a paramiko SSH connection.

        Returns:
            Connected paramiko SSHClient.

        Raises:
            SmolVMError: If connection fails.
        """
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        connect_kwargs: dict[str, object] = {
            "hostname": self.host,
            "port": self.port,
            "username": self.user,
            "timeout": float(self.connect_timeout),
            "allow_agent": False,
            "look_for_keys": False,
            "banner_timeout": float(self.connect_timeout),
        }

        if self.key_path:
            connect_kwargs["key_filename"] = self.key_path
        elif self.password:
            connect_kwargs["password"] = self.password
        else:
            # Fall back to trying the agent / default keys
            connect_kwargs["allow_agent"] = True
            connect_kwargs["look_for_keys"] = True

        try:
            client.connect(**connect_kwargs)  # type: ignore[arg-type]
        except Exception as e:
            client.close()
            raise SmolVMError(f"SSH connection failed: {e}") from e

        return client

    def _ensure_connected(self) -> paramiko.SSHClient:
        """Return existing connection, or create a new one.

        Automatically reconnects if the previous connection is dead.
        """
        if self._client is not None:
            transport = self._client.get_transport()
            if transport is not None and transport.is_active():
                return self._client
            # Transport is dead — close and reconnect.
            logger.debug("SSH transport is dead, reconnecting to %s:%d", self.host, self.port)
            self._client.close()
            self._client = None

        self._client = self._connect()
        return self._client

    @property
    def connected(self) -> bool:
        """Check if the SSH connection is alive."""
        if self._client is None:
            return False
        transport = self._client.get_transport()
        return transport is not None and transport.is_active()

    def close(self) -> None:
        """Close the SSH connection and release resources."""
        if self._client is not None:
            try:
                self._client.close()
            except Exception:
                pass
            self._client = None

    # ── Command execution ───────────────────────────────────────

    @staticmethod
    def _wrap_login_shell_command(command: str) -> str:
        """Wrap a command so it runs inside a login shell."""
        quoted_command = shlex.quote(command)
        return f'SHELL_BIN="${{SHELL:-/bin/sh}}"; exec "$SHELL_BIN" -lc {quoted_command}'

    def _prepare_remote_command(self, command: str, shell: ShellMode) -> str:
        """Prepare a command string based on desired shell mode."""
        if shell == "raw":
            return command
        if shell == "login":
            return self._wrap_login_shell_command(command)
        raise ValueError("shell must be 'login' or 'raw'")

    def run(
        self,
        command: str,
        timeout: int = 30,
        shell: ShellMode = "login",
    ) -> CommandResult:
        """Execute a command on the guest VM.

        Uses the persistent SSH connection.  If the connection is not
        yet established or has been lost, it is (re)created transparently.

        Args:
            command: Shell command to execute.
            timeout: Maximum seconds to wait for the command.
            shell: Command execution mode:
                - ``"login"`` (default): run via guest login shell.
                - ``"raw"``: execute command directly with no shell wrapping.

        Returns:
            :class:`~smolvm.types.CommandResult` with exit code, stdout,
            and stderr.

        Raises:
            ValueError: If *command* is empty.
            OperationTimeoutError: If the command exceeds *timeout*.
            SmolVMError: If the SSH connection cannot be established.
        """
        if not command or not command.strip():
            raise ValueError("command cannot be empty")
        if timeout < 1:
            raise ValueError("timeout must be >= 1")

        remote_command = self._prepare_remote_command(command, shell=shell)

        logger.debug(
            "SSH exec on %s:%d (shell=%s): %s",
            self.host,
            self.port,
            shell,
            command,
        )

        client = self._ensure_connected()

        try:
            _, stdout_ch, stderr_ch = client.exec_command(remote_command, timeout=timeout)
            # Read all output
            stdout = stdout_ch.read().decode("utf-8", errors="replace")
            stderr = stderr_ch.read().decode("utf-8", errors="replace")
            exit_code = stdout_ch.channel.recv_exit_status()

            return CommandResult(
                exit_code=exit_code,
                stdout=stdout,
                stderr=stderr,
            )
        except socket.timeout as e:
            raise OperationTimeoutError(
                f"ssh {self.host}:{self.port}: {command}",
                timeout,
            ) from e
        except paramiko.SSHException as e:
            # Connection may have died mid-command — invalidate it
            self.close()
            raise SmolVMError(f"SSH command failed: {e}") from e
        except Exception as e:
            raise SmolVMError(f"SSH command failed: {e}") from e

    # ── Readiness detection ─────────────────────────────────────

    def _tcp_port_open(self, timeout: float = 0.1) -> bool:
        """Check if the SSH port is accepting TCP connections.

        This is much faster than a full SSH handshake (~1ms vs ~100ms)
        and is used to efficiently poll for port readiness before attempting
        a paramiko connect.
        """
        try:
            with socket.create_connection((self.host, self.port), timeout=timeout) as sock:
                # Read the SSH banner to confirm sshd is actually ready,
                # not just TCP backlog accepting connections.
                sock.settimeout(timeout)
                data = sock.recv(32)
                return data.startswith(b"SSH-")
        except (OSError, socket.timeout):
            return False

    def wait_for_ssh(self, timeout: float = 60.0, interval: float = 0.1) -> None:
        """Wait for the SSH daemon to become reachable on the guest.

        Uses a two-phase approach for fast detection:

        1. **TCP probe** — lightweight ``socket.connect()`` calls (~1ms each)
           with a short fixed polling interval to detect when sshd is
           listening and sending its banner.
        2. **Paramiko connect** — full SSH handshake + auth.  The resulting
           connection is kept open for subsequent :meth:`run` calls.

        Args:
            timeout: Maximum seconds to wait.
            interval: Seconds between TCP probe attempts.

        Raises:
            OperationTimeoutError: If SSH does not become available
                within *timeout* seconds.
        """
        if timeout <= 0:
            raise ValueError("timeout must be > 0")

        deadline = time.monotonic() + timeout

        logger.info(
            "Waiting for SSH on %s:%d (timeout=%.0fs)",
            self.host,
            self.port,
            timeout,
        )

        # Phase 1: Fast TCP probe — detect when sshd port is open.
        while time.monotonic() < deadline:
            if self._tcp_port_open(timeout=min(0.5, deadline - time.monotonic())):
                logger.debug("TCP port %d is open on %s", self.port, self.host)
                break

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise OperationTimeoutError(
                    f"wait_for_ssh({self.host}:{self.port}): port never opened",
                    timeout,
                )
            time.sleep(min(interval, remaining))

        # Phase 2: Establish persistent paramiko connection.
        last_error: str = ""
        while time.monotonic() < deadline:
            try:
                self._client = self._connect()
                logger.info("SSH is ready on %s:%d", self.host, self.port)
                return
            except SmolVMError as e:
                last_error = str(e)
                logger.debug(
                    "SSH connect failed on %s:%d: %s",
                    self.host,
                    self.port,
                    last_error,
                )

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break
            time.sleep(min(0.2, remaining))

        raise OperationTimeoutError(
            f"wait_for_ssh({self.host}:{self.port}): last error: {last_error}",
            timeout,
        )
