import asyncio
import re
import time
from typing import Annotated, 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: Annotated[str, Form()], profile: Annotated[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: Annotated[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,
)