#!/usr/bin/env python3
"""Baut nodes.json aus den öffentlichen Replicate-Schema-Seiten (kein API-Token nötig).

Aufruf:  python scripts/build_catalog.py
Schreibt: nodes.json im Repo-Root.

Neue Modelle: unten in META eintragen, Script laufen lassen, fertig.
"""
import json
import re
import time
import urllib.request
from pathlib import Path

# Modell-ID → (Anzeigename, Kategorie, Output-Art: image|video|audio)
META = {
    "bytedance/seedream-4.5":      ("Seedream 4.5",         "image",   "image"),
    "google/nano-banana-pro":      ("Nano Banana Pro",      "image",   "image"),
    "qwen/qwen-image-edit-plus":   ("Qwen Image Edit Plus", "image",   "image"),
    "prunaai/z-image-turbo":       ("Z-Image Turbo",        "image",   "image"),
    "prunaai/z-image-turbo-img2img": ("Z-Image Turbo I2I",  "image",   "image"),
    "kwaivgi/kling-v3-video":      ("Kling v3 Video",       "video",   "video"),
    "lightricks/ltx-2-distilled":  ("LTX-2 Distilled",      "video",   "video"),
    "wan-video/wan-2.2-t2v-fast":  ("Wan 2.2 T2V Fast",     "video",   "video"),
    "wan-video/wan-2.2-i2v-fast":  ("Wan 2.2 I2V Fast",     "video",   "video"),
    "google/veo-3.1":              ("Veo 3.1",              "video",   "video"),
    "bytedance/seedance-2.0-fast": ("Seedance 2.0 Fast",    "video",   "video"),
    "cjwbw/supir":                 ("SUPIR Upscale",        "upscale", "image"),
    "cjwbw/real-esrgan":           ("Real-ESRGAN Upscale",  "upscale", "image"),
    "minimax/speech-2.8-hd":       ("MiniMax Speech 2.8 HD","tts",     "audio"),
    "inworld/realtime-tts-2":      ("Inworld TTS 2",        "tts",     "audio"),
    "elevenlabs/v3":               ("ElevenLabs v3",        "tts",     "audio"),
    "veed/fabric-1.0":             ("VEED Fabric 1.0",      "lipsync", "video"),
    "prunaai/p-video-avatar":      ("Pruna Video Avatar",   "lipsync", "video"),
    "sync/lipsync-2-pro":          ("Sync Lipsync 2 Pro",   "lipsync", "video"),
    "heygen/lipsync-precision":    ("HeyGen Lipsync Precision", "lipsync", "video"),
}


def fetch_schema(model_id: str) -> dict | None:
    """Schema aus dem eingebetteten JSON der öffentlichen Schema-Seite ziehen."""
    req = urllib.request.Request(
        f"https://replicate.com/{model_id}/api/schema",
        headers={"User-Agent": "Mozilla/5.0"},
    )
    html = urllib.request.urlopen(req, timeout=30).read().decode()
    for m in re.finditer(r'<script[^>]*type="application/json"[^>]*>(.*?)</script>', html, re.DOTALL):
        try:
            data = json.loads(m.group(1))
        except Exception:
            continue
        version = data.get("version") or data
        schema = version.get("_extras", {}).get("dereferenced_openapi_schema") or {}
        schemas = schema.get("components", {}).get("schemas", {})
        if schemas.get("Input"):
            return {"input": schemas["Input"], "version_id": version.get("id")}
    return None


def classify_uri(name: str, is_list: bool) -> dict:
    """URI-Inputs nach Medientyp einsortieren."""
    lname = name.lower()
    if any(k in lname for k in ("audio", "voice_sample", "speech")):
        return {"kind": "audio"}
    if any(k in lname for k in ("video", "motion")):
        return {"kind": "uri_string"}  # Video-Inputs: URL als String (zu groß für Upload-Tricks)
    if any(k in lname for k in ("image", "img", "frame", "mask", "reference")):
        return {"kind": "image", "is_list": is_list}
    return {"kind": "uri_string"}


def distill(model_id: str, input_schema: dict) -> dict:
    props = input_schema.get("properties", {})
    required = set(input_schema.get("required", []))
    inputs = []
    for name, p in props.items():
        if "allOf" in p and p["allOf"]:
            merged = dict(p["allOf"][0])
            merged.update({k: v for k, v in p.items() if k != "allOf"})
            p = merged
        entry = {
            "name": name,
            "order": p.get("x-order", 999),
            "required": name in required,
            "description": (p.get("description") or "")[:200],
        }
        t = p.get("type")
        fmt = p.get("format", "")
        items = p.get("items", {}) if isinstance(p.get("items"), dict) else {}
        is_uri = fmt == "uri" or items.get("format") == "uri"
        is_list = t == "array"
        lname = name.lower()

        if p.get("enum"):
            entry["kind"] = "combo"
            entry["options"] = p["enum"]
            entry["default"] = p.get("default", p["enum"][0])
        elif t == "boolean":
            entry["kind"] = "bool"
            entry["default"] = bool(p.get("default", False))
        elif t == "integer":
            entry["kind"] = "int"
            entry["default"] = p.get("default", 0)
            entry["min"] = p.get("minimum", -2**31)
            entry["max"] = p.get("maximum", 2**31 - 1)
        elif t == "number":
            entry["kind"] = "float"
            entry["default"] = p.get("default", 0.0)
            entry["min"] = p.get("minimum", -1e9)
            entry["max"] = p.get("maximum", 1e9)
        elif is_uri:
            entry.update(classify_uri(name, is_list))
        else:
            entry["kind"] = "string"
            entry["default"] = p.get("default") or ""
            if any(k in lname for k in ("prompt", "text", "instruction", "script")):
                entry["multiline"] = True
        inputs.append(entry)
    inputs.sort(key=lambda e: (not e["required"], e["order"]))
    name, cat, out = META[model_id]
    return {"id": model_id, "name": name, "category": cat, "output": out, "inputs": inputs}


def main():
    models = []
    for model_id in META:
        try:
            fetched = fetch_schema(model_id)
        except Exception as exc:
            print(f"✗ {model_id}: {exc}")
            continue
        if not fetched:
            print(f"✗ {model_id}: kein Schema gefunden")
            continue
        entry = distill(model_id, fetched["input"])
        models.append(entry)
        print(f"✓ {model_id}: {len(entry['inputs'])} inputs")
        time.sleep(1)

    out_path = Path(__file__).resolve().parent.parent / "nodes.json"
    out_path.write_text(json.dumps({"models": models}, indent=1, ensure_ascii=False))
    print(f"→ {out_path} ({len(models)}/{len(META)} Modelle)")


if __name__ == "__main__":
    main()