import asyncio
import re
import time
from typing import Any

from fastapi import APIRouter, Form
from fastapi.responses import HTMLResponse

from app.collect.gpu import _shorten
from app.plugins.base import Plugin
from app.render import render

router = APIRouter(prefix="/api/plugins/lact", tags=["plugins"])

LIST_TIMEOUT = 5
SET_TIMEOUT = 15
GPU_CACHE_TTL = 60

_gpu_cache: tuple[float, list[dict[str, str]]] | None = None
_set_lock = asyncio.Lock()


async def _run(args: list[str], timeout: float) -> tuple[str, str]:
    try:
        proc = await asyncio.create_subprocess_exec(
            "lact", "cli", *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
    except FileNotFoundError:
        return "", "lact not found in PATH"
    except OSError as e:
        return "", str(e)[:200]
    try:
        out, err = await asyncio.wait_for(proc.communicate(), timeout)
    except TimeoutError:
        try:
            proc.kill()
        except ProcessLookupError:
            pass
        return "", "lact timed out"
    if proc.returncode != 0:
        return "", (err.decode(errors="replace").strip() or f"lact failed (rc={proc.returncode})")[:200]
    return out.decode(errors="replace"), ""


def _parse_gpus(out: str) -> list[dict[str, str]]:
    gpus: list[dict[str, str]] = []
    for line in out.splitlines():
        m = re.match(r"^\s*(\d+):\s+(.*)$", line)
        if not m:
            continue
        rest = m.group(2)
        name = re.search(r"\(([^)]*)\)", rest)
        gpu_type = re.search(r"\[([^\]]*)\]\s*$", rest)
        gpus.append({
            "id": m.group(1),
            "name": _shorten(name.group(1)) if name else rest.strip(),
            "type": gpu_type.group(1) if gpu_type else "",
        })
    return gpus


async def _gpus(force: bool = False) -> tuple[list[dict[str, str]], str]:
    global _gpu_cache
    if not force and _gpu_cache is not None:
        ts, cached = _gpu_cache
        if time.monotonic() - ts < GPU_CACHE_TTL:
            return [dict(g) for g in cached], ""
    out, err = await _run(["list"], LIST_TIMEOUT)
    if err:
        return [], err
    gpus = _parse_gpus(out)
    _gpu_cache = (time.monotonic(), gpus)
    return gpus, ""


async def _gpu_entry(g: dict[str, str], with_profiles: bool) -> dict[str, Any]:
    entry: dict[str, Any] = {**g, "profiles": [], "active": None, "error": ""}
    base = ["--gpu-id", g["id"], "profile"]
    if with_profiles:
        active_p, profiles_p = await asyncio.gather(
            _run([*base, "get"], LIST_TIMEOUT),
            _run([*base, "list"], LIST_TIMEOUT),
        )
    else:
        active_p, profiles_p = await _run([*base, "get"], LIST_TIMEOUT), None
    out, err = active_p
    if err:
        entry["error"] = err
        return entry
    entry["active"] = out.strip() or None
    if profiles_p is not None:
        out2, err2 = profiles_p
        entry["error"] = err2
        if not err2:
            entry["profiles"] = [line.strip() for line in out2.splitlines() if line.strip()]
    return entry


async def _gather(with_profiles: bool, force_gpus: bool = False) -> dict[str, Any]:
    gpus, err = await _gpus(force=force_gpus)
    if err:
        return {"gpus": [], "error": err}
    entries = await asyncio.gather(*[_gpu_entry(g, with_profiles) for g in gpus])
    return {"gpus": list(entries), "error": ""}


async def _state(message: str = "", error: str = "") -> str:
    data = await _gather(with_profiles=False)
    data["message"] = message
    data["error"] = error or data["error"]
    return render("plugins/lact_state.html", **data)


async def _skeleton(message: str = "", error: str = "") -> str:
    data = await _gather(with_profiles=True, force_gpus=True)
    data["message"] = message
    data["error"] = error or data["error"]
    return render("plugins/lact_skeleton.html", **data)


@router.get("/state")
async def state():
    return HTMLResponse(await _state())


@router.post("/set")
async def set_profile(gpu_id: str = Form(...), profile: str = Form(...)):
    async with _set_lock:
        data = await _gather(with_profiles=True, force_gpus=True)
        if data["error"]:
            return HTMLResponse(await _skeleton(error=data["error"]))
        gpu = next((g for g in data["gpus"] if g["id"] == gpu_id), None)
        if gpu is None:
            return HTMLResponse(await _skeleton(error=f"unknown gpu id: {gpu_id}"))
        if profile not in gpu["profiles"]:
            return HTMLResponse(await _skeleton(error=f"unknown profile: {profile}"))
        if profile == gpu["active"]:
            return HTMLResponse(await _skeleton(message=f"{gpu['name']}: {profile} already active"))
        _out, err = await _run(["--gpu-id", gpu_id, "profile", "set", profile], SET_TIMEOUT)
        if err:
            return HTMLResponse(await _skeleton(error=err))
        return HTMLResponse(await _skeleton(message=f"{gpu['name']}: set profile {profile}"))


@router.post("/reload")
async def reload(gpu_id: str | None = Form(None)):
    _ = gpu_id
    return HTMLResponse(await _skeleton(message="profiles refreshed"))


plugin = Plugin(
    id="lact",
    title="GPU power profiles",
    description="Active profile per GPU and a profile switcher (lact cli).",
    skeleton_fn=_skeleton,
)