"""SSH local-forward tunnel context manager for WinRM connectivity."""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING, Any, Self

import asyncssh

if TYPE_CHECKING:
    from types import TracebackType

    from mcp_vm_blackbox.targets import TargetConfig


class SshTunnel:
    """Async context manager that opens an SSH local-forward tunnel via asyncssh.

    For local targets (``ssh_host=""``), this is a no-op — no connection is
    opened and no port-forwarding takes place.

    Example:
        >>> async with SshTunnel(cfg):
        ...     client = WinRMClient(host="127.0.0.1", port=cfg.tunnel_local_port)
        ...     client.run_ps("Get-Date")
    """

    #: Seconds to wait after establishing the tunnel before the caller
    #: attempts a connection through it.
    _CONNECT_DELAY: float = 0.5

    def __init__(self, cfg: TargetConfig) -> None:
        """Initialise tunnel parameters without opening the connection.

        Args:
            cfg: Resolved target configuration. When ``is_remote`` is False,
                the tunnel is a no-op.
        """
        self._cfg = cfg

    # ------------------------------------------------------------------
    # Async context manager protocol
    # ------------------------------------------------------------------

    async def __aenter__(self) -> Self:
        """Open the SSH tunnel (no-op for local targets).

        Returns:
            This instance, allowing ``async with SshTunnel(cfg) as tunnel:``.
        """
        if not self._cfg.is_remote:
            return self

        if "@" in self._cfg.ssh_host:
            username, host = self._cfg.ssh_host.rsplit("@", 1)
        else:
            host = self._cfg.ssh_host
            username = self._cfg.ssh_user or ""

        conn_kwargs: dict[str, Any] = {"host": host, "port": self._cfg.ssh_port or 22, "known_hosts": None}
        if username:
            conn_kwargs["username"] = username
        if self._cfg.ssh_identity_file:
            conn_kwargs["client_keys"] = [self._cfg.ssh_identity_file]

        self._conn = await asyncssh.connect(**conn_kwargs)
        self._listener = await self._conn.forward_local_port(
            "", self._cfg.tunnel_local_port, self._cfg.winrm_host, self._cfg.winrm_port
        )
        await asyncio.sleep(self._CONNECT_DELAY)
        return self

    async def __aexit__(
        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
    ) -> None:
        """Close the SSH tunnel connection and listener."""
        if not self._cfg.is_remote:
            return
        listener = getattr(self, "_listener", None)
        conn = getattr(self, "_conn", None)
        try:
            if listener is not None:
                listener.close()
        finally:
            if conn is not None:
                conn.close()
                await conn.wait_closed()
