import asyncio
from asyncio.subprocess import PIPE


class GPUMonitor:
    """Monitors a remote GPU training job via SSH polling."""

    def __init__(self, poll_interval: int = 300):
        self.poll_interval = poll_interval

    async def check_pid(self, host: str, pid: int) -> bool:
        """Check whether a process with the given PID is still running on host.

        Returns True if the process is running, False otherwise.
        """
        command = f"kill -0 {pid} 2>/dev/null && echo running"
        proc = await asyncio.create_subprocess_exec(
            "ssh", host, command,
            stdout=PIPE,
            stderr=PIPE,
        )
        await proc.communicate()
        return proc.returncode == 0

    async def fetch_log(self, host: str, log_path: str) -> str:
        """Fetch the contents of a log file from the remote host.

        Returns the log output as a string.
        """
        command = f"cat {log_path}"
        proc = await asyncio.create_subprocess_exec(
            "ssh", host, command,
            stdout=PIPE,
            stderr=PIPE,
        )
        stdout, _ = await proc.communicate()
        return stdout.decode()

    async def wait_for_completion(self, host: str, pid: int, log_path: str) -> bool:
        """Poll check_pid every poll_interval seconds until the process exits.

        Returns True once the PID is no longer running.
        """
        while True:
            running = await self.check_pid(host, pid)
            if not running:
                return True
            await asyncio.sleep(self.poll_interval)
