from __future__ import annotations

import asyncio
from typing import Protocol, runtime_checkable


@runtime_checkable
class Connection(Protocol):
    async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]: ...
    async def push(self, local_path: str, remote_path: str) -> None: ...
    async def pull(self, remote_path: str, local_path: str) -> None: ...
    async def close(self) -> None: ...


@runtime_checkable
class Transport(Protocol):
    async def connect(self, host: str, port: int) -> Connection: ...


class SSHConnection:
    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port

    async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]:
        proc = await asyncio.create_subprocess_exec(
            "ssh",
            "-p", str(self.port),
            "-o", "StrictHostKeyChecking=no",
            "-o", "ConnectTimeout=5",
            "-o", "BatchMode=yes",
            f"root@{self.host}",
            cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=timeout
            )
            return proc.returncode, (stdout + stderr).decode()
        except asyncio.TimeoutError:
            proc.kill()
            return 1, "timeout"

    async def push(self, local_path: str, remote_path: str) -> None:
        proc = await asyncio.create_subprocess_exec(
            "scp",
            "-P", str(self.port),
            "-o", "StrictHostKeyChecking=no",
            local_path,
            f"root@{self.host}:{remote_path}",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        _, stderr = await proc.communicate()
        if proc.returncode != 0:
            raise ConnectionError(
                f"scp push failed to {self.host}:{self.port}: {stderr.decode()}"
            )

    async def pull(self, remote_path: str, local_path: str) -> None:
        proc = await asyncio.create_subprocess_exec(
            "scp",
            "-P", str(self.port),
            "-o", "StrictHostKeyChecking=no",
            f"root@{self.host}:{remote_path}",
            local_path,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        _, stderr = await proc.communicate()
        if proc.returncode != 0:
            raise ConnectionError(
                f"scp pull failed from {self.host}:{self.port}: {stderr.decode()}"
            )

    async def close(self) -> None:
        pass


class SSHTransport:
    async def connect(self, host: str, port: int) -> SSHConnection:
        return SSHConnection(host, port)