"""GPU tenants: the applications competing for the card, described as data.

The point of this service is fast handoff of a single GPU between applications. It grew
up around the two on this box, and their names ended up compiled into process matching,
VRAM attribution, busy detection and release calls alike -- roughly 385 references. That
makes it a script for Ollama and ComfyUI rather than a GPU arbitrator.

A tenant is described here instead:

  * how to recognise its processes            (match)
  * how to tell whether it is actually working (busy probe)
  * how to ask it to give VRAM back            (release strategy)
  * how much it matters when two want the card (priority)

Ollama and ComfyUI ship as defaults so behaviour is unchanged, but nothing about the
arbitration logic knows their names. A third application -- a training run, a speech
model, another inference server -- is a config entry, not a code change. A tenant that
cannot be released (no API to ask) is still worth declaring, because naming it turns
"unmanaged VRAM" into "held by X, which cannot be reclaimed".
"""
import json
import logging
import os
import time
from dataclasses import dataclass, field, asdict
from typing import Any, Dict, List, Optional

import httpx
import psutil

logger = logging.getLogger("tenants")

_BASE = os.path.dirname(os.path.abspath(__file__))
CONFIG_PATH = os.environ.get("HYPERSWAP_TENANTS", os.path.join(_BASE, "tenants.json"))

# Kinds are advisory: they drive presentation and sensible defaults, never control flow.
KIND_LLM, KIND_DIFFUSION, KIND_DESKTOP, KIND_OTHER = "llm", "diffusion", "desktop", "other"


@dataclass
class ProcessMatch:
    """How to recognise a tenant's processes among those NVML reports."""
    names: List[str] = field(default_factory=list)      # matched against process name
    cmdline: List[str] = field(default_factory=list)    # substrings of the full cmdline
    cmdline_endswith: List[str] = field(default_factory=list)

    def matches(self, pname: str, cmdline: str) -> bool:
        pname, cmdline = pname.lower(), cmdline.lower()
        if any(n.lower() in pname for n in self.names):
            return True
        if any(c.lower() in cmdline for c in self.cmdline):
            return True
        return any(cmdline.rstrip().endswith(c.lower()) for c in self.cmdline_endswith)


@dataclass
class BusyProbe:
    """How to tell whether a tenant is genuinely working.

    `vram_floor_gb` exists because a queue flag can lie: ComfyUI leaves dead jobs in
    queue_running, and only its VRAM reveals that nothing is loaded. GPU utilisation is
    deliberately unavailable as a signal -- it is shared by every tenant, so it cannot
    attribute work to one of them.
    """
    type: str = "none"                  # none | http_count | vram
    url: Optional[str] = None
    count_keys: List[str] = field(default_factory=list)  # keys whose lists are summed
    vram_busy_gb: float = 0.0           # busy when its VRAM exceeds this
    vram_floor_gb: float = 0.0          # below this it holds no real work
    stale_after_s: float = 90.0


@dataclass
class ReleaseStrategy:
    """How to ask a tenant to give VRAM back."""
    type: str = "none"                  # none | http_post
    url: Optional[str] = None
    body: Dict[str, Any] = field(default_factory=dict)
    # Set when the call must name the loaded model (Ollama unloads per model).
    per_model: bool = False
    timeout_s: float = 120.0
    confirm: bool = True                # wait for NVML to show the memory released


@dataclass
class GpuTenant:
    name: str
    kind: str = KIND_OTHER
    enabled: bool = True
    # Higher wins contention; a tenant yields to anything above it.
    priority: int = 50
    match: ProcessMatch = field(default_factory=ProcessMatch)
    busy: BusyProbe = field(default_factory=BusyProbe)
    release: ReleaseStrategy = field(default_factory=ReleaseStrategy)
    notes: str = ""

    @property
    def reclaimable(self) -> bool:
        return self.release.type != "none"

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


def _tenant_from_dict(d: Dict[str, Any]) -> GpuTenant:
    return GpuTenant(
        name=d["name"],
        kind=d.get("kind", KIND_OTHER),
        enabled=d.get("enabled", True),
        priority=int(d.get("priority", 50)),
        match=ProcessMatch(**(d.get("match") or {})),
        busy=BusyProbe(**(d.get("busy") or {})),
        release=ReleaseStrategy(**(d.get("release") or {})),
        notes=d.get("notes", ""),
    )


# Defaults reproduce today's behaviour exactly; they are data, not special cases.
DEFAULT_TENANTS: List[Dict[str, Any]] = [
    {
        "name": "ollama",
        "kind": KIND_LLM,
        "priority": 60,
        "match": {"names": ["ollama"], "cmdline": ["llama-server", "ollama"]},
        "busy": {"type": "http_count", "url": "http://localhost:11434/api/ps",
                 "count_keys": ["models"]},
        "release": {"type": "http_post", "url": "http://localhost:11434/api/generate",
                    "body": {"keep_alive": 0}, "per_model": True, "timeout_s": 120.0},
        "notes": "Unloads per model. With OLLAMA_NUM_PARALLEL=1 the request queues "
                 "behind any running generation and applies when it finishes.",
    },
    {
        "name": "comfyui",
        "kind": KIND_DIFFUSION,
        "priority": 50,
        "match": {"cmdline": ["comfyui", "comfy"], "cmdline_endswith": ["main.py"]},
        "busy": {"type": "http_count", "url": "http://127.0.0.1:8188/queue",
                 "count_keys": ["queue_running", "queue_pending"],
                 "vram_floor_gb": 1.5, "stale_after_s": 90.0},
        "release": {"type": "http_post", "url": "http://127.0.0.1:8188/free",
                    "body": {"unload_models": True, "free_memory": True},
                    "timeout_s": 30.0},
        "notes": "Leaves dead jobs in queue_running; the queue flag is corroborated "
                 "against its own VRAM before being believed.",
    },
    {
        "name": "desktop",
        "kind": KIND_DESKTOP,
        "priority": 90,
        "match": {"names": ["gnome-shell", "xorg", "mutter", "kwin", "plasmashell",
                            "gnome-remote-desktop", "sddm", "gdm", "picom", "weston"]},
        "release": {"type": "none"},
        "notes": "Compositor and display server. Small, permanent, never reclaimable.",
    },
]


_cache: Dict[str, Any] = {"ts": 0.0, "tenants": None, "mtime": None}
CACHE_TTL_S = 10.0


def load_tenants(force: bool = False) -> List[GpuTenant]:
    """Load tenant definitions, writing the defaults out on first run."""
    now = time.time()
    try:
        mtime = os.path.getmtime(CONFIG_PATH) if os.path.exists(CONFIG_PATH) else None
    except OSError:
        mtime = None
    if (not force and _cache["tenants"] is not None
            and mtime == _cache["mtime"] and (now - _cache["ts"]) < CACHE_TTL_S):
        return _cache["tenants"]

    raw: List[Dict[str, Any]]
    if os.path.exists(CONFIG_PATH):
        try:
            with open(CONFIG_PATH) as f:
                raw = json.load(f)
        except Exception as e:
            logger.error(f"could not read {CONFIG_PATH}, using defaults: {e}")
            raw = DEFAULT_TENANTS
    else:
        raw = DEFAULT_TENANTS
        try:
            with open(CONFIG_PATH, "w") as f:
                json.dump(DEFAULT_TENANTS, f, indent=2)
            logger.info(f"wrote default tenant definitions to {CONFIG_PATH}")
        except Exception as e:
            logger.warning(f"could not write {CONFIG_PATH}: {e}")

    tenants = []
    for d in raw:
        try:
            tenants.append(_tenant_from_dict(d))
        except Exception as e:
            logger.error(f"skipping malformed tenant {d!r}: {e}")
    _cache.update({"ts": now, "tenants": tenants, "mtime": mtime})
    return tenants


def get_tenant(name: str) -> Optional[GpuTenant]:
    return next((t for t in load_tenants() if t.name == name), None)


def save_tenants(tenants: List[Dict[str, Any]]) -> bool:
    try:
        with open(CONFIG_PATH, "w") as f:
            json.dump(tenants, f, indent=2)
        _cache["tenants"] = None
        return True
    except Exception as e:
        logger.error(f"save_tenants failed: {e}")
        return False


def classify_process(pname: str, cmdline: str) -> str:
    """Return the owning tenant's name, or 'unmanaged'.

    'unmanaged' is meaningful rather than a dumping ground: it is VRAM this service has
    no way to reclaim, and it is reported as such.
    """
    for t in load_tenants():
        if t.enabled and t.match.matches(pname, cmdline):
            return t.name
    return "unmanaged"


def classify_pid(pid: int) -> str:
    try:
        proc = psutil.Process(pid)
        return classify_process(proc.name(), " ".join(proc.cmdline()))
    except Exception:
        return "unmanaged"


async def probe_busy(tenant: GpuTenant, vram_gb: float = 0.0,
                     state: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Is this tenant actually working? Returns {busy, reason, stale}."""
    probe = tenant.busy
    if probe.type == "vram":
        busy = vram_gb > probe.vram_busy_gb
        return {"busy": busy, "reason": f"{vram_gb:.2f} GB held", "stale": False}
    if probe.type != "http_count" or not probe.url:
        return {"busy": False, "reason": "no busy probe configured", "stale": False}

    try:
        async with httpx.AsyncClient(timeout=3.0) as c:
            r = await c.get(probe.url)
        if r.status_code != 200:
            return {"busy": False, "reason": f"probe HTTP {r.status_code}", "stale": False}
        data = r.json()
        count = sum(len(data.get(k) or []) for k in probe.count_keys)
    except Exception as e:
        return {"busy": False, "reason": f"probe failed: {str(e)[:60]}", "stale": False}

    if count == 0:
        return {"busy": False, "reason": "queue empty", "stale": False}
    # A queue that claims work while the tenant holds no VRAM is not doing work.
    if probe.vram_floor_gb and vram_gb < probe.vram_floor_gb:
        return {"busy": True, "reason": f"{count} queued, holding {vram_gb:.2f} GB",
                "stale": None, "below_floor": True}
    return {"busy": True, "reason": f"{count} queued/running", "stale": False}


async def release_vram(tenant: GpuTenant, models: Optional[List[str]] = None
                       ) -> Dict[str, Any]:
    """Ask a tenant to give its VRAM back, however that tenant expects to be asked."""
    strategy = tenant.release
    if strategy.type == "none" or not strategy.url:
        return {"success": False, "tenant": tenant.name, "released": False,
                "reason": "this tenant exposes no way to release VRAM"}

    t0 = time.perf_counter()
    payloads: List[Dict[str, Any]] = []
    if strategy.per_model:
        for m in (models or []):
            payloads.append({**strategy.body, "model": m})
        if not payloads:
            return {"success": True, "tenant": tenant.name, "released": False,
                    "reason": "nothing loaded to release"}
    else:
        payloads.append(dict(strategy.body))

    errors = []
    try:
        async with httpx.AsyncClient(timeout=strategy.timeout_s) as c:
            for body in payloads:
                try:
                    await c.post(strategy.url, json=body)
                except Exception as e:
                    errors.append(str(e)[:80])
    except Exception as e:
        errors.append(str(e)[:80])

    return {
        "success": not errors,
        "tenant": tenant.name,
        "released": True,
        "requests": len(payloads),
        "duration_ms": round((time.perf_counter() - t0) * 1000, 2),
        "errors": errors or None,
    }


def describe() -> List[Dict[str, Any]]:
    """Tenant definitions for the API, with what each can and cannot do."""
    out = []
    for t in sorted(load_tenants(), key=lambda x: -x.priority):
        d = t.to_dict()
        d["reclaimable"] = t.reclaimable
        d["busy_probe"] = t.busy.type
        d["release_via"] = t.release.type
        out.append(d)
    return out