#!/usr/bin/env python3
"""
llm-models — deterministic local model registry utility.
Subcommands: discover, list, status, test, benchmark, sync, doctor, logs.

Stable model IDs, no filename-based capability guessing.
Every model is tested before its capability flags are set.
"""

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.request
import urllib.error
from pathlib import Path

LLAMA_SWAP_URL = os.environ.get("LLAMA_SWAP_URL", "http://127.0.0.1:9292")
MODEL_DIRS = [
    Path("/models/downloads"),
    Path("/models"),
    Path.home() / "models",
    Path.home() / "downloads",
]
LLAMA_SWAP_CONFIG = Path.home() / ".config" / "llama-swap" / "config.yaml"
REGISTRY_FILE = Path.home() / ".config" / "llama-models" / "registry.json"
TEST_CACHE = Path.home() / ".config" / "llama-models" / "test_cache.json"


def eprint(*a, **kw):
    print(*a, file=sys.stderr, **kw)


# ─── discover ────────────────────────────────────────────────────────────────

def cmd_discover(args):
    """Scan model dirs + llama-swap config + running servers for GGUF files."""
    registry = {"discovered": [], "sources": {}}

    # 1. File system scan
    for d in MODEL_DIRS:
        if d.exists():
            for f in sorted(d.glob("*.gguf")):
                size_gb = f.stat().st_size / (1024**3)
                entry = {
                    "model_id": _infer_id(f),
                    "display_name": f.stem,
                    "path": str(f),
                    "size_gb": round(size_gb, 2),
                    "quantization": _infer_quant(f.stem),
                    "architecture": _infer_arch(f.stem),
                    "source": "filesystem",
                    "in_config": False,
                    "in_swap": False,
                    "gguf_header": _read_gguf_header(f),
                }
                registry["discovered"].append(entry)

    # 2. Check llama-swap config
    if LLAMA_SWAP_CONFIG.exists():
        registry["sources"]["llama_swap_config"] = str(LLAMA_SWAP_CONFIG)
        with open(LLAMA_SWAP_CONFIG) as fh:
            text = fh.read()
        # Extract model names from config
        for m in re.finditer(r'^\s{2}(\S+):\s*\n\s+cmd:', text, re.MULTILINE):
            mid = m.group(1)
            for e in registry["discovered"]:
                if mid == e["model_id"] or mid in e["model_id"]:
                    e["in_config"] = True
                    break
            else:
                # In config but not on disk? Mark anyway
                registry.setdefault("config_only", []).append(mid)

    # 3. Check running llama-swap
    try:
        resp = urllib.request.urlopen(f"{LLAMA_SWAP_URL}/v1/models", timeout=5)
        data = json.loads(resp.read())
        for m in data.get("data", []):
            mid = m["id"]
            for e in registry["discovered"]:
                if mid == e["model_id"] or mid in e["model_id"] or e["display_name"] == mid:
                    e["in_swap"] = True
                    break
    except Exception:
        pass

    _save_registry(registry)

    print(json.dumps(registry, indent=2))
    print(f"\n{'='*60}")
    print(f"Discovered {len(registry['discovered'])} models, saved to {REGISTRY_FILE}")


# ─── list ────────────────────────────────────────────────────────────────────

def cmd_list(args):
    """List all discovered models with key metadata."""
    reg = _load_registry()
    if not reg.get("discovered"):
        print("No models discovered yet. Run 'llm-models discover' first.")
        return

    print(f"{'ID':<22} {'Quant':<8} {'Size':>6} {'Config':<6} {'Swap':<6} Arch")
    print("-" * 80)
    for e in reg["discovered"]:
        print(
            f"{e['model_id']:<22} "
            f"{e['quantization']:<8} "
            f"{e['size_gb']:>5.1f}G "
            f"{'YES' if e.get('in_config') else 'no':<6} "
            f"{'YES' if e.get('in_swap') else 'no':<6} "
            f"{e.get('architecture','?')}"
        )


# ─── status ──────────────────────────────────────────────────────────────────

def cmd_status(args):
    """Show current runtime status of all known models."""
    try:
        resp = urllib.request.urlopen(f"{LLAMA_SWAP_URL}/v1/models", timeout=5)
        data = json.loads(resp.read())
        print(f"llama-swap: {LLAMA_SWAP_URL} — {len(data.get('data',[]))} model(s) registered\n")
        for m in data.get("data", []):
            print(f"  {m['id']}")
            for k, v in m.items():
                if k != "id":
                    print(f"    {k}: {v}")
    except Exception as e:
        print(f"Cannot reach llama-swap at {LLAMA_SWAP_URL}: {e}")

    # Show system info
    print("\n--- System ---")
    try:
        nvs = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,name,memory.total,memory.used",
             "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10
        )
        for line in nvs.stdout.strip().split("\n"):
            print(f"  GPU: {line}")
    except Exception:
        print("  nvidia-smi unavailable")


# ─── test ────────────────────────────────────────────────────────────────────

_TEST_PROMPTS = {
    "basic": "Reply with exactly: OK",
    "streaming": "Count from 1 to 5, one per line.",
    "long_input": "Repeat after me: " + "hello world " * 2000,
    "structured_json": 'Reply with JSON: {"status":"ok","value":42}',
}


def cmd_test(args):
    """Run validation tests against a specific model or all."""
    reg = _load_registry()
    models = [e for e in reg.get("discovered", []) if e.get("in_swap")]
    if args.model_id:
        models = [m for m in models if args.model_id in m["model_id"]]

    if not models:
        print("No models to test. Run discover first, or check llama-swap.")
        return

    results = {"tested": [], "failed": [], "capabilities": {}}

    for model in models:
        mid = model["model_id"]
        eprint(f"\n{'='*60}")
        eprint(f"Testing: {mid} ({model['display_name']})")
        eprint(f"{'='*60}")

        caps = {
            "chat_completion": False,
            "streaming": False,
            "long_input": False,
            "structured_json": False,
            "tool_call": False,
            "responses_api": False,
        }

        # Chat completion
        try:
            r = _llm_request(mid, _TEST_PROMPTS["basic"])
            caps["chat_completion"] = "OK" in r
            eprint(f"  basic chat: {'PASS' if caps['chat_completion'] else 'FAIL'}")
        except Exception as e:
            eprint(f"  basic chat: FAIL ({e})")

        # Streaming
        try:
            r = _llm_request(mid, _TEST_PROMPTS["streaming"], stream=True)
            caps["streaming"] = r
            eprint(f"  streaming:  {'PASS' if caps['streaming'] else 'FAIL'}")
        except Exception as e:
            eprint(f"  streaming:  FAIL ({e})")

        # Long input
        try:
            r = _llm_request(mid, _TEST_PROMPTS["long_input"], timeout=120)
            caps["long_input"] = "hello world" in r[:100]
            eprint(f"  long input: {'PASS' if caps['long_input'] else 'FAIL'}")
        except Exception as e:
            eprint(f"  long input: FAIL ({e})")

        # Structured JSON (response_format)
        try:
            r = _llm_request(mid, "Return JSON: {\"test\": true}", response_format="json")
            caps["structured_json"] = '"test"' in r
            eprint(f"  json mode:  {'PASS' if caps['structured_json'] else 'FAIL'}")
        except Exception as e:
            eprint(f"  json mode:  FAIL ({e})")

        # Tool call test
        try:
            r = _llm_request(mid, "What's 2+2?", tools=_get_test_tools())
            caps["tool_call"] = r
            eprint(f"  tool call:  {'PASS' if caps['tool_call'] else 'FAIL'}")
        except Exception as e:
            eprint(f"  tool call:  FAIL ({e})")

        # Responses API
        try:
            r = _llm_request_responses(mid)
            caps["responses_api"] = r
            eprint(f"  responses:  {'PASS' if caps['responses_api'] else 'FAIL'}")
        except Exception as e:
            eprint(f"  responses:  FAIL ({e})")

        # Classify
        classification = _classify(caps)
        eprint(f"  => Classification: {classification}")

        model["capabilities"] = caps
        model["classification"] = classification
        results["capabilities"][mid] = caps

        if all(caps.values()) or sum(1 for v in caps.values() if v) >= 3:
            results["tested"].append(mid)
        else:
            results["failed"].append(mid)

    _save_test_cache(results)
    print(json.dumps(results, indent=2))
    eprint(f"\nTested: {len(results['tested'])}, Failed: {len(results['failed'])}")


def _get_test_tools():
    return [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get weather for a city",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "city": {"type": "string"}
                    },
                    "required": ["city"]
                }
            }
        }
    ]


def _classify(caps):
    """Classify model based on tested capabilities."""
    score = sum(1 for v in caps.values() if v)
    if score >= 5 and caps.get("tool_call") and caps.get("structured_json"):
        return "agent-full"
    if score >= 4 and caps.get("chat_completion"):
        return "coding-edit"
    if score >= 3:
        return "analysis-only"
    if caps.get("chat_completion"):
        return "fast-helper"
    return "unsupported"


def _llm_request(model_id, prompt, stream=False, response_format=None, tools=None, timeout=60):
    """Make a chat completion request to llama-swap."""
    body = {
        "model": model_id,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 256,
        "temperature": 0.01,
    }
    if stream:
        body["stream"] = True
    if response_format == "json":
        body["response_format"] = {"type": "json_object"}
    if tools:
        body["tools"] = tools
        body["tool_choice"] = "auto"

    data = json.dumps(body).encode()
    req = urllib.request.Request(
        f"{LLAMA_SWAP_URL}/v1/chat/completions",
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        resp = urllib.request.urlopen(req, timeout=timeout)
    except urllib.error.HTTPError as e:
        if stream:
            return False
        raise

    if stream:
        # Just check we can get first chunk
        for line in resp:
            if line.strip():
                return True
        return False

    result = json.loads(resp.read())

    if tools:
        return "tool_calls" in (result.get("choices", [{}])[0].get("message", {}))

    content = result.get("choices", [{}])[0].get("message", {}).get("content", "")
    return content or ""


def _llm_request_responses(model_id, timeout=30):
    """Test the Responses API endpoint."""
    body = {
        "model": model_id,
        "input": "Reply with exactly: RESPONSES_OK",
        "max_output_tokens": 64,
    }
    data = json.dumps(body).encode()
    req = urllib.request.Request(
        f"{LLAMA_SWAP_URL}/v1/responses",
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        resp = urllib.request.urlopen(req, timeout=timeout)
        result = json.loads(resp.read())
        return "RESPONSES_OK" in json.dumps(result)
    except (urllib.error.HTTPError, Exception):
        return False


# ─── benchmark ───────────────────────────────────────────────────────────────

def cmd_benchmark(args):
    """Quick performance benchmark for a model."""
    reg = _load_registry()
    mid = args.model_id or (reg.get("discovered", [{}])[0].get("model_id") if reg.get("discovered") else None)
    if not mid:
        print("Specify a model ID or run discover first.")
        return

    prompt = "What is the capital of France? Answer concisely.\n" * 10
    eprint(f"Benchmarking {mid}...")

    # Warmup
    try:
        _llm_request(mid, prompt[:100])
    except Exception:
        pass

    # Timed run
    times = []
    for _ in range(3):
        t0 = time.time()
        try:
            r = _llm_request(mid, prompt)
            elapsed = time.time() - t0
            tokens = len(r.split()) + len(prompt.split())
            tps = tokens / elapsed if elapsed > 0 else 0
            times.append({"elapsed": round(elapsed, 2), "tokens_est": tokens, "tps": round(tps, 1)})
            eprint(f"  Run {len(times)}: {elapsed:.1f}s, ~{tps:.0f} t/s")
        except Exception as e:
            eprint(f"  Run {len(times)+1}: FAIL ({e})")

    avg_tps = sum(t["tps"] for t in times) / len(times) if times else 0
    result = {"model": mid, "runs": times, "avg_tps": round(avg_tps, 1)}
    print(json.dumps(result, indent=2))
    eprint(f"\nAverage: {avg_tps:.0f} tokens/sec")


# ─── sync ────────────────────────────────────────────────────────────────────

def cmd_sync(args):
    """Sync discovered models into llama-swap config."""
    reg = _load_registry()
    if not LLAMA_SWAP_CONFIG.exists():
        eprint(f"Cannot find {LLAMA_SWAP_CONFIG}")
        return

    with open(LLAMA_SWAP_CONFIG) as fh:
        config_text = fh.read()

    changes = 0
    for e in reg.get("discovered", []):
        if e.get("in_config"):
            continue
        # Add a basic model entry
        model_block = (
            f"\n  # Auto-discovered: {e['display_name']}\n"
            f"  {e['model_id']}:\n"
            f"    cmd: >\n"
            f"      /usr/local/bin/llama-server\n"
            f"      -m {e['path']}\n"
            f"      --port ${{PORT}} --host 127.0.0.1\n"
            f"      -ngl 99 --no-mmap --mlock\n"
            f"      --flash-attn on\n"
            f"      -c 32768 --ctx-size 32768\n"
            f"      --batch-size 1024 --ubatch-size 256\n"
            f"      --chat-template-kwargs '{{\"enable_thinking\":false}}'\n"
            f"      --jinja\n"
            f"      --temp 0.2 --top-p 0.95\n"
            f"      --metrics\n"
            f"    aliases: []\n"
            f"    ttl: 1800\n"
            f"    concurrency: 1\n"
        )
        config_text += model_block
        changes += 1

    if changes:
        backup = LLAMA_SWAP_CONFIG.with_suffix(".yaml.bak")
        shutil.copy2(LLAMA_SWAP_CONFIG, backup)
        with open(LLAMA_SWAP_CONFIG, "w") as fh:
            fh.write(config_text)
        eprint(f"Added {changes} models to {LLAMA_SWAP_CONFIG} (backup at {backup})")
        eprint("Restart llama-swap: sudo systemctl restart llama-swap")
    else:
        eprint("All discovered models already in config.")


# ─── doctor ──────────────────────────────────────────────────────────────────

def cmd_doctor(args):
    """Diagnose common issues."""
    checks = []

    # Check llama-swap
    r = subprocess.run(["systemctl", "is-active", "llama-swap"], capture_output=True, text=True, timeout=5)
    checks.append(("llama-swap service", r.stdout.strip() == "active" if r.stdout.strip() else "inactive"))

    # Check binary
    r = subprocess.run(["llama-swap", "--version"], capture_output=True, text=True, timeout=5)
    checks.append(("llama-swap binary", r.stdout.strip() or "missing"))

    # Check ports
    try:
        resp = urllib.request.urlopen(f"{LLAMA_SWAP_URL}/v1/models", timeout=5)
        checks.append(("API reachable", f"yes ({resp.status})"))
    except Exception as e:
        checks.append(("API reachable", f"no ({e})"))

    # Check GPU
    r = subprocess.run(["nvidia-smi"], capture_output=True, text=True, timeout=10)
    driver_ok = "NVIDIA-SMI" in r.stdout
    checks.append(("NVIDIA driver", "loaded" if driver_ok else "missing"))

    # Check models dirs
    for d in MODEL_DIRS:
        if d.exists():
            gguvs = list(d.glob("*.gguf"))
            checks.append((f"Models in {d}", f"{len(gguvs)} GGUF files"))
            break
    else:
        checks.append(("Model directories", "none found"))

    # Ulimits
    r = subprocess.run(["ulimit", "-l"], capture_output=True, text=True, timeout=5)
    checks.append(("memlock limit", r.stdout.strip()))

    print("llm-models doctor report\n")
    for name, status in checks:
        icon = "✓" if "no" not in status and "missing" not in status and "inactive" not in status else "✗"
        print(f"  {icon} {name}: {status}")
    print()


# ─── logs ────────────────────────────────────────────────────────────────────

def cmd_logs(args):
    """Show recent llama-swap logs."""
    lines = args.lines or 50
    r = subprocess.run(
        ["journalctl", "-u", "llama-swap", "--no-pager", "-n", str(lines)],
        capture_output=True, text=True, timeout=10
    )
    print(r.stdout if r.stdout else "(no logs)")
    if r.stderr:
        eprint(r.stderr)


# ─── helpers ─────────────────────────────────────────────────────────────────

def _infer_id(path):
    """Generate a stable model ID from filename."""
    stem = path.stem
    # Remove common patterns
    stem = re.sub(r'-[Qq][0-9]_[A-Za-z0-9_.]+', '', stem)
    stem = re.sub(r'-[Uu][Dd]', '', stem)
    stem = re.sub(r'-MTP', '', stem)
    stem = re.sub(r'[-_]it$', '', stem)
    return stem.lower().replace(" ", "-")


def _infer_quant(stem):
    """Extract quantization from filename."""
    m = re.search(r'(Q[0-9]_[A-Za-z0-9]+|[IQ][0-9]_[0-9])', stem)
    return m.group(1) if m else "unknown"


def _infer_arch(stem):
    """Guess architecture from filename."""
    low = stem.lower()
    if "qwen" in low: return "Qwen"
    if "gemma" in low: return "Gemma"
    if "bonsai" in low: return "Bonsai (ternary)"
    if "nex" in low or "n2" in low: return "Nex (MoE)"
    if "llama" in low: return "Llama"
    if "deepseek" in low: return "DeepSeek"
    if "mistral" in low: return "Mistral"
    if "qwopus" in low: return "Qwen (Qwopus)"
    return "unknown"


def _read_gguf_header(path):
    """Quick GGUF header read — limited to small reads only."""
    try:
        file_size = path.stat().st_size
        if file_size > 100 * 1024 * 1024:  # Skip detailed header for files > 100MB
            return {"valid": True, "size_gb": round(file_size / (1024**3), 2), "note": "large file, header skip"}

        with open(path, "rb") as f:
            magic = f.read(4)
            if magic != b"GGUF":
                return {"valid": False, "error": "not a GGUF file"}
            return {"valid": True, "magic": "GGUF"}
    except Exception as e:
        return {"valid": False, "error": str(e)}


def _load_registry():
    if REGISTRY_FILE.exists():
        try:
            with open(REGISTRY_FILE) as f:
                return json.load(f)
        except Exception:
            pass
    return {"discovered": [], "sources": {}}


def _save_registry(reg):
    REGISTRY_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(REGISTRY_FILE, "w") as f:
        json.dump(reg, f, indent=2, default=str)


def _save_test_cache(results):
    TEST_CACHE.parent.mkdir(parents=True, exist_ok=True)
    with open(TEST_CACHE, "w") as f:
        json.dump(results, f, indent=2)


# ─── main ────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Local model registry and testing utility")
    sub = parser.add_subparsers(dest="command", required=True)

    p = sub.add_parser("discover", help="Scan filesystem + config for models")
    p = sub.add_parser("list", help="List discovered models")
    p = sub.add_parser("status", help="Show runtime status")

    p = sub.add_parser("test", help="Run validation tests against models")
    p.add_argument("model_id", nargs="?", help="Specific model ID to test (default: all)")

    p = sub.add_parser("benchmark", help="Benchmark a model")
    p.add_argument("model_id", nargs="?", help="Model ID to benchmark")

    p = sub.add_parser("sync", help="Add undiscovered models to llama-swap config")
    p = sub.add_parser("doctor", help="Diagnose common issues")

    p = sub.add_parser("logs", help="Show llama-swap logs")
    p.add_argument("-n", "--lines", type=int, default=50, help="Number of log lines")

    args = parser.parse_args()

    command_map = {
        "discover": cmd_discover,
        "list": cmd_list,
        "status": cmd_status,
        "test": cmd_test,
        "benchmark": cmd_benchmark,
        "sync": cmd_sync,
        "doctor": cmd_doctor,
        "logs": cmd_logs,
    }

    fn = command_map.get(args.command)
    if fn:
        fn(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()