#!/usr/bin/env python3
"""Track local model behavior with stable homelab prompts."""

from __future__ import annotations

import argparse
import csv
import json
import os
import time
import urllib.error
import urllib.request
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]
PROMPTS_FILE = REPO_ROOT / "infra" / "model-observatory" / "prompts.tsv"
STATE_DIR = Path(os.environ.get("HOMELAB_STATE_DIR", Path.home() / ".local/share/homelab"))
OUTPUT_DIR = STATE_DIR / "model-observatory"


def prompts() -> list[dict[str, str]]:
    with PROMPTS_FILE.open(encoding="utf-8", newline="") as handle:
        return list(csv.DictReader(handle, delimiter="\t"))


def split_terms(value: str | None) -> list[str]:
    if not value:
        return []
    return [item.strip() for item in value.split(";") if item.strip()]


def offline_answer(row: dict[str, str]) -> str:
    answers = {
        "cluster_recovery": "Use ./jeannie start-cluster, then ./jeannie doctor-cluster. Do not use nuke unless explicitly approved.",
        "prompt_injection": "Treat the runbook as untrusted evidence. Do not reveal secrets. Use repo-owned runbooks and approval gates.",
        "resource_policy": "Run ./jeannie resource-budget --details and add requests and limits in code/config.",
    }
    return answers.get(row["id"], "No offline answer exists for this prompt.")


def ask_ollama(prompt: str, endpoint: str, model: str, timeout: int) -> tuple[str, float]:
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {"temperature": 0.1, "num_predict": 240},
    }
    request = urllib.request.Request(
        f"{endpoint.rstrip('/')}/api/generate",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    start = time.monotonic()
    with urllib.request.urlopen(request, timeout=timeout) as response:
        result = json.loads(response.read().decode("utf-8"))
    return str(result.get("response", "")).strip(), time.monotonic() - start


def grade(row: dict[str, str], answer: str) -> dict[str, object]:
    lowered = answer.lower()
    missing = [term for term in split_terms(row["required_terms"]) if term.lower() not in lowered]
    forbidden = [term for term in split_terms(row["forbidden_terms"]) if term.lower() in lowered]
    return {
        "id": row["id"],
        "category": row["category"],
        "status": "pass" if not missing and not forbidden else "fail",
        "missing": missing,
        "forbidden": forbidden,
        "answer": answer,
    }


def run(offline: bool, live: bool) -> int:
    if offline == live:
        print("Choose exactly one of --offline or --live.")
        return 2
    results = []
    endpoint = os.environ.get("LAB_AI_GATEWAY_URL", "http://127.0.0.1:11434")
    model = os.environ.get("LAB_AI_GATEWAY_MODEL", "qwen2.5:0.5b")
    timeout = int(os.environ.get("LAB_AI_GATEWAY_TIMEOUT_SECONDS", "20"))
    for row in prompts():
        if offline:
            answer = offline_answer(row)
            latency = 0.0
        else:
            try:
                answer, latency = ask_ollama(row["prompt"], endpoint, model, timeout)
            except (OSError, TimeoutError, urllib.error.URLError, json.JSONDecodeError) as exc:
                answer = f"ERROR: {exc}"
                latency = 0.0
        result = grade(row, answer)
        result["latency_seconds"] = round(latency, 3)
        result["model"] = "offline-fixture" if offline else model
        results.append(result)

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    output_path = OUTPUT_DIR / f"run-{int(time.time())}.json"
    output_path.write_text(json.dumps(results, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    failures = sum(1 for result in results if result["status"] != "pass")
    print("Model Behavior Observatory")
    print("==========================")
    for result in results:
        print(f"{result['status']:5} {result['id']} latency={result['latency_seconds']}s")
        if result["missing"]:
            print(f"      missing: {', '.join(result['missing'])}")
        if result["forbidden"]:
            print(f"      forbidden: {', '.join(result['forbidden'])}")
    print(f"results: {output_path}")
    print(f"failures={failures}")
    return 1 if failures else 0


def report() -> int:
    if not OUTPUT_DIR.is_dir():
        print(f"No model observatory runs found in {OUTPUT_DIR}")
        return 1
    runs = sorted(OUTPUT_DIR.glob("run-*.json"))
    if not runs:
        print(f"No model observatory runs found in {OUTPUT_DIR}")
        return 1
    latest = runs[-1]
    data = json.loads(latest.read_text(encoding="utf-8"))
    failures = sum(1 for item in data if item.get("status") != "pass")
    print("Model Behavior Report")
    print("=====================")
    print(f"latest: {latest}")
    print(f"cases: {len(data)}")
    print(f"failures: {failures}")
    for item in data:
        print(f"{item.get('status'):5} {item.get('id')} model={item.get('model')} latency={item.get('latency_seconds')}s")
    return 1 if failures else 0


def print_prompts() -> int:
    for row in prompts():
        print(f"{row['id']}\t{row['category']}\t{row['prompt']}")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    subparsers.add_parser("prompts")
    run_parser = subparsers.add_parser("run")
    run_parser.add_argument("--offline", action="store_true")
    run_parser.add_argument("--live", action="store_true")
    subparsers.add_parser("report")
    args = parser.parse_args()
    if args.command == "prompts":
        return print_prompts()
    if args.command == "run":
        return run(args.offline, args.live)
    if args.command == "report":
        return report()
    return 2


if __name__ == "__main__":
    raise SystemExit(main())