"""Shared helper for fal.ai operators: upload input image, call an endpoint with
streamed progress, and download every 3D file found in the result.

Runs in the server's own venv (fal-client is dependency-light). Kept generic so
operators only supply endpoint id + argument mapping; unknown/new fal args pass
through via an `extra_args` JSON param so the UI keeps full control."""
import argparse
import json
import sys
import urllib.request
from pathlib import Path

MODEL_EXTS = {".glb", ".gltf", ".obj", ".fbx", ".usdz", ".ply", ".stl", ".zip"}


def _log(msg):
    print(msg, flush=True)


def _find_file_urls(obj, acc):
    """Recursively collect (url, ext) for any file-like url in the result JSON."""
    if isinstance(obj, dict):
        url = obj.get("url")
        if isinstance(url, str):
            ext = Path(url.split("?")[0]).suffix.lower()
            if ext in MODEL_EXTS:
                acc.append((url, ext, obj.get("file_name")))
        for v in obj.values():
            _find_file_urls(v, acc)
    elif isinstance(obj, list):
        for v in obj:
            _find_file_urls(v, acc)


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", action="append", default=[])
    ap.add_argument("--outdir", required=True)
    ap.add_argument("--params", default="{}")
    a = ap.parse_args()
    return a, json.loads(a.params)


def run(endpoint, arguments, outdir, image_path=None, image_arg="image_url",
        image_as_list=False, extra_args=None, out_stem="model"):
    import fal_client

    if image_path:
        _log(f"uploading {Path(image_path).name} to fal ...")
        url = fal_client.upload_file(image_path)
        arguments[image_arg] = [url] if image_as_list else url

    if extra_args:
        try:
            arguments.update(json.loads(extra_args) if isinstance(extra_args, str) else extra_args)
        except ValueError as e:
            _log(f"warning: could not parse extra_args ({e}); ignoring")

    # drop None values so we send only what the user set
    arguments = {k: v for k, v in arguments.items() if v is not None and v != ""}
    _log(f"calling {endpoint}")
    _log(f"arguments: {json.dumps({k: v for k, v in arguments.items() if k != image_arg})}")

    def on_update(update):
        for entry in getattr(update, "logs", None) or []:
            msg = entry.get("message") if isinstance(entry, dict) else str(entry)
            if msg:
                _log(f"  {msg}")

    result = fal_client.subscribe(endpoint, arguments=arguments, with_logs=True,
                                  on_queue_update=on_update)

    outdir = Path(outdir)
    (outdir / "fal_result.json").write_text(json.dumps(result, indent=2))

    urls = []
    _find_file_urls(result, urls)
    if not urls:
        _log("ERROR: no 3D file url found in fal result. Raw result saved to fal_result.json:")
        _log(json.dumps(result, indent=2)[:2000])
        sys.exit(1)

    outputs = []
    for i, (url, ext, fname) in enumerate(urls):
        name = fname or (f"{out_stem}{ext}" if i == 0 else f"{out_stem}_{i}{ext}")
        dest = outdir / name
        _log(f"downloading {name} ...")
        urllib.request.urlretrieve(url, dest)
        outputs.append({"path": name, "meta": {"endpoint": endpoint, "source_url": url}})

    (outdir / "result.json").write_text(json.dumps({"outputs": outputs}))
    _log(f"done — {len(outputs)} file(s)")