from __future__ import annotations

import subprocess
from collections.abc import Sequence
from dataclasses import dataclass


@dataclass
class CollectResult:
    """Raw collection output for one host.

    Attributes:
        host: Host name used for collection.
        stdout: Command standard output.
        stderr: Command standard error.
        error: Normalized error message when collection failed; otherwise None.
    """

    host: str
    stdout: str
    stderr: str
    error: str | None = None


def collect_from_host(
    host: str,
    ssh_timeout: float,
    cmd_timeout: float,
    hy_smi_args: Sequence[str],
) -> CollectResult:
    """Run hy-smi locally or via SSH and return raw output.

    Args:
        host: Hostname or localhost alias.
        ssh_timeout: SSH connect timeout in seconds.
        cmd_timeout: Command timeout in seconds.

    Returns:
        Raw command output with normalized error information.
    """
    local_names = {"localhost", "127.0.0.1", "::1"}
    if host in local_names:
        cmd = ["hy-smi", *hy_smi_args]
    else:
        connect_timeout = max(1, round(ssh_timeout))
        cmd = [
            "ssh",
            "-o",
            "BatchMode=yes",
            "-o",
            f"ConnectTimeout={connect_timeout}",
            host,
            "hy-smi",
            *hy_smi_args,
        ]

    try:
        proc = subprocess.run(
            cmd,
            check=False,
            capture_output=True,
            text=True,
            timeout=cmd_timeout,
        )
    except subprocess.TimeoutExpired:
        return CollectResult(
            host=host,
            stdout="",
            stderr="",
            error=f"timeout after {cmd_timeout:.1f}s",
        )
    except OSError as exc:
        return CollectResult(host=host, stdout="", stderr="", error=str(exc))

    if proc.returncode != 0:
        stderr = proc.stderr.strip() or "unknown error"
        return CollectResult(
            host=host,
            stdout=proc.stdout,
            stderr=proc.stderr,
            error=f"exit {proc.returncode}: {stderr}",
        )

    return CollectResult(host=host, stdout=proc.stdout, stderr=proc.stderr, error=None)
