#!/usr/bin/env python3
"""Together.ai-Serverless-Chat-Modelle per Checkbox-Dialog in die OpenCode-Config übernehmen.

Vorauswahl = bereits konfigurierte Modelle (providers.togetherai.models); bei leerer
Config die im OpenCode-Katalog verfügbaren. Esc/Abbruch oder leere Auswahl ändert
nichts; bestätigte Auswahl ersetzt als Voll-Sync die models-Map (Backup vorher).
"""
import json
import os
import shutil
import subprocess
import sys
import urllib.request
from datetime import datetime

CONFIG = os.path.expanduser("~/.config/opencode/opencode.jsonc")
AUTH = os.path.expanduser("~/.local/share/opencode/auth.json")
# urllib-Default-UA wird von Together (Cloudflare) gebannt -> eigener UA nötig
HEADERS = {"User-Agent": "text-tool/1.0"}


def api_key():
    if os.environ.get("TOGETHER_API_KEY"):
        return os.environ["TOGETHER_API_KEY"]
    try:
        return json.load(open(AUTH))["togetherai"]["key"]
    except Exception:
        sys.exit("Kein API-Key: TOGETHER_API_KEY setzen oder %s vorhanden." % AUTH)


def fetch(url, key):
    req = urllib.request.Request(url, headers={**HEADERS, "Authorization": "Bearer " + key})
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read())


def serverless_chat(key):
    ms = fetch("https://api.together.xyz/v1/models?serverless=true", key)
    if isinstance(ms, dict):
        ms = ms.get("data", [])
    chat = [m for m in ms if m.get("type") == "chat"]
    chat.sort(key=lambda m: (-m.get("created", 0), m["id"]))
    return chat


def configured_models():
    """Bereits konfigurierte togetherai-Modelle (die Vorauswahl)."""
    try:
        m = json.load(open(CONFIG))["providers"]["togetherai"]["models"]
        return set(m)
    except Exception:
        return set()


def opencode_models():
    """Erstlauf-Fallback: Modell-IDs, die OpenCode aus dem Katalog anbietet."""
    try:
        out = subprocess.run(["opencode", "api", "get", "/api/model"],
                             capture_output=True, text=True, timeout=30, check=True).stdout
        data = json.loads(out).get("data", [])
        return {m["id"] for m in data
                if m.get("providerID") == "togetherai" and m.get("enabled") is not False}
    except Exception as e:
        print("warn: OpenCode-Modelle nicht ermittelbar (%s) — keine Vorauswahl" % e,
              file=sys.stderr)
        return set()


def entry(m):
    e = {"name": m.get("display_name") or m["id"]}
    if m.get("context_length"):
        e["limit"] = {"context": m["context_length"]}
    p = m.get("pricing") or {}
    c = {"input": round(p.get("input", 0), 6), "output": round(p.get("output", 0), 6)}
    if p.get("cached_input"):
        c["cache_read"] = round(p["cached_input"], 6)
    e["cost"] = c
    return e


def zenity_pick(chat, avail):
    rows = []
    for m in chat:
        p = m.get("pricing") or {}
        rows += ["TRUE" if m["id"] in avail else "FALSE", m["id"],
                 str(m.get("context_length") or "?"),
                 "%g / %g" % (p.get("input", 0), p.get("output", 0))]
    cmd = ["zenity", "--list", "--checklist",
           "--title=Together.ai → OpenCode",
           "--text=Zu übernehmende Modelle anhaken:",
           "--column=Anwahl", "--column=Modell", "--column=Kontext",
           "--column=$/1M in→out", "--print-column=2",
           "--width=980", "--height=720"] + rows
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        return None  # Esc/Abbruch
    return [x for x in r.stdout.strip().split("|") if x]


def terminal_pick(chat, avail):
    for i, m in enumerate(chat, 1):
        p = m.get("pricing") or {}
        print("%2d %-45s ctx %-8s %g / %g %s" % (
            i, m["id"], m.get("context_length") or "?",
            p.get("input", 0), p.get("output", 0), "*" if m["id"] in avail else ""))
    raw = input("Nummern (z. B. 1,3,7; a=alle; leer=Abbruch): ").strip()
    if not raw:
        return None
    if raw == "a":
        return [m["id"] for m in chat]
    try:
        return [chat[int(n) - 1]["id"] for n in raw.replace(" ", "").split(",")]
    except (ValueError, IndexError):
        return None


def merge(selected_ids, chat):
    with open(CONFIG) as f:
        text = f.read()
    try:
        cfg = json.loads(text)
    except ValueError as e:
        sys.exit("Config ist kein reines JSON (%s) — bitte manuell einfügen." % e)
    old_keys = set(cfg)
    entry_by_id = {m["id"]: entry(m) for m in chat}
    models = {i: entry_by_id[i] for i in sorted(selected_ids)}
    prov = cfg.setdefault("providers", {}).setdefault("togetherai", {})
    if prov.get("models") == models:
        print("Config unverändert.", file=sys.stderr)
        return
    bak = CONFIG + ".bak-" + datetime.now().strftime("%Y%m%d-%H%M%S")
    shutil.copy2(CONFIG, bak)
    prov["models"] = models
    tmp = CONFIG + ".tmp"
    with open(tmp, "w") as f:
        f.write(json.dumps(cfg, indent=2, ensure_ascii=False) + "\n")
    os.replace(tmp, CONFIG)
    # Self-Check: Config parst, alte Keys unangetastet, Auswahl vollständig
    new = json.load(open(CONFIG))
    assert old_keys <= set(new), "Top-Level-Keys verloren!"
    assert new["providers"]["togetherai"]["models"] == models
    print("%d Modelle übernommen, Backup: %s" % (len(models), bak), file=sys.stderr)


def main():
    chat = serverless_chat(api_key())
    assert chat, "Keine Chat-Modelle in der serverless-Liste?"
    ids = {m["id"] for m in chat}
    cfg_ids = configured_models()
    if cfg_ids:
        pre = cfg_ids & ids
        src = "letzte Auswahl"
    else:
        pre = opencode_models() & ids
        src = "OpenCode-Katalog (Erstlauf)"
    print("%d Serverless-Chat-Modelle, %d vorausgewählt (%s)"
          % (len(chat), len(pre), src), file=sys.stderr)
    pick = (zenity_pick if shutil.which("zenity") else terminal_pick)(chat, pre)
    if not pick:
        sys.exit("Abgebrochen — nichts geändert.")
    assert all(i in ids for i in pick), "Unbekannte Modell-ID in Auswahl!"
    merge(pick, chat)


if __name__ == "__main__":
    main()