#!/usr/bin/env python3
"""
Trellis 3D model generator connector — Gradio API wrapper.

Talks to the Trellis Gradio app at tower-of-joy:11510.
Pipeline: upload image → start session → image_to_3d → extract_glb → download .glb

Usage:
    python3 trellis_connector.py health
    python3 trellis_connector.py generate image.png [--output model.glb] [--simplify 0.95] [--texture-size 1024] [--seed 42] [--timeout 600]

Gradio API Parameter Reference (TRELLIS v1, microsoft/TRELLIS):

    /image_to_3d — 9 inputs:
        0: image           (Image)   preprocessed image from /preprocess_image_1
        1: multiimages     (Gallery) [] for single-image mode
        2: is_multiimage   (State)   False for single-image, True for multi-image
        3: seed            (Slider)  int, 0-2147483647
        4: ss_guidance      (Slider)  float, sparse structure guidance strength (default 7.5)
        5: ss_steps         (Slider)  int, sparse structure sampling steps (default 12)
        6: slat_guidance    (Slider)  float, structured latent guidance strength (default 3.0)
        7: slat_steps       (Slider)  int, structured latent sampling steps (default 12)
        8: multiimage_algo  (Radio)   "stochastic" or "multidiffusion"

    /extract_glb — 3 inputs:
        0: output_buf      (State)   None — server uses internal state from image_to_3d
        1: simplify         (Slider)  float, mesh simplification ratio (default 0.95)
        2: texture_size     (Slider)  int, texture resolution (default 1024)

    Common failure modes:
        - "needed 9, got 8": missing is_multiimage (position 2) — must pass False
        - "needed 3, got 2": missing output_buf (position 0) — must pass None
        - "'float' cannot be interpreted as int": numpy version issue on server,
          Gradio Sliders send all values as float. Fix: patch flow_euler.py on
          the server to cast steps to int, or pin numpy < 2.0
        - CUDA device mismatch after crash: restart the container to clear GPU state
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import ensure_venv  # noqa: E402

ensure_venv()

import json
import os
import shutil
import time
import urllib.error
import urllib.parse
import urllib.request

CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")


def load_config():
    with open(CONFIG_PATH) as f:
        return json.load(f)


def get_base_url():
    config = load_config()
    return config.get("trellis_url", "http://tower-of-joy:11510")


def health():
    """Check if the Trellis API is reachable."""
    base = get_base_url()
    try:
        req = urllib.request.Request(f"{base}/info", method="GET")
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read())
            endpoints = list(data.get("named_endpoints", {}).keys())
            print(json.dumps({
                "ok": True,
                "url": base,
                "endpoints": endpoints
            }, indent=2))
    except Exception as e:
        print(json.dumps({
            "ok": False,
            "url": base,
            "error": str(e)
        }, indent=2))
        sys.exit(1)


def _call_api(base, endpoint, data, timeout=600, session_hash=None):
    """Call a Gradio API endpoint with optional session tracking.

    Gradio gr.State components are stored per session_hash on the server.
    All calls in a pipeline (image_to_3d → extract_glb) must share the same
    session_hash so the server can pass state between them.
    """
    api_url = f"{base}/api{endpoint}"

    body_dict = {"data": data}
    if session_hash:
        body_dict["session_hash"] = session_hash

    payload = json.dumps(body_dict)
    req = urllib.request.Request(
        api_url,
        data=payload.encode(),
        headers={"Content-Type": "application/json"},
        method="POST"
    )

    print(f"  Calling {endpoint}...", file=sys.stderr)

    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            result = json.loads(resp.read())
            if isinstance(result, dict) and "data" in result:
                return result["data"]
            return result
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"{endpoint} failed ({e.code}): {body[:300]}")


def _upload_image(base, image_path):
    """Upload an image file to the Gradio server and return the file reference."""
    upload_url = f"{base}/upload"

    with open(image_path, "rb") as f:
        image_data = f.read()

    filename = os.path.basename(image_path)

    # Gradio upload expects multipart/form-data with a 'files' field
    boundary = "----TrellisConnectorBoundary"
    body = (
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="files"; filename="{filename}"\r\n'
        f"Content-Type: image/png\r\n"
        f"\r\n"
    ).encode() + image_data + f"\r\n--{boundary}--\r\n".encode()

    req = urllib.request.Request(
        upload_url,
        data=body,
        headers={
            "Content-Type": f"multipart/form-data; boundary={boundary}",
        },
        method="POST"
    )

    print(f"  Uploading {filename}...", file=sys.stderr)
    with urllib.request.urlopen(req, timeout=30) as resp:
        result = json.loads(resp.read())
        # Gradio returns a list of uploaded file paths
        if isinstance(result, list) and len(result) > 0:
            return result[0]
        raise RuntimeError(f"Upload failed: {result}")


def _download_file(url, output_path, base):
    """Download a file from the Gradio server."""
    if url.startswith("/"):
        url = f"{base}{url}"
    elif not url.startswith("http"):
        url = f"{base}/file={url}"

    print(f"  Downloading to {output_path}...", file=sys.stderr)
    req = urllib.request.Request(url, method="GET")
    with urllib.request.urlopen(req, timeout=120) as resp:
        with open(output_path, "wb") as f:
            shutil.copyfileobj(resp, f)

    return os.path.getsize(output_path)


def _check_available(base):
    """Quick check if Trellis is reachable. Fail fast with a clear message."""
    try:
        req = urllib.request.Request(f"{base}/info", method="GET")
        urllib.request.urlopen(req, timeout=5)
    except Exception:
        print(json.dumps({
            "ok": False,
            "error": f"Trellis is not available at {base}. The service may be switched off to save system resources. Start it before generating 3D models."
        }, indent=2))
        sys.exit(1)


def generate(image_path, output=None, simplify=0.95, texture_size=1024,
             seed=0, timeout=600):
    """
    Generate a 3D model from an image.

    Pipeline:
        1. Start session
        2. Upload and preprocess image
        3. Generate 3D from image
        4. Extract GLB
        5. Download GLB file

    Args:
        image_path: Path to the input image (PNG recommended)
        output: Output .glb file path (default: auto-named)
        simplify: Mesh simplification factor (0.9-0.98, default 0.95)
        texture_size: Texture resolution (512-2048, default 1024)
        seed: Random seed (default 0)
        timeout: Max wait time per step in seconds
    """
    base = get_base_url()
    _check_available(base)
    start_time = time.time()

    if not os.path.isfile(image_path):
        print(json.dumps({"ok": False, "error": f"Image not found: {image_path}"}), indent=2)
        sys.exit(1)

    if output is None:
        name = os.path.splitext(os.path.basename(image_path))[0]
        output = f"{name}.glb"

    # Generate a session hash — Gradio uses this to maintain gr.State between
    # separate API calls. Without it, image_to_3d's output state is lost before
    # extract_glb can read it.
    import random, string
    session = ''.join(random.choices(string.ascii_lowercase + string.digits, k=12))
    print(f"  Session: {session}", file=sys.stderr)

    # Step 1: Start session
    print("Step 1/5: Starting session...", file=sys.stderr)
    session_result = _call_api(base, "/start_session", [], timeout=30, session_hash=session)

    # Step 2: Upload and preprocess image
    print("Step 2/5: Uploading and preprocessing image...", file=sys.stderr)
    uploaded_path = _upload_image(base, image_path)
    file_ref = {
        "path": uploaded_path,
        "meta": {"_type": "gradio.FileData"}
    }
    preprocess_result = _call_api(base, "/preprocess_image_1", [file_ref], timeout=60, session_hash=session)

    # _call_api returns the "data" array directly
    if isinstance(preprocess_result, list) and len(preprocess_result) > 0:
        preprocessed_ref = preprocess_result[0]
    else:
        preprocessed_ref = preprocess_result

    # Step 3: Get seed
    print("Step 3/5: Generating 3D model...", file=sys.stderr)
    seed_result = _call_api(base, "/get_seed", [True, seed], timeout=10, session_hash=session)
    if isinstance(seed_result, list) and seed_result:
        actual_seed = seed_result[0]
    else:
        actual_seed = seed

    # Step 4: Image to 3D
    # Gradio app has 9 inputs: Image, Gallery, State, Slider×5, Radio
    # State is a hidden session component at position 3 — pass None.
    # Slider order (from /info): seed, ss_guidance, ss_steps, slat_guidance, slat_steps
    # Note: server-side numpy bug requires patching flow_euler.py to cast steps to int.
    gen_result = _call_api(base, "/image_to_3d", [
        preprocessed_ref,  # image
        [],                # multiimages (empty)
        False,             # is_multiimage (boolean flag, not session state)
        actual_seed,       # seed (Slider, 0-2147483647)
        7.5,               # ss_guidance_strength
        12,                # ss_sampling_steps
        3.0,               # slat_guidance_strength
        12,                # slat_sampling_steps
        "stochastic",      # multiimage_algo
    ], timeout=timeout, session_hash=session)

    # Step 5: Extract GLB
    # State (output_buf) is maintained server-side via session_hash.
    # extract_glb reads it automatically. Pass None as placeholder for the State component.
    print("Step 4/5: Extracting GLB...", file=sys.stderr)
    glb_result = _call_api(base, "/extract_glb", [None, simplify, texture_size], timeout=120, session_hash=session)

    # _call_api returns the "data" array: [model_viewer_data, download_button_data]
    glb_url = None
    if isinstance(glb_result, list):
        for item in glb_result:
            if isinstance(item, dict):
                url = item.get("url") or item.get("path")
                if url:
                    glb_url = url
                    break

    if not glb_url:
        print(json.dumps({
            "ok": False,
            "error": "Could not extract GLB URL from response",
            "response": glb_result
        }, indent=2))
        sys.exit(1)

    # Step 6: Download
    print("Step 5/5: Downloading GLB...", file=sys.stderr)
    os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True)
    file_size = _download_file(glb_url, output, base)

    elapsed = round(time.time() - start_time, 1)
    print(json.dumps({
        "ok": True,
        "file": output,
        "size_bytes": file_size,
        "simplify": simplify,
        "texture_size": texture_size,
        "seed": actual_seed,
        "generation_time_s": elapsed,
        "source_image": image_path
    }, indent=2))


def main():
    if len(sys.argv) < 2:
        print("Usage:")
        print("  trellis_connector.py health")
        print("  trellis_connector.py generate image.png [--output model.glb] [--simplify 0.95] [--texture-size 1024] [--seed N] [--timeout N]")
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "health":
        health()
    elif cmd == "generate":
        if len(sys.argv) < 3:
            print("Error: image path required", file=sys.stderr)
            sys.exit(1)

        image_path = sys.argv[2]
        output = None
        simplify = 0.95
        texture_size = 1024
        seed = 0
        timeout = 600

        i = 3
        while i < len(sys.argv):
            if sys.argv[i] == "--output" and i + 1 < len(sys.argv):
                output = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == "--simplify" and i + 1 < len(sys.argv):
                simplify = float(sys.argv[i + 1])
                i += 2
            elif sys.argv[i] == "--texture-size" and i + 1 < len(sys.argv):
                texture_size = int(sys.argv[i + 1])
                i += 2
            elif sys.argv[i] == "--seed" and i + 1 < len(sys.argv):
                seed = int(sys.argv[i + 1])
                i += 2
            elif sys.argv[i] == "--timeout" and i + 1 < len(sys.argv):
                timeout = int(sys.argv[i + 1])
                i += 2
            else:
                print(f"Unknown argument: {sys.argv[i]}", file=sys.stderr)
                sys.exit(1)

        generate(image_path, output=output, simplify=simplify,
                 texture_size=texture_size, seed=seed, timeout=timeout)
    else:
        print(f"Unknown command: {cmd}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()