#!/usr/bin/env python3
"""First-boot Swarm/Comfy performance tune: sageattention ExtraArgs + pip libs.

Idempotent. Marker: /mnt/swarm_data/.gpu-rent-perf-tuned
Re-runs if GPU uuid changed.
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys
from pathlib import Path

DATA = Path("/mnt/swarm_data")
GPU_JSON = DATA / ".gpu-rent-gpu.json"
MARKER = DATA / ".gpu-rent-perf-tuned"
BACKENDS = DATA / "Data" / "Backends.fds"
COMFY_VENV_CANDIDATES = [
    DATA / "dlbackend" / "ComfyUI" / "venv" / "bin" / "pip",
    DATA / "dlbackend" / "comfy" / "venv" / "bin" / "pip",
    DATA / "dlbackend" / "comfy" / "ComfyUI" / "venv" / "bin" / "pip",
    Path("/opt/swarmui/dlbackend/ComfyUI/venv/bin/pip"),
    Path("/opt/swarmui/dlbackend/comfy/venv/bin/pip"),
    Path("/opt/swarmui/dlbackend/comfy/ComfyUI/venv/bin/pip"),
]


def load_gpu() -> dict:
    if not GPU_JSON.is_file():
        return {}
    try:
        return json.loads(GPU_JSON.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return {}


def tier_notes(gpu: dict) -> dict:
    """Mirror gpu_rent.perf_tiers.swarm_tune_for (stdlib-only on VM)."""
    vram = int(gpu.get("vram_mib") or 0)
    gib = vram / 1024.0
    if gib < 16:
        tier = "low"
    elif gib < 24:
        tier = "mid"
    elif gib < 48:
        tier = "high"
    else:
        tier = "ultra"
    cap = str(gpu.get("compute_cap") or "0.0")
    try:
        parts = cap.split(".")
        maj, mnr = int(parts[0]), int(parts[1]) if len(parts) > 1 else 0
        ampere = (maj, mnr) >= (8, 0)
    except ValueError:
        ampere = False
    use_sage = ampere and tier in {"mid", "high", "ultra"}
    return {
        "tier": tier,
        "use_sage": use_sage,
        "extra_args": "--use-sage-attention" if use_sage else "",
        "uuid": str(gpu.get("uuid") or ""),
        "name": str(gpu.get("name") or ""),
    }


def find_pip() -> Path | None:
    for p in COMFY_VENV_CANDIDATES:
        if p.is_file():
            return p
    # glob
    for p in DATA.glob("dlbackend/**/venv/bin/pip"):
        return p
    return None


def patch_backends_extra_args(extra: str) -> bool:
    """Ensure ExtraArgs contains sage flag for Comfy self-start backends."""
    if not extra:
        return False
    if not BACKENDS.is_file():
        print(f"no {BACKENDS} yet — skip ExtraArgs (Comfy not registered)")
        return False
    text = BACKENDS.read_text(encoding="utf-8")
    if "--use-sage-attention" in text:
        print("Backends.fds already has --use-sage-attention")
        return False
    lines = text.splitlines()
    changed = False
    out = []
    for line in lines:
        if re.match(r"^(\s*)ExtraArgs:\s*$", line) or re.match(r"^(\s*)ExtraArgs:\s*\"\"\s*$", line):
            indent = re.match(r"^(\s*)", line).group(1)
            out.append(f"{indent}ExtraArgs: {extra}")
            changed = True
        elif re.match(r"^(\s*)ExtraArgs:\s+", line) and "--use-sage-attention" not in line:
            out.append(line.rstrip() + f" {extra}")
            changed = True
        else:
            out.append(line)
    if not changed:
        print("Backends.fds: no ExtraArgs field patched")
        return False
    BACKENDS.write_text("\n".join(out) + "\n", encoding="utf-8")
    print(f"patched {BACKENDS} ExtraArgs += {extra}")
    return True


def pip_install_sage(pip: Path) -> bool:
    print(f"pip install triton sageattention via {pip}")
    env = dict(os.environ)
    env["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
    cmd = [str(pip), "install", "-U", "triton", "sageattention"]
    try:
        subprocess.check_call(cmd, env=env)
        print("triton + sageattention installed")
        return True
    except subprocess.CalledProcessError as exc:
        print(f"WARN: pip install failed ({exc}) — ExtraArgs NOT applied")
        return False


def main() -> int:
    gpu = load_gpu()
    plan = tier_notes(gpu)
    prev = {}
    if MARKER.is_file():
        try:
            prev = json.loads(MARKER.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            prev = {}
    same_gpu = bool(prev.get("uuid") and prev.get("uuid") == plan["uuid"])
    if same_gpu and prev.get("extra_args") == plan["extra_args"]:
        if prev.get("pip_ok") or not plan["use_sage"]:
            print(f"perf tune already applied for {plan['name']} ({plan['tier']})")
            return 0

    if same_gpu and plan["use_sage"] and prev.get("pip_ok") is False:
        print("perf tune: retry (previous pip_ok=false — sage/triton ещё не встали)")

    print(f"perf tune: {plan['name']} tier={plan['tier']} sage={plan['use_sage']}")
    pip_ok = not plan["use_sage"]
    restarted_needed = False
    if plan["use_sage"]:
        pip = find_pip()
        if pip:
            pip_ok = pip_install_sage(pip)
        else:
            print("Comfy venv pip not found yet — will retry next up")
            pip_ok = False
        # Only patch ExtraArgs when wheels installed — otherwise Comfy may break.
        if pip_ok and patch_backends_extra_args(plan["extra_args"]):
            restarted_needed = True

    marker = {
        "uuid": plan["uuid"],
        "name": plan["name"],
        "tier": plan["tier"],
        "extra_args": plan["extra_args"] if pip_ok else "",
        "pip_ok": pip_ok,
        "restart_needed": restarted_needed,
    }
    MARKER.write_text(json.dumps(marker, indent=2) + "\n", encoding="utf-8")
    print("wrote", MARKER)
    if restarted_needed:
        print("RESTART_SWARMUI=1")
    return 0


if __name__ == "__main__":
    sys.exit(main())