"""SSH blocks — session management, exec, and probe."""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass

from loguru import logger

from cutip_blocks.decorator import block


@dataclass
class ExecResult:
    """Result of a remote command execution."""

    exit_code: int
    stdout: str
    stderr: str = ""


class SSHSession:
    """Persistent SSH session through a container to a remote host.

    Holds a paramiko channel open for the lifetime of the ``with`` block.
    All commands execute over the same connection — no per-call handshake.

    Credentials passed at init are stored only for connection setup and
    are redacted from all log output.
    """

    def __init__(
        self,
        host: str,
        username: str,
        password: str,
        *,
        container_runtime=None,
        container_name: str | None = None,
        port: int = 22,
    ) -> None:
        self._host = host
        self._username = username
        self._port = port
        self._container_runtime = container_runtime
        self._container_name = container_name
        self._redact: set[str] = {password}
        self._client = None

    def connect(self) -> None:
        """Open the SSH connection."""
        import paramiko

        self._client = paramiko.SSHClient()
        self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        if self._container_runtime and self._container_name:
            # Execute SSH through container — the container has network access
            self._via_container = True
        else:
            self._via_container = False

        # Direct paramiko connection (from host or from container context)
        self._client.connect(
            hostname=self._host,
            port=self._port,
            username=self._username,
            password=next(iter(self._redact)),  # the original password
            timeout=30,
            allow_agent=False,
            look_for_keys=False,
        )

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

    def exec(self, cmd: str, *, block_name: str = "") -> ExecResult:
        """Execute a command over the SSH session.

        Args:
            cmd: Shell command string (pipes, redirects, etc. are valid).
            block_name: Optional block name for log prefix.
        """
        if self._client is None:
            raise RuntimeError("SSH session is not connected")

        prefix = f"[{block_name}] " if block_name else ""
        display_cmd = self._redact_str(cmd)
        logger.info("{}ssh {}@{} :: {}", prefix, self._username, self._host, display_cmd)

        _, stdout_ch, stderr_ch = self._client.exec_command(cmd, timeout=300)
        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()

        if exit_code != 0:
            logger.warning(
                "{}ssh {}@{} :: {} [exit {}]",
                prefix,
                self._username,
                self._host,
                display_cmd,
                exit_code,
            )

        return ExecResult(exit_code=exit_code, stdout=stdout, stderr=stderr)

    def _redact_str(self, text: str) -> str:
        """Replace known secrets in text with ****."""
        result = text
        for secret in self._redact:
            if secret and secret in result:
                result = result.replace(secret, "****")
        return result


@contextmanager
def session(
    ctx,
    *,
    container: str,
    host: str,
    username: str,
    password: str,
    port: int = 22,
):
    """Context manager that opens a persistent SSH session through a container.

    Usage::

        with ssh.session(ctx, container="my-container",
                         host="10.0.0.1", username="root",
                         password=ctx.config["password"]) as sesh:
            sesh.exec("whoami")
            k8s.get_pod(ctx, sesh, namespace="default", deployment="web")

    The session opens one SSH connection and reuses it for all commands
    within the ``with`` block.
    """
    sesh = SSHSession(
        host=host,
        username=username,
        password=password,
        port=port,
        container_runtime=getattr(ctx, "runtime", None),
        container_name=container,
    )
    logger.info("Opening SSH session to {}@{} via container '{}'", username, host, container)
    sesh.connect()
    try:
        yield sesh
    finally:
        sesh.close()
        logger.info("Closed SSH session to {}@{}", username, host)


@block(name="SSH Exec", category="ssh", action="exec")
def exec(ctx, sesh: SSHSession, *, cmd: str) -> ExecResult:
    """Execute a command over an existing SSH session."""
    return sesh.exec(cmd, block_name="SSH Exec")


@block(name="SSH Probe", category="ssh", action="probe")
def probe(ctx, sesh: SSHSession) -> ExecResult:
    """Verify SSH session is authenticated and responsive."""
    result = sesh.exec("whoami", block_name="SSH Probe")
    if result.exit_code != 0:
        raise RuntimeError(f"SSH probe failed: {result.stderr.strip()}")
    logger.info("[SSH Probe] Authenticated as: {}", result.stdout.strip())
    return result
