"""ComfyUI-Replicate-Select — kuratierte Replicate-Modelle als ComfyUI-Nodes.

Architektur nach Runware-Vorbild: nodes.json (Katalog) + Generator + Custom-Node.
Auth: REPLICATE_API_TOKEN aus der Umgebung (Modal-Secret mcp-keys).
Katalog aktualisieren: python scripts/build_catalog.py
"""
import base64
import io
import json
import os
import time
from pathlib import Path

import numpy as np
import torch

CATALOG_PATH = Path(__file__).parent / "nodes.json"
CATEGORY_ROOT = "Replicate Select"
DATA_URI_LIMIT = 200 * 1024  # größere Dateien gehen über die Files-API


# ---------- Replicate-Client ----------

def _client():
    import replicate
    if not os.environ.get("REPLICATE_API_TOKEN"):
        raise RuntimeError(
            "REPLICATE_API_TOKEN fehlt. In Modal: Secret 'mcp-keys' prüfen; "
            "lokal: export REPLICATE_API_TOKEN=r8_..."
        )
    return replicate


def _upload_bytes(data: bytes, filename: str, content_type: str) -> str:
    """Datei zu Replicate hochladen (Files-API) oder als data-URI inline."""
    if len(data) <= DATA_URI_LIMIT:
        b64 = base64.b64encode(data).decode()
        return f"data:{content_type};base64,{b64}"
    import replicate
    f = replicate.files.create(io.BytesIO(data), {"filename": filename, "content_type": content_type})
    return f.urls["get"]


# ---------- Medien-Konvertierung ----------

def _image_to_uri(image_tensor) -> str:
    """ComfyUI IMAGE (B,H,W,C float 0-1) → PNG-URI (erstes Bild im Batch)."""
    from PIL import Image
    arr = (image_tensor[0].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
    img = Image.fromarray(arr)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return _upload_bytes(buf.getvalue(), "input.png", "image/png")


def _audio_to_uri(audio: dict) -> str:
    """ComfyUI AUDIO ({'waveform': (B,C,T), 'sample_rate': int}) → WAV-URI."""
    import torchaudio
    buf = io.BytesIO()
    waveform = audio["waveform"][0]  # (C, T)
    torchaudio.save(buf, waveform.cpu(), audio["sample_rate"], format="wav")
    return _upload_bytes(buf.getvalue(), "input.wav", "audio/wav")


def _download(url_or_file) -> bytes:
    """Replicate-Output (FileOutput oder URL-String) → Bytes."""
    if hasattr(url_or_file, "read"):
        return url_or_file.read()
    import urllib.request
    with urllib.request.urlopen(str(url_or_file), timeout=600) as resp:
        return resp.read()


def _to_image_tensor(data: bytes):
    from PIL import Image
    img = Image.open(io.BytesIO(data)).convert("RGB")
    arr = np.asarray(img).astype(np.float32) / 255.0
    return torch.from_numpy(arr)[None,]


def _to_audio(data: bytes) -> dict:
    import torchaudio
    waveform, sr = torchaudio.load(io.BytesIO(data))
    return {"waveform": waveform[None,], "sample_rate": sr}


def _save_video(data: bytes, model_id: str):
    """Video in den ComfyUI-Output-Ordner schreiben → (Pfad, ui-Metadaten für /history)."""
    try:
        import folder_paths
        out_dir = Path(folder_paths.get_output_directory())
    except Exception:
        out_dir = Path("output")
    subfolder = "video"
    (out_dir / subfolder).mkdir(parents=True, exist_ok=True)
    slug = model_id.replace("/", "_")
    filename = f"replicate_{slug}_{int(time.time())}.mp4"
    (out_dir / subfolder / filename).write_bytes(data)
    meta = {"filename": filename, "subfolder": subfolder, "type": "output"}
    return str(out_dir / subfolder / filename), meta


def _first_output(result):
    """Replicate gibt je nach Modell FileOutput, Liste oder dict zurück."""
    if isinstance(result, (list, tuple)):
        return result[0] if result else None
    if isinstance(result, dict):
        for key in ("video", "audio", "image", "output", "url"):
            if key in result:
                return result[key]
        return next(iter(result.values()), None)
    return result


# ---------- Node-Generierung ----------

def _comfy_input(entry: dict):
    kind = entry["kind"]
    tooltip = {"tooltip": entry["description"]} if entry.get("description") else {}
    if kind == "combo":
        return ([str(o) for o in entry["options"]], {"default": str(entry.get("default", "")), **tooltip})
    if kind == "bool":
        return ("BOOLEAN", {"default": entry.get("default", False), **tooltip})
    if kind == "int":
        return ("INT", {"default": int(entry.get("default") or 0),
                        "min": int(max(entry.get("min", -2**31), -2**31)),
                        "max": int(min(entry.get("max", 2**31 - 1), 2**31 - 1)), **tooltip})
    if kind == "float":
        return ("FLOAT", {"default": float(entry.get("default") or 0.0),
                          "min": float(entry.get("min", -1e9)), "max": float(entry.get("max", 1e9)),
                          "step": 0.01, **tooltip})
    if kind == "image":
        return ("IMAGE", tooltip)
    if kind == "audio":
        return ("AUDIO", tooltip)
    # string / uri_string
    return ("STRING", {"default": str(entry.get("default") or ""),
                       "multiline": entry.get("multiline", False), **tooltip})


def _build_arguments(model_def: dict, kwargs: dict) -> dict:
    args = {}
    for entry in model_def["inputs"]:
        name = entry["name"]
        value = kwargs.get(name)
        if value is None:
            continue
        kind = entry["kind"]
        if kind == "image":
            if entry.get("is_list"):
                # Batch-Bilder einzeln hochladen (z.B. Qwen Edit Plus: Source + Referenzen)
                args[name] = [_image_to_uri(value[i:i + 1]) for i in range(value.shape[0])]
            else:
                args[name] = _image_to_uri(value)
        elif kind == "audio":
            args[name] = _audio_to_uri(value)
        elif kind == "combo":
            args[name] = value
        elif kind in ("string", "uri_string"):
            if str(value).strip():
                args[name] = value
        else:
            args[name] = value
    return args


_RETURN_BY_OUTPUT = {
    "image": (("IMAGE",), ("image",)),
    "audio": (("AUDIO",), ("audio",)),
    "video": (("STRING",), ("video_path",)),
}


def _make_node(model_def: dict):
    output_kind = model_def["output"]
    return_types, return_names = _RETURN_BY_OUTPUT[output_kind]

    required, optional = {}, {}
    for entry in model_def["inputs"]:
        target = required if entry["required"] else optional
        target[entry["name"]] = _comfy_input(entry)

    class ReplicateModelNode:
        CATEGORY = f"{CATEGORY_ROOT}/{model_def['category']}"
        FUNCTION = "run"
        RETURN_TYPES = return_types
        RETURN_NAMES = return_names
        OUTPUT_NODE = output_kind == "video"

        @classmethod
        def INPUT_TYPES(cls):
            spec = {"required": dict(required)}
            if optional:
                spec["optional"] = dict(optional)
            return spec

        def run(self, **kwargs):
            replicate = _client()
            args = _build_arguments(model_def, kwargs)
            print(f"[Replicate Select] {model_def['id']} ← {list(args.keys())}")
            result = replicate.run(model_def["id"], input=args)
            output = _first_output(result)
            if output is None:
                raise RuntimeError(f"{model_def['id']}: leere Antwort")
            data = _download(output)
            if output_kind == "image":
                return (_to_image_tensor(data),)
            if output_kind == "audio":
                return (_to_audio(data),)
            path, meta = _save_video(data, model_def["id"])
            # ui-Eintrag → Video erscheint in /history und ist via /view abrufbar (Client-Apps)
            return {"ui": {"images": [meta], "animated": (True,)}, "result": (path,)}

    ReplicateModelNode.__name__ = "Replicate_" + model_def["id"].replace("/", "_").replace("-", "_").replace(".", "_")
    return ReplicateModelNode


class ReplicateTextEncode:
    """Prompt-Passthrough (Text rein → Text raus).

    Client-Apps wie ComfyChair erkennen den positiven Prompt nur an Nodes,
    deren class_type 'TextEncode' oder 'Prompt' enthält — die generierten
    Replicate_*-Nodes fallen durch diesen Filter. Prompt daher über diesen
    Node in den prompt-Input des Modell-Nodes leiten.
    """

    CATEGORY = CATEGORY_ROOT
    FUNCTION = "run"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"text": ("STRING", {"multiline": True, "default": ""})}}

    def run(self, text):
        return (text,)


class ReplicateCustom:
    """Beliebiges Replicate-Modell per ID aufrufen — der Escape-Hatch."""

    CATEGORY = CATEGORY_ROOT
    FUNCTION = "run"
    RETURN_TYPES = ("STRING", "IMAGE")
    RETURN_NAMES = ("raw_output", "image")

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_id": ("STRING", {"default": "owner/model", "tooltip": "z.B. bytedance/seedream-4.5"}),
                "input_json": ("STRING", {"multiline": True, "default": '{\n  "prompt": ""\n}',
                                          "tooltip": "Input-Parameter als JSON"}),
            },
            "optional": {
                "image": ("IMAGE", {"tooltip": "Wird als 'image'-Parameter angehängt, falls gesetzt"}),
            },
        }

    def run(self, model_id, input_json, image=None):
        replicate = _client()
        args = json.loads(input_json)
        if image is not None:
            args["image"] = _image_to_uri(image)
        result = replicate.run(model_id.strip(), input=args)
        output = _first_output(result)
        raw = str(output)
        image_out = torch.zeros(1, 1, 1, 3)
        try:
            data = _download(output)
            if data[:8].startswith((b"\x89PNG", b"\xff\xd8")):
                image_out = _to_image_tensor(data)
            else:
                raw = f"{raw}\n(bytes: {len(data)})"
        except Exception:
            pass
        return (raw, image_out)


def _load_catalog() -> dict:
    with open(CATALOG_PATH, encoding="utf-8") as fh:
        return json.load(fh)


NODE_CLASS_MAPPINGS = {
    "ReplicateCustom": ReplicateCustom,
    "ReplicateTextEncode": ReplicateTextEncode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "ReplicateCustom": "Replicate Custom (beliebige Modell-ID)",
    "ReplicateTextEncode": "Prompt (Replicate Select)",
}

for _model in _load_catalog()["models"]:
    _cls = _make_node(_model)
    NODE_CLASS_MAPPINGS[_cls.__name__] = _cls
    NODE_DISPLAY_NAME_MAPPINGS[_cls.__name__] = f"{_model['name']} (Replicate)"