#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业 5 · 从 GPU 机拉取 vLLM / nvidia-smi / 排队指标。"""
from __future__ import annotations

from typing import Callable, Optional

import _config as C

Runner = Callable[[str], str]


def gpu_exec(cmd: str, runner: Optional[Runner] = None, timeout: int = 40) -> str:
    if runner is not None:
        return runner(cmd)
    import paramiko

    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(C.GPU_HOST, username=C.GPU_USER, password=C.GPU_PASS, timeout=20)
    try:
        _i, o, e = c.exec_command(cmd, timeout=timeout)
        out = o.read().decode("utf-8", "replace")
        err = e.read().decode("utf-8", "replace")
        if not out.strip() and err.strip():
            raise RuntimeError(err.strip())
        return out
    finally:
        c.close()


def parse_docker_status(raw: str) -> dict:
    line = (raw or "").strip().splitlines()[0] if (raw or "").strip() else ""
    parts = line.split("|")
    if len(parts) < 3:
        raise ValueError("unexpected docker inspect: %r" % raw)
    name = parts[0].lstrip("/")
    running = parts[1].strip().lower() == "true"
    return {
        "container": name,
        "running": running,
        "status": parts[2].strip(),
        "started_at": parts[3].strip() if len(parts) > 3 else "",
        "image": parts[4].strip() if len(parts) > 4 else "",
        "healthy": running and parts[2].strip() == "running",
    }


def parse_nvidia_smi(raw: str) -> dict:
    gpus = []
    for line in (raw or "").splitlines():
        line = line.strip()
        if not line:
            continue
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 6:
            continue
        mem_used = int(parts[4].replace("MiB", "").strip() or "0")
        mem_total = int(parts[5].replace("MiB", "").strip() or "0")
        gpus.append(
            {
                "index": int(parts[0]),
                "name": parts[1],
                "temperature_c": int(parts[2]),
                "utilization": parts[3],
                "memory_used_mib": mem_used,
                "memory_total_mib": mem_total,
                "memory": "%s/%s MiB" % (mem_used, mem_total),
            }
        )
    if not gpus:
        raise ValueError("unexpected nvidia-smi: %r" % raw)
    return {"host": C.GPU_HOST, "gpus": gpus}


def _metric_value(text: str, exact_name: str) -> Optional[float]:
    for line in (text or "").splitlines():
        s = line.strip()
        if not s or s.startswith("#"):
            continue
        name = s.split("{", 1)[0].split(" ", 1)[0]
        if name != exact_name:
            continue
        try:
            return float(s.rsplit(None, 1)[-1])
        except ValueError:
            continue
    return None


def parse_vllm_queue(raw: str) -> dict:
    waiting = _metric_value(raw, "vllm:num_requests_waiting")
    running = _metric_value(raw, "vllm:num_requests_running")
    kv = _metric_value(raw, "vllm:kv_cache_usage_perc")
    return {
        "waiting": 0.0 if waiting is None else waiting,
        "running": 0.0 if running is None else running,
        "kv_cache_usage_perc": kv,
        "source": "vLLM /metrics",
    }


def get_vllm_status(runner: Optional[Runner] = None) -> dict:
    cmd = (
        "docker inspect -f '{{.Name}}|{{.State.Running}}|{{.State.Status}}|"
        "{{.State.StartedAt}}|{{.Config.Image}}' %s"
        % C.GPU_VLLM_CONTAINER
    )
    raw = gpu_exec(cmd, runner=runner)
    out = parse_docker_status(raw)
    out["host"] = C.GPU_HOST
    return out


def get_gpu_info(runner: Optional[Runner] = None) -> dict:
    cmd = (
        "nvidia-smi --query-gpu=index,name,temperature.gpu,utilization.gpu,"
        "memory.used,memory.total --format=csv,noheader"
    )
    raw = gpu_exec(cmd, runner=runner)
    return parse_nvidia_smi(raw)


def get_request_queue(runner: Optional[Runner] = None) -> dict:
    cmd = (
        "curl -s -m 10 -H 'Authorization: Bearer %s' http://127.0.0.1:8000/metrics"
        % C.VLLM_KEY
    )
    raw = gpu_exec(cmd, runner=runner)
    out = parse_vllm_queue(raw)
    out["host"] = C.GPU_HOST
    out["model"] = C.QWEN_MODEL
    return out