"""
ARIA Node-Stats — Auslastungs-Monitor pro Box (GPU + optional Tokens).

Identische Kopie in jedem Worker-Build-Context (f5tts/whisper/voxtral/llm-adapter),
weil jeder Worker ein eigener Docker-Build-Context ist.

Aufgaben:
- Sampler-Loop (alle SAMPLE_SEC): nvidia-smi-Auslastung + Token-Delta → Ringpuffer
  (persistent als JSON auf der Box). Laeuft unabhaengig vom Modal.
- Live-Stream: bei node_stats_stream_start jede Sekunde rohes nvidia-smi + Werte
  senden (bis stop / Auto-Timeout).
- History-Request + Reset (Besen).

Reicht `handle(ws, mtype, payload)` in die Worker-Message-Loop ein; gibt True
zurueck, wenn die Nachricht eine node_stats_*-Nachricht war.
"""
import asyncio
import json
import os
import time

SAMPLE_SEC = int(os.getenv("STATS_SAMPLE_SEC", "15"))
HISTORY_CAP = int(os.getenv("STATS_HISTORY_CAP", "500"))   # ~2h bei 15s
STREAM_MAX_SEC = int(os.getenv("STATS_STREAM_MAX_SEC", "300"))


async def _run_cmd(*args, timeout=8) -> str:
    """Fuehrt ein Kommando aus, gibt stdout (str) zurueck; '' bei Fehler."""
    try:
        proc = await asyncio.create_subprocess_exec(
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.DEVNULL,
        )
        out, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout)
        return (out or b"").decode("utf-8", "replace")
    except Exception:
        return ""


class NodeStats:
    def __init__(self, instance_id: str, node_name: str, history_path: str,
                 token_getter=None, logger=None):
        self.instance_id = instance_id
        self.node_name = node_name
        self.history_path = history_path
        self.token_getter = token_getter   # callable -> kumulative Token-Zahl (oder None)
        self.log = logger
        self.samples = self._load()
        self._last_tokens = self._tokens_now()
        self._stream_task = None

    # ── Persistenz ──────────────────────────────────────────
    def _load(self) -> list:
        try:
            with open(self.history_path) as f:
                data = json.load(f)
            return data if isinstance(data, list) else []
        except Exception:
            return []

    def _persist(self) -> None:
        try:
            os.makedirs(os.path.dirname(self.history_path) or ".", exist_ok=True)
            tmp = self.history_path + ".tmp"
            with open(tmp, "w") as f:
                json.dump(self.samples[-HISTORY_CAP:], f)
            os.replace(tmp, self.history_path)
        except Exception:
            pass

    def _tokens_now(self) -> int:
        try:
            return int(self.token_getter()) if self.token_getter else 0
        except Exception:
            return 0

    # ── nvidia-smi ──────────────────────────────────────────
    async def _query_gpu(self) -> dict:
        """Aggregierte GPU-Werte ueber alle sichtbaren Karten."""
        out = await _run_cmd(
            "nvidia-smi",
            "--query-gpu=utilization.gpu,memory.used,memory.total",
            "--format=csv,noheader,nounits")
        utils, used, total = [], 0, 0
        for line in out.strip().splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) < 3:
                continue
            try:
                utils.append(float(parts[0]))
                used += float(parts[1])
                total += float(parts[2])
            except ValueError:
                continue
        gpu = round(sum(utils) / len(utils), 1) if utils else 0.0
        return {"gpu": gpu, "memUsed": int(used), "memTotal": int(total)}

    async def _nvidia_smi_text(self) -> str:
        txt = await _run_cmd("nvidia-smi")
        return txt or "nvidia-smi nicht verfuegbar"

    # ── Sampler (Verlauf) ───────────────────────────────────
    async def run_sampler(self) -> None:
        while True:
            try:
                g = await self._query_gpu()
                now_tok = self._tokens_now()
                dtok = max(0, now_tok - self._last_tokens)
                self._last_tokens = now_tok
                self.samples.append({
                    "ts": int(time.time()),
                    "gpu": g["gpu"], "memUsed": g["memUsed"],
                    "memTotal": g["memTotal"], "tokens": dtok,
                })
                if len(self.samples) > HISTORY_CAP:
                    self.samples = self.samples[-HISTORY_CAP:]
                self._persist()
            except Exception as e:
                if self.log:
                    self.log.debug("node_stats sample fehlgeschlagen: %s", e)
            await asyncio.sleep(SAMPLE_SEC)

    # ── Live-Stream ─────────────────────────────────────────
    async def _stream(self, ws, send) -> None:
        t0 = time.time()
        try:
            while time.time() - t0 < STREAM_MAX_SEC:
                g = await self._query_gpu()
                smi = await self._nvidia_smi_text()
                await send(ws, "node_stats", {
                    "instanceId": self.instance_id, "node": self.node_name,
                    "nvidiaSmi": smi, **g,
                })
                await asyncio.sleep(1)
        except asyncio.CancelledError:
            raise
        except Exception:
            return

    # ── Dispatch ────────────────────────────────────────────
    async def handle(self, ws, mtype: str, payload: dict, send) -> bool:
        if mtype == "node_stats_stream_start":
            if self._stream_task and not self._stream_task.done():
                self._stream_task.cancel()
            self._stream_task = asyncio.create_task(self._stream(ws, send))
            return True
        if mtype == "node_stats_stream_stop":
            if self._stream_task:
                self._stream_task.cancel()
                self._stream_task = None
            return True
        if mtype == "node_stats_history_request":
            await send(ws, "node_stats_history", {
                "instanceId": self.instance_id, "node": self.node_name,
                "samples": self.samples[-HISTORY_CAP:],
                "tokenCapable": self.token_getter is not None,
                "sampleSec": SAMPLE_SEC,
            })
            return True
        if mtype == "node_stats_reset":
            self.samples = []
            self._persist()
            await send(ws, "node_stats_reset_done",
                       {"instanceId": self.instance_id, "node": self.node_name})
            return True
        return False