#!/usr/bin/env python3
import json, os, sys, datetime

MODELS_FILE = "/opt/data/models_raw.json"
WHITELIST_FILE = "/opt/data/whitelist.json"
REPORT_TXT = "/opt/data/reports/latest_report.txt"
REPORT_MD = "/opt/data/reports/latest_report.md"

CATEGORIES = {
    "coding":    ["code", "coder", "codestral", "qwen2.5-coder", "deepseek-coder", "devstral", "starcoder"],
    "reasoning": ["reasoner", "thinking", "r1", "o1", "o3", "o4", "qwq", "deepseek-r"],
    "vision":    ["vision", "pixtral", "llava", "gpt-4o", "gemini", "claude-3", "qwen2-vl",
                  "llama-3.2-11b-vision", "llama-3.2-90b-vision"],
    "writing":   ["mistral", "mixtral", "command", "llama-3.1-70b", "llama-3.3", "qwen2.5-72b", "gpt-4"],
    "fast":      ["mini", "flash", "haiku", "tiny", "small", "8b", "3b", "1b"],
    "chat":      [],
}

TRUSTED = {"anthropic", "openai", "google", "meta-llama", "deepseek", "mistralai", "qwen", "microsoft", "cohere"}


def parse_price(val):
    try:
        return float(val)
    except Exception:
        return 0.0


def categorize(mid, mname):
    cats = []
    for cat, kws in CATEGORIES.items():
        if cat == "chat":
            continue
        for kw in kws:
            if kw in mid.lower() or kw in mname.lower():
                cats.append(cat)
                break
    return cats if cats else ["chat"]


def is_vision(m):
    mod = m.get("architecture", {}).get("modality", "")
    return "image" in mod


def score(m):
    pricing = m.get("pricing", {})
    p = parse_price(pricing.get("prompt", 0)) * 1_000_000
    c = parse_price(pricing.get("completion", 0)) * 1_000_000
    total_price = p + c + 0.01
    ctx = min(m.get("context_length", 4096) / 100_000, 2.0)
    bonus = 1.2 if m.get("id", "").split("/")[0] in TRUSTED else 1.0
    if total_price <= 0.01:
        return ctx * bonus * 50
    return (ctx / total_price) * bonus * 1000


if not os.path.exists(MODELS_FILE):
    print("FEHLER: models_raw.json fehlt. Erst fetch_models.py ausfuehren.", file=sys.stderr)
    sys.exit(1)

with open(MODELS_FILE) as f:
    data = json.load(f)
models = data["models"]
print(f"[compare] {len(models)} Modelle geladen")

best = {}
for m in models:
    mid = m.get("id", "")
    mname = m.get("name", mid)
    cats = categorize(mid, mname)
    if "vision" in cats and not is_vision(m):
        cats = [c for c in cats if c != "vision"] or ["chat"]
    s = score(m)
    for cat in cats:
        if cat not in best or s > best[cat][1]:
            best[cat] = (m, s)

if "chat" not in best:
    top = sorted(models, key=score, reverse=True)
    if top:
        best["chat"] = (top[0], score(top[0]))

selected = {}
for cat, (m, s) in best.items():
    pricing = m.get("pricing", {})
    pp = parse_price(pricing.get("prompt", 0)) * 1_000_000
    pc = parse_price(pricing.get("completion", 0)) * 1_000_000
    selected[cat] = {
        "id": m["id"],
        "name": m.get("name", m["id"]),
        "score": round(s, 2),
        "context_length": m.get("context_length", 0),
        "price_prompt_per_1M": round(pp, 4),
        "price_completion_per_1M": round(pc, 4),
        "provider": m["id"].split("/")[0]
    }

out = {
    "generated_at": datetime.datetime.utcnow().isoformat() + "Z",
    "model_count": len(selected),
    "selection": selected,
    "slugs": [v["id"] for v in selected.values()]
}
with open(WHITELIST_FILE, "w") as f:
    json.dump(out, f, indent=2)

now = datetime.datetime.utcnow().strftime("%d.%m.%Y %H:%M")
lines_txt = [
    f"AI MODEL SCOUT -- Report {now} UTC",
    "=" * 60,
    f"Analysiert: {len(models)} Modelle | Kategorien besetzt: {len(selected)}",
    "",
    "AUSWAHL (1 bestes Modell pro Kategorie):",
    "-" * 40,
]
lines_md = [
    f"# AI Model Scout -- {now}",
    "",
    f"Analysiert: **{len(models)}** | Kategorien: **{len(selected)}**",
    "",
    "| Kategorie | Modell | Score | Kontext | Preis $/1M |",
    "|-----------|--------|-------|---------|------------|",
]

for cat, info in sorted(selected.items()):
    pt = info["price_prompt_per_1M"] + info["price_completion_per_1M"]
    lines_txt.append(
        f"  [{cat.upper():<10}] {info['id']:<50} Score:{info['score']:<8} "
        f"Ctx:{info['context_length']:>8,} ${pt:.4f}/1M"
    )
    lines_md.append(
        f"| {cat} | `{info['id']}` | {info['score']} | {info['context_length']:,} | ${pt:.4f} |"
    )

lines_txt += ["", f"Whitelist-Slugs ({len(out['slugs'])}):", "  " + str(out["slugs"])]
rtxt = "\n".join(lines_txt)
rmd = "\n".join(lines_md)

with open(REPORT_TXT, "w") as f:
    f.write(rtxt)
with open(REPORT_MD, "w") as f:
    f.write(rmd)

print(rtxt)
print(f"\n[compare] Whitelist -> {WHITELIST_FILE}")
print(f"[compare] Report    -> {REPORT_TXT}")