"""
sacruise/threed-precision
VAST-AI TripoSG image-to-3D API space.
Best for: mechanical parts, tools, furniture, vehicles, precision/hard-surface objects.
"""

import base64
import io
import json
import logging
import os
import sys
import tempfile
import traceback

import gradio as gr
import spaces
import torch
from PIL import Image

logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)s: %(message)s")
logger = logging.getLogger("threed-precision")

# Clone TripoSG repo at startup (pip git+https install is unreliable on ZeroGPU)
TRIPOSG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "triposg_repo")
if not os.path.exists(TRIPOSG_DIR):
    logger.info("Cloning TripoSG repository…")
    ret = os.system(f"git clone https://github.com/VAST-AI-Research/TripoSG.git {TRIPOSG_DIR}")
    if ret != 0:
        logger.error(f"git clone failed (exit {ret})")
    else:
        logger.info("TripoSG cloned.")
else:
    logger.info("TripoSG repo present.")

sys.path.insert(0, TRIPOSG_DIR)
sys.path.insert(0, os.path.join(TRIPOSG_DIR, "scripts"))

PIPELINE_LOAD_ERROR: str | None = None
pipeline = None

try:
    logger.info("Loading TripoSG pipeline to CPU RAM…")
    from huggingface_hub import snapshot_download
    from triposg.pipelines.pipeline_triposg import TripoSGPipeline
    model_path = snapshot_download("VAST-AI/TripoSG")
    pipeline = TripoSGPipeline.from_pretrained(model_path)
    logger.info("TripoSG loaded on CPU. Ready.")
except Exception as e:
    PIPELINE_LOAD_ERROR = f"TripoSG load failed: {e}\n{traceback.format_exc()}"
    logger.error(PIPELINE_LOAD_ERROR)


@spaces.GPU(duration=120)
def generate_api(images_json: str, seed: int = 42) -> str:
    import rembg

    if pipeline is None:
        raise RuntimeError(f"Pipeline not available: {PIPELINE_LOAD_ERROR}")

    pipeline.to("cuda", dtype=torch.float16)

    try:
        logger.info("generate_api: decoding input image")
        images_data = json.loads(images_json)
        if not images_data or not isinstance(images_data, list):
            raise ValueError("images_json must be a non-empty JSON array")

        raw = images_data[0]
        if isinstance(raw, str) and raw.startswith("data:"):
            raw = raw.split(",", 1)[1]
        img_bytes = base64.b64decode(raw)
        image = Image.open(io.BytesIO(img_bytes)).convert("RGBA")
        logger.info(f"Image decoded: {image.size}")

        # Force CPU provider so onnxruntime doesn't conflict with ZeroGPU CUDA context
        session = rembg.new_session("u2net", providers=["CPUExecutionProvider"])
        image = rembg.remove(image, session=session)

        generator = torch.Generator(device="cuda").manual_seed(int(seed))

        logger.info("Running TripoSG pipeline…")
        outputs = pipeline(
            image=image,
            num_inference_steps=8,
            guidance_scale=7.5,
            generator=generator,
        )

        logger.info("Exporting GLB…")
        # Handle both output formats: obj.mesh[0] or outputs[0]
        if hasattr(outputs, "mesh"):
            mesh = outputs.mesh[0]
        elif hasattr(outputs, "meshes"):
            mesh = outputs.meshes[0]
        elif isinstance(outputs, (list, tuple)):
            mesh = outputs[0]
        else:
            mesh = outputs

        tmp = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
        mesh.export(tmp.name)
        tmp.close()
        logger.info(f"GLB written: {tmp.name}")
        return tmp.name

    except Exception as e:
        logger.error(f"generate_api error: {e}\n{traceback.format_exc()}")
        raise
    finally:
        pipeline.to("cpu")
        torch.cuda.empty_cache()


demo = gr.Interface(
    fn=generate_api,
    inputs=[
        gr.Textbox(label="Base64 Images JSON", placeholder='["<base64_string>"]'),
        gr.Slider(0, 65535, value=42, step=1, label="Seed"),
    ],
    outputs=gr.File(label="GLB File"),
    title="TripoSG Precision Generator",
    description=(
        "Generate clean, precise 3D meshes from images using VAST-AI TripoSG. "
        f"{'⚠️ Pipeline load error: ' + PIPELINE_LOAD_ERROR if PIPELINE_LOAD_ERROR else '✅ Pipeline loaded.'}"
    ),
    api_name="generate_api",
)

if __name__ == "__main__":
    demo.launch()