"""AsyncSSH-based transport for remote execution.

Service class pattern - dependencies bound at construction,
not passed on every call.
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import re
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from typing import Literal

import asyncssh

from skyward.observability.logger import logger

log = logger.bind(component="ssh")

# =============================================================================
# Stream Event Types (parsed from JSONL)
# =============================================================================


@dataclass(frozen=True, slots=True)
class RawBootstrapConsole:
    """Raw console output from JSONL (without instance info)."""

    content: str
    stream: Literal["stdout", "stderr"] = "stdout"


@dataclass(frozen=True, slots=True)
class RawBootstrapPhase:
    """Raw phase event from JSONL (without instance info)."""

    event: Literal["started", "completed", "failed"]
    phase: str
    elapsed: float | None = None
    error: str | None = None


@dataclass(frozen=True, slots=True)
class RawBootstrapCommand:
    """Raw command event from JSONL (without instance info)."""

    command: str


@dataclass(frozen=True, slots=True)
class RawMetricEvent:
    """Raw metric event from JSONL (without instance info)."""

    name: str
    value: float
    ts: float


@dataclass(frozen=True, slots=True)
class RawLogEvent:
    """Raw log event from JSONL (without instance info)."""

    content: str
    stream: Literal["stdout", "stderr"] = "stdout"


type RawStreamEvent = (
    RawBootstrapConsole
    | RawBootstrapPhase
    | RawBootstrapCommand
    | RawMetricEvent
    | RawLogEvent
)


# =============================================================================
# Exceptions
# =============================================================================


class BootstrapError(Exception):
    """Bootstrap failed."""

    def __init__(self, phase: str, error: str) -> None:
        self.phase = phase
        self.error = error
        super().__init__(f"Bootstrap phase '{phase}' failed: {error}")


# =============================================================================
# SSH Transport
# =============================================================================


@dataclass
class SSHTransport:
    """Async SSH transport using asyncssh.

    Service class that holds connection configuration.
    Dependencies (host, user, key) are bound at construction.

    Retry behavior is built-in to connect(). Configure via retry_* params:
    - retry_max_attempts: Max connection attempts (default: 60 = 5 min at 5s intervals)
    - retry_delay: Delay between attempts in seconds (default: 5.0)

    Example:
        >>> transport = SSHTransport(
        ...     host="10.0.0.1",
        ...     user="ubuntu",
        ...     key_path="~/.ssh/id_rsa",
        ... )
        >>> await transport.connect()  # Retries automatically
        >>> code, stdout, stderr = await transport.run("nvidia-smi")
        >>> await transport.close()

    As context manager:
        >>> async with SSHTransport(...) as t:
        ...     code, out, err = await t.run("ls")
    """

    host: str
    user: str
    key_path: str
    port: int = 22
    connect_timeout: float = 30.0
    retry_max_attempts: int = 150  # 5 min at 2s intervals
    retry_delay: float = 2.0

    _conn: asyncssh.SSHClientConnection | None = field(default=None, repr=False)

    async def connect(self) -> None:
        """Establish SSH connection with automatic retry.

        Uses exponential_base=1.0 (fixed delay) for predictable polling.
        Retries on any exception from asyncssh.connect().
        """
        if self._conn is not None:
            return

        from skyward.infra.retry import retry

        @retry(
            max_attempts=self.retry_max_attempts,
            base_delay=self.retry_delay,
            exponential_base=1.0,  # Fixed delay
            jitter=False,
        )
        async def do_connect() -> asyncssh.SSHClientConnection:
            return await asyncssh.connect(
                self.host,
                port=self.port,
                username=self.user,
                client_keys=[self.key_path],
                known_hosts=None,
                connect_timeout=self.connect_timeout,
            )

        self._conn = await do_connect()
        log.info("SSH connection established to {host}:{port}", host=self.host, port=self.port)

    async def close(self) -> None:
        """Close SSH connection."""
        if self._conn is not None:
            self._conn.close()
            with contextlib.suppress(TimeoutError):
                await asyncio.wait_for(self._conn.wait_closed(), timeout=5.0)
            self._conn = None
            log.info("SSH connection closed to {host}:{port}", host=self.host, port=self.port)

    async def reconnect(self) -> None:
        """Close and re-establish SSH connection."""
        await self.close()
        await self.connect()

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

    async def __aexit__(self, *_: object) -> None:
        await self.close()

    @property
    def is_connected(self) -> bool:
        """Whether connection is established."""
        return self._conn is not None

    def _require_connection(self) -> asyncssh.SSHClientConnection:
        """Get connection or raise."""
        if self._conn is None:
            raise RuntimeError("Not connected. Call connect() first.")
        return self._conn

    # -------------------------------------------------------------------------
    # Command Execution
    # -------------------------------------------------------------------------

    async def run(
        self,
        *command: str,
        timeout: float | None = None,
        check: bool = False,
    ) -> tuple[int, str, str]:
        """Execute command and return (exit_code, stdout, stderr).

        Args:
            command: Command and arguments.
            timeout: Execution timeout in seconds.
            check: If True, raise on non-zero exit code.

        Returns:
            Tuple of (exit_code, stdout, stderr).

        Raises:
            RuntimeError: If check=True and command fails.
        """
        conn = self._require_connection()
        cmd = " ".join(command)
        log.trace("Executing command: {cmd}", cmd=cmd[:200])

        result = await conn.run(cmd, timeout=timeout, check=False)

        code = result.exit_status or 0
        stdout = result.stdout or ""
        stderr = result.stderr or ""

        if check and code != 0:
            raise RuntimeError(f"Command failed ({code}): {stderr}")

        return code, str(stdout), str(stderr)

    async def run_stream(
        self,
        *command: str,
    ) -> AsyncIterator[tuple[str, str]]:
        """Execute command and stream output lines.

        Yields:
            Tuples of ("stdout" | "stderr", line).
        """
        conn = self._require_connection()
        cmd = " ".join(command)

        async with conn.create_process(cmd) as proc:
            # Read stdout and stderr concurrently
            async def read_stream(
                stream: asyncssh.SSHReader[str],
                name: str,
            ) -> AsyncIterator[tuple[str, str]]:
                async for line in stream:
                    yield (name, line.rstrip("\n"))

            stdout_iter = read_stream(proc.stdout, "stdout")
            stderr_iter = read_stream(proc.stderr, "stderr")

            # Simple alternating read
            async for item in self._merge_streams(stdout_iter, stderr_iter):
                yield item

    async def _merge_streams(
        self,
        *iterators: AsyncIterator[tuple[str, str]],
    ) -> AsyncIterator[tuple[str, str]]:
        """Merge multiple async iterators."""
        import asyncio

        pending: set[asyncio.Task[tuple[str, str] | None]] = set()
        iters = {id(it): it.__aiter__() for it in iterators}

        async def get_next(
            it_id: int, it: AsyncIterator[tuple[str, str]],
        ) -> tuple[str, str] | None:
            try:
                return await it.__anext__()
            except StopAsyncIteration:
                iters.pop(it_id, None)
                return None

        # Start initial fetches
        for it_id, it in iters.items():
            task = asyncio.create_task(get_next(it_id, it))
            task.it_id = it_id  # type: ignore[attr-defined]
            pending.add(task)

        while pending:
            done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)

            for task in done:
                result = task.result()
                if result is not None:
                    yield result

                    # Schedule next from same iterator
                    it_id = task.it_id  # type: ignore[attr-defined]
                    if it_id in iters:
                        new_task = asyncio.create_task(get_next(it_id, iters[it_id]))
                        new_task.it_id = it_id  # type: ignore[attr-defined]
                        pending.add(new_task)

    # -------------------------------------------------------------------------
    # File Transfer
    # -------------------------------------------------------------------------

    async def upload(self, local: str, remote: str) -> None:
        """Upload file to remote host.

        Args:
            local: Local file path.
            remote: Remote destination path.
        """
        conn = self._require_connection()
        log.info("Uploading {local} -> {remote}", local=local, remote=remote)
        await asyncssh.scp(local, (conn, remote))

    async def download(self, remote: str, local: str) -> None:
        """Download file from remote host.

        Args:
            remote: Remote file path.
            local: Local destination path.
        """
        conn = self._require_connection()
        log.info("Downloading {remote} -> {local}", remote=remote, local=local)
        await asyncssh.scp((conn, remote), local)

    async def read_file(self, remote: str) -> str:
        """Read remote file contents."""
        _, stdout, _ = await self.run("cat", remote)
        return stdout

    async def write_file(self, remote: str, content: str) -> None:
        """Write content to remote file using SFTP."""
        conn = self._require_connection()

        async with conn.start_sftp_client() as sftp, sftp.open(remote, "w") as f:
            await f.write(content)

    async def write_bytes(self, remote: str, content: bytes) -> None:
        """Write binary content to remote file using SFTP."""
        conn = self._require_connection()

        async with conn.start_sftp_client() as sftp, sftp.open(remote, "wb") as f:
            await f.write(content)

    async def file_exists(self, remote: str) -> bool:
        """Check if remote file exists."""
        code, _, _ = await self.run("test", "-f", remote)
        return code == 0

    # -------------------------------------------------------------------------
    # Convenience Methods
    # -------------------------------------------------------------------------

    async def wait_for_file(
        self,
        remote: str,
        timeout: float = 300,
        poll_interval: float = 5,
    ) -> bool:
        """Wait for remote file to exist.

        Returns:
            True if file appeared, False if timeout.
        """
        import asyncio

        start = asyncio.get_event_loop().time()
        while asyncio.get_event_loop().time() - start < timeout:
            if await self.file_exists(remote):
                return True
            await asyncio.sleep(poll_interval)
        return False

    async def tail_file(
        self,
        remote: str,
        follow: bool = True,
    ) -> AsyncIterator[str]:
        """Stream lines from remote file.

        Args:
            remote: Remote file path.
            follow: If True, use tail -f (stream indefinitely).

        Yields:
            Lines from the file.
        """
        cmd = ["tail", "-f", remote] if follow else ["tail", remote]
        async for stream, line in self.run_stream(*cmd):
            if stream == "stdout":
                yield line

    # -------------------------------------------------------------------------
    # Events Streaming (JSONL)
    # -------------------------------------------------------------------------

    async def stream_events(
        self,
        log_path: str = "/opt/skyward/events.jsonl",
        timeout: float = 600.0,
        wait_for_file_timeout: float = 120.0,
        start_line: int = 0,
    ) -> AsyncIterator[tuple[int, RawStreamEvent]]:
        """Stream events from JSONL log file via tail -F with offset tracking.

        Uses tail -n +{start_line+1} -F to read from a specific line offset,
        ensuring no events are lost on first connection or reconnection.

        The timeout is adaptive: it resets on each event received.

        Args:
            log_path: Path to events.jsonl on remote.
            timeout: Maximum time between events before timeout.
            wait_for_file_timeout: Max time to wait for log file to exist.
            start_line: Line offset to start reading from (0 = beginning).

        Yields:
            Tuples of (lines_read, event) where lines_read is the cumulative
            count of lines processed (for reconnection tracking).

        Raises:
            BootstrapError: If a phase fails.
            TimeoutError: If no events received within timeout.
            FileNotFoundError: If log file doesn't appear in time.
        """
        await self._wait_for_log_file(log_path, wait_for_file_timeout)

        conn = self._require_connection()
        tail_offset = start_line + 1
        logger.debug("Starting tail -n +{offset} -F {path}", offset=tail_offset, path=log_path)

        async with conn.create_process(f"tail -n +{tail_offset} -F {log_path}") as proc:
            deadline = asyncio.get_event_loop().time() + timeout
            buffer = ""
            lines_read = start_line

            while True:
                remaining = deadline - asyncio.get_event_loop().time()
                if remaining <= 0:
                    raise TimeoutError(f"Events stream timeout after {timeout}s of inactivity")

                try:
                    chunk = await asyncio.wait_for(
                        proc.stdout.read(4096),
                        timeout=min(1.0, remaining),
                    )

                    if not chunk:
                        break

                    buffer += chunk

                    while "\n" in buffer:
                        line, buffer = buffer.split("\n", 1)
                        lines_read += 1
                        line = line.strip()
                        if not line:
                            continue

                        event = _parse_jsonl_line(line)
                        if event is None:
                            continue

                        yield lines_read, event

                        deadline = asyncio.get_event_loop().time() + timeout

                        match event:
                            case RawBootstrapPhase(event="failed", phase=phase, error=error):
                                raise BootstrapError(phase, error or "unknown")

                except TimeoutError:
                    if asyncio.get_event_loop().time() >= deadline:
                        raise TimeoutError(
                            f"Events stream timeout after {timeout}s of inactivity"
                        ) from None

    async def stream_bootstrap(
        self,
        log_path: str = "/opt/skyward/events.jsonl",
        timeout: float = 600.0,
        start_line: int = 0,
    ) -> AsyncIterator[tuple[int, RawStreamEvent]]:
        """Stream bootstrap events until completion.

        Convenience wrapper around stream_events() that stops when
        bootstrap phase completes successfully.

        Args:
            log_path: Path to events.jsonl on remote.
            timeout: Maximum time between events.
            start_line: Line offset to start reading from.

        Yields:
            Tuples of (lines_read, event) until bootstrap completion.

        Raises:
            BootstrapError: If bootstrap fails.
            TimeoutError: If timeout exceeded.
        """
        async for lines_read, event in self.stream_events(log_path, timeout, start_line=start_line):
            yield lines_read, event

            match event:
                case RawBootstrapPhase(phase="bootstrap", event="completed"):
                    return

    async def _wait_for_log_file(self, log_path: str, timeout: float) -> None:
        """Wait for log file to exist on remote."""
        start = asyncio.get_event_loop().time()
        check_count = 0

        logger.debug("Waiting for log file {path}", path=log_path)

        while asyncio.get_event_loop().time() - start < timeout:
            check_count += 1
            if await self.file_exists(log_path):
                elapsed = asyncio.get_event_loop().time() - start
                logger.debug(
                    "Log file found after {elapsed:.1f}s ({checks} checks)",
                    elapsed=elapsed, checks=check_count,
                )
                return
            await asyncio.sleep(1.0)

        raise FileNotFoundError(f"Log file {log_path} not created within {timeout}s")


# =============================================================================
# JSONL Parsing
# =============================================================================


_ANSI_RE = re.compile(r"\x1b\[[0-9;?]*[a-zA-Z]")
_CTRL_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")


def _parse_jsonl_line(line: str) -> RawStreamEvent | None:
    """Parse a single JSONL line from events log.

    Args:
        line: Raw line from events.jsonl.

    Returns:
        Parsed event or None if line is invalid.
    """
    try:
        data = json.loads(line)
    except json.JSONDecodeError:
        sanitized = _CTRL_RE.sub("", _ANSI_RE.sub("", line))
        try:
            data = json.loads(sanitized)
        except json.JSONDecodeError:
            log.warning("Failed to parse JSONL line: {line}", line=line[:200])
            return None

    match data.get("type"):
        case "console":
            return RawBootstrapConsole(
                content=data.get("content", ""),
                stream=data.get("stream", "stdout"),
            )
        case "phase":
            return RawBootstrapPhase(
                event=data.get("event", "started"),
                phase=data.get("phase", ""),
                elapsed=data.get("elapsed"),
                error=data.get("error"),
            )
        case "command":
            return RawBootstrapCommand(command=data.get("command", ""))
        case "metric":
            value = data.get("value")
            if value is None:
                return None
            try:
                return RawMetricEvent(
                    name=data.get("name", ""),
                    value=float(value),
                    ts=data.get("ts", 0.0),
                )
            except (ValueError, TypeError):
                return None
        case "log":
            return RawLogEvent(
                content=data.get("content", ""),
                stream=data.get("stream", "stdout"),
            )
        case _:
            return None
