import gradio as gr
import requests
import tempfile
import os
import trimesh
import json

# ------------------------------
# Load default models from JSON
# ------------------------------
# Option 1: Load external models.json (if exists)
def load_default_models():
    try:
        with open("models.json", "r") as f:
            data = json.load(f)
            return data.get("default_models", [])
    except FileNotFoundError:
        # Fallback embedded list
        return [
            {"name": "Damaged Helmet", "url": "https://threejs.org/examples/models/gltf/DamagedHelmet/glTF/DamagedHelmet.gltf"},
            {"name": "Flight Helmet", "url": "https://threejs.org/examples/models/gltf/FlightHelmet/glTF/FlightHelmet.gltf"},
            {"name": "Horse", "url": "https://threejs.org/examples/models/gltf/Horse.glb"},
            {"name": "Fox", "url": "https://threejs.org/examples/models/gltf/Fox.glb"},
            {"name": "Camera", "url": "https://raw.githubusercontent.com/ektogamat/camera-webgi/master/assets/camera.glb"},
        ]

DEFAULT_MODELS = load_default_models()

# ------------------------------
# Helper functions
# ------------------------------
def download_model(url):
    try:
        response = requests.get(url, stream=True, timeout=15)
        response.raise_for_status()
        # Determine suffix from URL or content-type
        suffix = ".glb"
        if url.lower().endswith(".gltf"):
            suffix = ".gltf"
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
            for chunk in response.iter_content(chunk_size=8192):
                tmp.write(chunk)
            return tmp.name
    except Exception as e:
        print(f"Download error: {e}")
        return None

def load_and_convert_model(url, target_format):
    if not url:
        return None
    # Download original
    original_path = download_model(url)
    if not original_path:
        return None
    try:
        mesh = trimesh.load(original_path)
        # Export
        ext = target_format.lower()
        with tempfile.NamedTemporaryFile(suffix=f".{ext}", delete=False) as tmp:
            out_path = tmp.name
        if ext in ["obj", "stl", "ply", "glb", "gltf"]:
            mesh.export(out_path, file_type=ext)
        else:
            return None
        return out_path
    except Exception as e:
        print(f"Conversion error: {e}")
        return None
    finally:
        if os.path.exists(original_path):
            os.unlink(original_path)

def display_model_from_url(url):
    if not url:
        return None
    return download_model(url)

# ------------------------------
# Gradio UI
# ------------------------------
with gr.Blocks(title="3D Model Viewer & Converter") as demo:
    gr.Markdown("## 🌐 3D Model Hub")
    gr.Markdown("Select a default model or paste any GLB/GLTF URL.")

    with gr.Row():
        # Dropdown for default models
        default_choices = [m["name"] for m in DEFAULT_MODELS]
        default_dropdown = gr.Dropdown(
            choices=default_choices,
            label="📦 Default Models",
            value=default_choices[0] if default_choices else None
        )
        # Custom URL input
        url_input = gr.Textbox(
            label="🔗 Or paste custom URL",
            placeholder="https://example.com/model.glb",
            scale=3
        )
        load_btn = gr.Button("Load Model", variant="primary")

    model_output = gr.Model3D(
        label="3D Preview",
        clear_color=[0.2, 0.2, 0.2, 0.8],
        camera_position=(2, 2, 3)
    )

    gr.Markdown("### 💾 Download as...")
    with gr.Row():
        obj_btn = gr.Button("OBJ")
        stl_btn = gr.Button("STL")
        ply_btn = gr.Button("PLY")
        glb_btn = gr.Button("GLB")
        gltf_btn = gr.Button("glTF")

    # Hidden file outputs
    obj_file = gr.File(label="Download OBJ", visible=False)
    stl_file = gr.File(label="Download STL", visible=False)
    ply_file = gr.File(label="Download PLY", visible=False)
    glb_file = gr.File(label="Download GLB", visible=False)
    gltf_file = gr.File(label="Download glTF", visible=False)

    # State to hold current URL
    current_url = gr.State()

    # When dropdown changes, update the custom URL textbox (optional)
    def update_url_from_dropdown(selected_name):
        for model in DEFAULT_MODELS:
            if model["name"] == selected_name:
                return model["url"]
        return ""
    
    default_dropdown.change(fn=update_url_from_dropdown, inputs=default_dropdown, outputs=url_input)

    # Load model from either dropdown or custom URL
    def load_model(selected_name, custom_url):
        # Prioritize custom URL if non-empty, else use selected default
        if custom_url and custom_url.strip():
            url = custom_url.strip()
        else:
            for model in DEFAULT_MODELS:
                if model["name"] == selected_name:
                    url = model["url"]
                    break
            else:
                return None
        # Download and return file path for preview
        return download_model(url)

    load_btn.click(
        fn=load_model,
        inputs=[default_dropdown, url_input],
        outputs=model_output
    ).then(
        fn=lambda selected_name, custom_url: (
            custom_url.strip() if custom_url and custom_url.strip()
            else next((m["url"] for m in DEFAULT_MODELS if m["name"] == selected_name), "")
        ),
        inputs=[default_dropdown, url_input],
        outputs=current_url
    )

    # Conversion buttons
    def convert_and_download(url, fmt):
        if not url:
            return None
        return load_and_convert_model(url, fmt)

    for fmt, btn, out in [
        ("obj", obj_btn, obj_file),
        ("stl", stl_btn, stl_file),
        ("ply", ply_btn, ply_file),
        ("glb", glb_btn, glb_file),
        ("gltf", gltf_btn, gltf_file),
    ]:
        btn.click(
            fn=lambda url, f=fmt: convert_and_download(url, f),
            inputs=current_url,
            outputs=out
        ).then(
            fn=lambda: gr.update(visible=True),
            outputs=out
        )

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