"""Remote test runner — executes commands on remote devices via SSH.

Uses asyncssh for async SSH connections. Handles PowerShell (Windows)
and bash (Linux/macOS) shell differences, background process management,
and file transfer.
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import asyncssh

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


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

    command: str
    exit_code: int
    stdout: str
    stderr: str
    duration_s: float
    timed_out: bool = False

    @property
    def success(self) -> bool:
        return self.exit_code == 0 and not self.timed_out


@dataclass
class BackgroundProcess:
    """Handle for a background process running on the remote host."""

    command: str
    pid: int | None = None
    _process: asyncssh.SSHClientProcess | None = field(default=None, repr=False)
    started_at: float = field(default_factory=time.monotonic)


@dataclass
class DeviceInfo:
    """Minimal device connection info needed by the runner."""

    hostname: str
    username: str
    shell: str  # "powershell" | "bash"
    os: str  # "windows" | "linux" | "macos"
    ssh_key: str | None = None
    port: int = 22


class SSHConnectionError(Exception):
    """Raised when SSH connection fails."""


class CommandTimeoutError(Exception):
    """Raised when a command exceeds its timeout."""


# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------


class RemoteTestRunner:
    """Async SSH test runner for remote device testing.

    Usage::

        device = DeviceInfo(hostname="alienware.local", username="user",
                            shell="powershell", os="windows")
        async with RemoteTestRunner(device) as runner:
            result = await runner.run_command("llmhosts --version")
            print(result.stdout)
    """

    def __init__(self, device: DeviceInfo, connect_timeout: float = 30.0) -> None:
        self.device = device
        self.connect_timeout = connect_timeout
        self._conn: asyncssh.SSHClientConnection | None = None
        self._background_procs: list[BackgroundProcess] = []

    # -- Context manager -----------------------------------------------------

    async def __aenter__(self) -> RemoteTestRunner:
        await self.connect()
        return self

    async def __aexit__(self, exc_type: type | None, exc_val: Exception | None, exc_tb: object) -> None:
        await self.close()

    # -- Connection ----------------------------------------------------------

    async def connect(self) -> None:
        """Establish SSH connection to the remote device."""
        try:
            import asyncssh
        except ImportError as exc:
            raise ImportError(
                "asyncssh is required for remote testing. Install with: pip install llmhosts[dev]"
            ) from exc

        connect_kwargs: dict[str, object] = {
            "host": self.device.hostname,
            "port": self.device.port,
            "username": self.device.username,
            "known_hosts": None,  # Accept any host key for test infra
        }

        if self.device.ssh_key:
            key_path = Path(self.device.ssh_key).expanduser()
            if key_path.exists():
                connect_kwargs["client_keys"] = [str(key_path)]
            else:
                logger.warning("SSH key %s not found, falling back to agent auth", key_path)

        try:
            self._conn = await asyncio.wait_for(
                asyncssh.connect(**connect_kwargs),
                timeout=self.connect_timeout,
            )
            logger.info("Connected to %s@%s", self.device.username, self.device.hostname)
        except asyncio.TimeoutError as exc:
            raise SSHConnectionError(
                f"SSH connection to {self.device.hostname} timed out after {self.connect_timeout}s"
            ) from exc
        except OSError as exc:
            raise SSHConnectionError(f"SSH connection to {self.device.hostname} failed: {exc}") from exc

    async def close(self) -> None:
        """Stop all background processes and close the SSH connection."""
        for proc in self._background_procs:
            try:
                await self.stop_background(proc)
            except Exception:
                logger.warning("Failed to stop background process: %s", proc.command)
        self._background_procs.clear()

        if self._conn is not None:
            self._conn.close()
            self._conn = None
            logger.info("Disconnected from %s", self.device.hostname)

    def _ensure_connected(self) -> asyncssh.SSHClientConnection:
        if self._conn is None:
            raise SSHConnectionError("Not connected. Call connect() first or use as context manager.")
        return self._conn

    # -- Shell helpers -------------------------------------------------------

    def _wrap_command(self, cmd: str) -> str:
        """Wrap a command for the device's shell type.

        PowerShell needs explicit invocation; bash commands pass through.
        """
        if self.device.shell == "powershell":
            # Escape double quotes inside the command for PowerShell
            escaped = cmd.replace('"', '`"')
            return f'powershell.exe -NoProfile -NonInteractive -Command "{escaped}"'
        return cmd

    def _wrap_background_command(self, cmd: str) -> str:
        """Wrap a command to run in the background and return its PID."""
        if self.device.shell == "powershell":
            escaped = cmd.replace('"', '`"')
            return (
                f"powershell.exe -NoProfile -NonInteractive -Command "
                f'"$p = Start-Process -PassThru -NoNewWindow -FilePath powershell '
                f"-ArgumentList '-NoProfile','-Command','{escaped}'; $p.Id\""
            )
        # bash: nohup + background, echo PID
        return f"nohup {cmd} > /dev/null 2>&1 & echo $!"

    def _kill_command(self, pid: int) -> str:
        """Build a command to kill a process by PID."""
        if self.device.shell == "powershell":
            return f'powershell.exe -NoProfile -NonInteractive -Command "Stop-Process -Id {pid} -Force -ErrorAction SilentlyContinue"'
        return f"kill {pid} 2>/dev/null; sleep 1; kill -9 {pid} 2>/dev/null || true"

    # -- Command execution ---------------------------------------------------

    async def run_command(self, cmd: str, timeout: float = 60.0) -> CommandResult:
        """Execute a command on the remote device and wait for completion.

        Args:
            cmd: The command to execute.
            timeout: Maximum seconds to wait. 0 means no timeout.

        Returns:
            CommandResult with exit code, stdout, stderr, and timing.
        """
        conn = self._ensure_connected()
        wrapped = self._wrap_command(cmd)
        start = time.monotonic()

        try:
            if timeout > 0:
                result = await asyncio.wait_for(
                    conn.run(wrapped),
                    timeout=timeout,
                )
            else:
                result = await conn.run(wrapped)

            duration = time.monotonic() - start
            return CommandResult(
                command=cmd,
                exit_code=result.exit_status or 0,
                stdout=(result.stdout or "").strip(),
                stderr=(result.stderr or "").strip(),
                duration_s=round(duration, 3),
            )

        except asyncio.TimeoutError:
            duration = time.monotonic() - start
            logger.warning("Command timed out after %.1fs: %s", timeout, cmd)
            return CommandResult(
                command=cmd,
                exit_code=-1,
                stdout="",
                stderr=f"Command timed out after {timeout}s",
                duration_s=round(duration, 3),
                timed_out=True,
            )

    async def run_background(self, cmd: str) -> BackgroundProcess:
        """Start a command in the background on the remote device.

        Returns a BackgroundProcess handle that can be passed to stop_background().
        The command is started via nohup (bash) or Start-Process (PowerShell).
        """
        conn = self._ensure_connected()
        wrapped = self._wrap_background_command(cmd)

        result = await asyncio.wait_for(conn.run(wrapped), timeout=30.0)
        stdout = (result.stdout or "").strip()

        pid: int | None = None
        try:
            pid = int(stdout.splitlines()[-1])
        except (ValueError, IndexError):
            logger.warning("Could not parse PID from background command output: %r", stdout)

        proc = BackgroundProcess(command=cmd, pid=pid)
        self._background_procs.append(proc)
        logger.info("Started background process: %s (PID=%s)", cmd, pid)
        return proc

    async def stop_background(self, proc: BackgroundProcess) -> CommandResult:
        """Stop a background process on the remote device."""
        if proc.pid is None:
            return CommandResult(
                command=f"stop({proc.command})",
                exit_code=-1,
                stdout="",
                stderr="No PID available to stop",
                duration_s=0.0,
            )

        kill_cmd = self._kill_command(proc.pid)
        conn = self._ensure_connected()
        start = time.monotonic()

        try:
            result = await asyncio.wait_for(conn.run(kill_cmd), timeout=15.0)
            duration = time.monotonic() - start
            if proc in self._background_procs:
                self._background_procs.remove(proc)
            return CommandResult(
                command=f"stop(PID={proc.pid})",
                exit_code=result.exit_status or 0,
                stdout=(result.stdout or "").strip(),
                stderr=(result.stderr or "").strip(),
                duration_s=round(duration, 3),
            )
        except asyncio.TimeoutError:
            duration = time.monotonic() - start
            return CommandResult(
                command=f"stop(PID={proc.pid})",
                exit_code=-1,
                stdout="",
                stderr="Kill command timed out",
                duration_s=round(duration, 3),
                timed_out=True,
            )

    # -- File transfer -------------------------------------------------------

    async def upload_file(self, local_path: str | Path, remote_path: str) -> None:
        """Upload a local file to the remote device via SFTP."""
        try:
            import asyncssh
        except ImportError as exc:
            raise ImportError("asyncssh is required for file transfer") from exc

        conn = self._ensure_connected()
        async with conn.start_sftp_client() as sftp:
            await asyncssh.scp(str(local_path), (sftp, remote_path))
        logger.info("Uploaded %s -> %s:%s", local_path, self.device.hostname, remote_path)

    async def download_file(self, remote_path: str, local_path: str | Path) -> None:
        """Download a file from the remote device via SFTP."""
        try:
            import asyncssh
        except ImportError as exc:
            raise ImportError("asyncssh is required for file transfer") from exc

        conn = self._ensure_connected()
        async with conn.start_sftp_client() as sftp:
            await asyncssh.scp((sftp, remote_path), str(local_path))
        logger.info("Downloaded %s:%s -> %s", self.device.hostname, remote_path, local_path)

    # -- Convenience ---------------------------------------------------------

    async def check_shell(self) -> str:
        """Detect the remote shell type by probing.

        Returns "powershell" or "bash" (or "unknown").
        """
        conn = self._ensure_connected()

        # Try PowerShell version check
        try:
            result = await asyncio.wait_for(
                conn.run('powershell.exe -NoProfile -Command "$PSVersionTable.PSVersion.Major"'),
                timeout=10.0,
            )
            if result.exit_status == 0 and (result.stdout or "").strip().isdigit():
                return "powershell"
        except Exception:
            pass

        # Try bash version check
        try:
            result = await asyncio.wait_for(
                conn.run("echo $BASH_VERSION"),
                timeout=10.0,
            )
            if result.exit_status == 0 and (result.stdout or "").strip():
                return "bash"
        except Exception:
            pass

        return "unknown"

    async def is_alive(self) -> bool:
        """Quick connectivity check."""
        try:
            result = await self.run_command("echo alive", timeout=10.0)
            return result.success and "alive" in result.stdout
        except Exception:
            return False
