"""
AI 3D Generation Pipeline — Hugging Face Space (ZeroGPU)
========================================================
A production-ready Gradio application that generates high-fidelity,
manifold 3D models from multi-angle images + text instructions.

Pipeline:
  1. Background removal (RMBG-1.4)
  2. Multi-image composite / best-view selection
  3. TripoSG 3D reconstruction
  4. Manifold repair + scale normalization
  5. GLB/OBJ/STL export

Designed for Hugging Face ZeroGPU (A100-40GB).
"""
from __future__ import annotations

import os
import io
import re
import sys
import json
import uuid
import time
import shutil
import base64
import logging
import tempfile
import zipfile
from pathlib import Path
from typing import Optional, List, Tuple, Dict

import gradio as gr
import numpy as np
import torch
from PIL import Image

# HuggingFace ZeroGPU decorator
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False
    # Fallback no-op decorator for local dev
    class spaces:
        @staticmethod
        def GPU(fn=None, duration=None):
            if fn is None:
                return lambda f: f
            return fn

# Mesh processing
import trimesh

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_ID = "VAST-AI/TripoSG"
RMBG_MODEL_ID = "briaai/RMBG-1.4"
TRIPOSG_REPO_URL = "https://github.com/VAST-AI-Research/TripoSG.git"
TRIPOSG_CODE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "triposg_repo")
MAX_FACES = 90000
DEFAULT_FACES = 50000
OUTPUT_DIR = Path(tempfile.gettempdir()) / "triposg_output"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

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

# Clone TripoSG repo if not present (needed for pipeline imports)
if not os.path.exists(TRIPOSG_CODE_DIR):
    logger.info("Cloning TripoSG repository...")
    os.system(f"git clone {TRIPOSG_REPO_URL} {TRIPOSG_CODE_DIR}")
sys.path.insert(0, TRIPOSG_CODE_DIR)
sys.path.insert(0, os.path.join(TRIPOSG_CODE_DIR, "scripts"))

# ---------------------------------------------------------------------------
# Global model references (loaded once)
# ---------------------------------------------------------------------------
triposg_pipeline = None
rmbg_model = None
device = "cuda" if torch.cuda.is_available() else "cpu"


def load_models():
    """Load TripoSG + RMBG models into GPU memory."""
    global triposg_pipeline, rmbg_model

    if triposg_pipeline is not None:
        return  # already loaded

    logger.info("Loading TripoSG pipeline from %s ...", MODEL_ID)
    try:
        from huggingface_hub import snapshot_download
        from triposg.pipelines.pipeline_triposg import TripoSGPipeline

        # Download model weights
        model_path = snapshot_download(MODEL_ID)
        logger.info("Model downloaded to: %s", model_path)

        # Load pipeline onto GPU
        triposg_pipeline = TripoSGPipeline.from_pretrained(model_path).to(
            device, torch.float16 if device == "cuda" else torch.float32
        )
        logger.info("TripoSG pipeline loaded successfully")

    except Exception as e:
        logger.error("Failed to load TripoSG: %s", e)
        triposg_pipeline = None

    # Load background remover
    logger.info("Loading RMBG background remover...")
    try:
        from transformers import pipeline as hf_pipeline
        rmbg_model = hf_pipeline("image-segmentation", model=RMBG_MODEL_ID, device=device)
        logger.info("RMBG loaded successfully")
    except Exception as e:
        logger.warning("RMBG not available, will skip background removal: %s", e)
        rmbg_model = None


# ---------------------------------------------------------------------------
# Image Processing Utilities
# ---------------------------------------------------------------------------
def remove_background(image: Image.Image) -> Image.Image:
    """Remove background from an image using RMBG-1.4."""
    if rmbg_model is None:
        return image
    try:
        results = rmbg_model(image)
        # Find the mask and apply it
        for result in results:
            if result.get("label") in ["foreground", "person", "object"]:
                mask = result["mask"]
                image = image.convert("RGBA")
                image.putalpha(mask)
                return image
        # If no specific label, use the first result
        if results:
            mask = results[0]["mask"]
            image = image.convert("RGBA")
            image.putalpha(mask)
        return image
    except Exception as e:
        logger.warning("Background removal failed: %s", e)
        return image


def select_best_view(images: list[Image.Image]) -> Image.Image:
    """
    From multiple input images, select the best single view for reconstruction.
    Heuristic: pick the largest image (highest resolution = most detail).
    In a production pipeline, you would use a multi-view reconstruction model.
    """
    if not images:
        raise ValueError("No images provided")
    if len(images) == 1:
        return images[0]

    # Score images by resolution and sharpness
    scored = []
    for img in images:
        w, h = img.size
        resolution_score = w * h
        # Simple Laplacian variance for sharpness
        gray = np.array(img.convert("L"), dtype=np.float32)
        laplacian = np.abs(np.roll(gray, 1, 0) + np.roll(gray, -1, 0) +
                          np.roll(gray, 1, 1) + np.roll(gray, -1, 1) - 4 * gray)
        sharpness = float(np.var(laplacian))
        scored.append((img, resolution_score * 0.3 + sharpness * 0.7))

    scored.sort(key=lambda x: x[1], reverse=True)
    logger.info("Selected best view (score=%.1f) from %d images", scored[0][1], len(images))
    return scored[0][0]


def create_composite_view(images: list[Image.Image], max_images: int = 4) -> Image.Image:
    """
    Create a multi-view composite image (2x2 grid) from multiple input images.
    Some models work better with a single composite showing multiple angles.
    """
    images = images[:max_images]
    n = len(images)

    if n == 1:
        return images[0]

    # Resize all to same size
    target_size = 512
    resized = [img.resize((target_size, target_size), Image.LANCZOS) for img in images]

    if n == 2:
        composite = Image.new("RGB", (target_size * 2, target_size))
        composite.paste(resized[0], (0, 0))
        composite.paste(resized[1], (target_size, 0))
    elif n <= 4:
        cols, rows = 2, 2
        composite = Image.new("RGB", (target_size * cols, target_size * rows), (255, 255, 255))
        for i, img in enumerate(resized):
            x = (i % cols) * target_size
            y = (i // cols) * target_size
            composite.paste(img, (x, y))
    else:
        composite = resized[0]

    return composite


# ---------------------------------------------------------------------------
# Text-to-Dimension Extraction
# ---------------------------------------------------------------------------
def extract_dimensions_from_prompt(prompt: str) -> dict:
    """
    Parse dimensional info from text prompt.
    Supports patterns like:
      - "10cm x 5cm x 3cm"
      - "base is 10cm"
      - "height: 200mm"
      - "scale 1:10"
    Returns dict with extracted measurements.
    """
    dims = {}

    # Match "NxMxP" patterns (cm, mm, m, inches)
    pattern_3d = r'(\d+(?:\.\d+)?)\s*(?:x|×|by)\s*(\d+(?:\.\d+)?)\s*(?:x|×|by)\s*(\d+(?:\.\d+)?)\s*(cm|mm|m|in|inch|inches)?'
    match = re.search(pattern_3d, prompt, re.IGNORECASE)
    if match:
        dims["width"] = float(match.group(1))
        dims["depth"] = float(match.group(2))
        dims["height"] = float(match.group(3))
        dims["unit"] = (match.group(4) or "cm").lower().rstrip("es").rstrip("ch")
        return dims

    # Match individual dimension patterns
    for dim_name in ["width", "height", "depth", "length", "base", "diameter", "radius"]:
        pattern = rf'{dim_name}\s*(?:is|:|\=)?\s*(\d+(?:\.\d+)?)\s*(cm|mm|m|in|inch|inches)?'
        match = re.search(pattern, prompt, re.IGNORECASE)
        if match:
            dims[dim_name] = float(match.group(1))
            dims["unit"] = (match.group(2) or "cm").lower().rstrip("es").rstrip("ch")

    # Match scale ratios
    scale_pattern = r'(?:scale|ratio)\s*(?:is|:|\=)?\s*1\s*:\s*(\d+(?:\.\d+)?)'
    match = re.search(scale_pattern, prompt, re.IGNORECASE)
    if match:
        dims["scale_ratio"] = float(match.group(1))

    return dims


# ---------------------------------------------------------------------------
# Mesh Post-Processing (Manifold + Scale)
# ---------------------------------------------------------------------------
def repair_mesh(mesh: trimesh.Trimesh) -> trimesh.Trimesh:
    """Make a mesh manifold (watertight) for 3D printing."""
    logger.info("Repairing mesh: %d vertices, %d faces", len(mesh.vertices), len(mesh.faces))

    # Fix normals
    trimesh.repair.fix_normals(mesh)
    trimesh.repair.fix_inversion(mesh)
    trimesh.repair.fix_winding(mesh)

    # Fill holes
    if not mesh.is_watertight:
        mesh.fill_holes()
        logger.info("Filled holes → watertight: %s", mesh.is_watertight)

    # Remove degenerate faces
    mask = mesh.nondegenerate_faces()
    if not mask.all():
        mesh.update_faces(mask)
        logger.info("Removed %d degenerate faces", (~mask).sum())

    # Merge close vertices
    mesh.merge_vertices(merge_tex=True, merge_norm=True)

    return mesh


def scale_mesh(mesh: trimesh.Trimesh, dimensions: dict) -> trimesh.Trimesh:
    """
    Scale mesh based on extracted dimensions.
    The model outputs are in arbitrary units — we map the largest extent
    to the largest provided dimension.
    """
    if not dimensions:
        return mesh

    # Convert everything to millimeters for consistency
    unit_to_mm = {"mm": 1.0, "cm": 10.0, "m": 1000.0, "in": 25.4}
    unit = dimensions.get("unit", "cm")
    multiplier = unit_to_mm.get(unit, 10.0)

    # Find the target size
    target_dims = []
    for key in ["width", "height", "depth", "length", "base", "diameter"]:
        if key in dimensions:
            target_dims.append(dimensions[key] * multiplier)  # convert to mm

    if not target_dims:
        if "scale_ratio" in dimensions:
            # Apply scale ratio directly
            scale_factor = 1.0 / dimensions["scale_ratio"]
            mesh.apply_scale(scale_factor)
            return mesh
        return mesh

    # Scale the mesh so its largest extent matches the largest target dimension
    current_extents = mesh.extents
    max_current = max(current_extents)
    max_target = max(target_dims)

    if max_current > 0:
        scale_factor = max_target / max_current
        mesh.apply_scale(scale_factor)
        logger.info("Scaled mesh by %.3f (target: %.1fmm)", scale_factor, max_target)

    return mesh


def segment_mesh_parts(mesh: trimesh.Trimesh) -> list[trimesh.Trimesh]:
    """
    Attempt to split a mesh into separate connected components (modular parts).
    This is useful for complex objects where parts should be independently editable.
    """
    try:
        components = mesh.split(only_watertight=False)
        if len(components) > 1:
            logger.info("Mesh split into %d modular parts", len(components))
            return components
        else:
            logger.info("Mesh is a single connected component")
            return [mesh]
    except Exception as e:
        logger.warning("Mesh segmentation failed: %s", e)
        return [mesh]


# ---------------------------------------------------------------------------
# Export Utilities
# ---------------------------------------------------------------------------
def export_mesh(mesh: trimesh.Trimesh, fmt: str = "glb", parts: list = None) -> str:
    """
    Export mesh to file. If parts are provided, creates a ZIP with individual files.
    """
    job_id = str(uuid.uuid4())[:8]
    output_base = OUTPUT_DIR / job_id
    output_base.mkdir(parents=True, exist_ok=True)

    if parts and len(parts) > 1:
        # Export each part separately and zip them
        part_files = []
        for i, part in enumerate(parts):
            fname = f"part_{i+1:02d}.{fmt}"
            fpath = output_base / fname
            part.export(str(fpath))
            part_files.append(str(fpath))

        # Also export the combined model
        combined_path = output_base / f"combined.{fmt}"
        mesh.export(str(combined_path))
        part_files.append(str(combined_path))

        # Create zip
        zip_path = output_base / f"model_parts_{job_id}.zip"
        with zipfile.ZipFile(str(zip_path), "w") as zf:
            for pf in part_files:
                zf.write(pf, os.path.basename(pf))

        logger.info("Exported %d parts + combined to %s", len(parts), zip_path)
        return str(combined_path), str(zip_path)
    else:
        fpath = output_base / f"model.{fmt}"
        mesh.export(str(fpath))
        logger.info("Exported single mesh to %s", fpath)
        return str(fpath), str(fpath)


# ---------------------------------------------------------------------------
# Main Generation Pipeline
# ---------------------------------------------------------------------------
@spaces.GPU(duration=300)
def generate_3d_model(
    input_images: list,
    text_prompt: str,
    num_faces: int = DEFAULT_FACES,
    export_format: str = "glb",
    enable_modular: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """
    End-to-end 3D generation pipeline.

    Args:
        input_images: List of uploaded image file paths
        text_prompt: Text instructions (dimensions, details, style)
        num_faces: Target face count for the mesh
        export_format: Output format (glb, obj, stl)
        enable_modular: Whether to split into modular parts
    """
    if not input_images and not text_prompt:
        raise gr.Error("Please provide at least one image or a text prompt.")

    start_time = time.time()
    status_log = []

    def log(msg):
        status_log.append(msg)
        logger.info(msg)

    # --- Step 1: Load models ---
    log("🔄 Loading models...")
    load_models()

    if triposg_pipeline is None:
        raise gr.Error("Model failed to load. Please try again or check the Space logs.")

    # --- Step 2: Process images ---
    log("🖼️ Processing input images...")
    pil_images = []
    if input_images:
        for img_path in input_images:
            if isinstance(img_path, str):
                img = Image.open(img_path).convert("RGB")
            elif hasattr(img_path, "name"):
                img = Image.open(img_path.name).convert("RGB")
            else:
                img = img_path.convert("RGB") if isinstance(img_path, Image.Image) else None
            if img:
                pil_images.append(img)

    # --- Step 3: Remove backgrounds ---
    log("✂️ Removing backgrounds...")
    clean_images = [remove_background(img) for img in pil_images]

    # --- Step 4: Select best view / create composite ---
    if clean_images:
        if len(clean_images) == 1:
            primary_image = clean_images[0]
            log("📷 Using single image input")
        else:
            primary_image = select_best_view(clean_images)
            log(f"📷 Selected best view from {len(clean_images)} images")
    else:
        primary_image = None

    # --- Step 5: Extract dimensions from prompt ---
    dimensions = {}
    if text_prompt:
        dimensions = extract_dimensions_from_prompt(text_prompt)
        if dimensions:
            log(f"📏 Extracted dimensions: {dimensions}")

    # --- Step 6: Run TripoSG inference ---
    log("🧠 Running 3D reconstruction (this may take 1-3 minutes)...")

    try:
        if not primary_image:
            raise gr.Error("Image input is required for 3D generation.")

        # Run TripoSG pipeline (matches official VAST-AI/TripoSG Space)
        import random
        seed = random.randint(0, 2**31 - 1)
        outputs = triposg_pipeline(
            image=primary_image,
            generator=torch.Generator(device=triposg_pipeline.device).manual_seed(seed),
            num_inference_steps=50,
            guidance_scale=7.0,
        ).samples[0]

        # outputs is (vertices, faces) tuple
        mesh = trimesh.Trimesh(
            vertices=outputs[0].astype(np.float32),
            faces=np.ascontiguousarray(outputs[1]),
        )
        log(f"Raw mesh: {len(mesh.vertices)} vertices, {len(mesh.faces)} faces")

    except Exception as e:
        raise gr.Error(f"3D reconstruction failed: {str(e)}")

    # --- Step 7: Post-process mesh ---
    log("🔧 Repairing mesh for manifold output...")
    mesh = repair_mesh(mesh)

    # --- Step 8: Apply dimensional scaling ---
    if dimensions:
        log("📐 Applying dimensional scaling...")
        mesh = scale_mesh(mesh, dimensions)

    # --- Step 9: Modular segmentation ---
    parts = None
    if enable_modular:
        log("🧩 Segmenting into modular parts...")
        parts = segment_mesh_parts(mesh)

    # --- Step 10: Export ---
    log(f"💾 Exporting as .{export_format}...")
    preview_path, download_path = export_mesh(mesh, export_format, parts)

    elapsed = time.time() - start_time
    log(f"✅ Complete in {elapsed:.1f}s — manifold: {mesh.is_watertight}, "
        f"vertices: {len(mesh.vertices)}, faces: {len(mesh.faces)}")

    status_text = "\n".join(status_log)
    return preview_path, download_path, status_text


# ---------------------------------------------------------------------------
# Gradio API Endpoint (for n8n / external calls)
# ---------------------------------------------------------------------------
@spaces.GPU(duration=300)
def api_generate(
    image_b64_list: str,
    prompt: str,
    faces: int = DEFAULT_FACES,
    fmt: str = "glb",
    modular: bool = False,
):
    """
    JSON API endpoint for programmatic access (n8n, curl, etc.).

    Args:
        image_b64_list: JSON array of base64-encoded images
        prompt: Text instructions
        faces: Target face count
        fmt: Export format (glb/obj/stl)
        modular: Split into parts
    Returns:
        Tuple of (model file path, download file path, status text)
    """
    # Decode images from base64
    images_data = json.loads(image_b64_list) if image_b64_list else []
    tmp_files = []

    for i, b64_str in enumerate(images_data):
        # Handle data URIs
        if "," in b64_str:
            b64_str = b64_str.split(",", 1)[1]
        img_bytes = base64.b64decode(b64_str)
        tmp_path = OUTPUT_DIR / f"api_input_{uuid.uuid4().hex[:6]}.png"
        tmp_path.write_bytes(img_bytes)
        tmp_files.append(str(tmp_path))

    return generate_3d_model(
        input_images=tmp_files,
        text_prompt=prompt,
        num_faces=faces,
        export_format=fmt,
        enable_modular=modular,
    )


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
HEADER_MD = """
# 🚀 AI 3D Model Generator — Print-Ready Pipeline

Generate **high-fidelity, manifold 3D models** from multiple images + text instructions.
Upload photos from different angles, add dimensional requirements, and get a
**watertight mesh** ready for 3D printing or game engines.

**Pipeline**: Background Removal → Multi-View Selection → TripoSG Reconstruction → Manifold Repair → Scale Calibration → Export

> 💡 **Tip**: For best results, provide 2-4 images of the object from different angles on a clean background.
"""

with gr.Blocks(
    title="AI 3D Model Generator",
    theme=gr.themes.Soft(
        primary_hue=gr.themes.colors.indigo,
        secondary_hue=gr.themes.colors.purple,
    ),
    css="""
    .gradio-container { max-width: 1200px !important; }
    .status-box textarea { font-family: monospace; font-size: 12px; }
    """,
) as demo:

    gr.Markdown(HEADER_MD)

    with gr.Row():
        # ----- Left Column: Inputs -----
        with gr.Column(scale=1):
            gr.Markdown("### 📥 Input")

            input_images = gr.File(
                label="Upload Images (multi-angle)",
                file_count="multiple",
                file_types=["image"],
                type="filepath",
            )

            text_prompt = gr.Textbox(
                label="Text Instructions",
                placeholder=(
                    "Describe details, dimensions, materials...\n"
                    "e.g.: Industrial bracket, base is 10cm × 5cm, "
                    "matte steel finish, M6 bolt holes on each corner"
                ),
                lines=4,
            )

            with gr.Row():
                num_faces = gr.Slider(
                    minimum=5000,
                    maximum=MAX_FACES,
                    value=DEFAULT_FACES,
                    step=5000,
                    label="Target Face Count",
                )
                export_format = gr.Dropdown(
                    choices=["glb", "obj", "stl"],
                    value="glb",
                    label="Export Format",
                )

            enable_modular = gr.Checkbox(
                label="🧩 Split into modular parts (experimental)",
                value=False,
            )

            generate_btn = gr.Button(
                "🚀 Generate 3D Model",
                variant="primary",
                size="lg",
            )

        # ----- Right Column: Outputs -----
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Output")

            model_preview = gr.Model3D(
                label="3D Model Preview",
                clear_color=[0.1, 0.1, 0.15, 1.0],
            )

            download_file = gr.File(label="⬇️ Download Model")

            status_output = gr.Textbox(
                label="Pipeline Status",
                lines=10,
                interactive=False,
                elem_classes=["status-box"],
            )

    # ----- Hidden API interface -----
    with gr.Row(visible=False):
        api_image_input = gr.Textbox(label="Base64 Images JSON")
        api_prompt_input = gr.Textbox(label="Prompt")
        api_faces_input = gr.Number(label="Faces", value=DEFAULT_FACES)
        api_fmt_input = gr.Textbox(label="Format", value="glb")
        api_modular_input = gr.Checkbox(label="Modular", value=False)
        api_model_output = gr.File(label="Model")
        api_download_output = gr.File(label="Download")
        api_status_output = gr.Textbox(label="Status")

    # ----- Event handlers -----
    generate_btn.click(
        fn=generate_3d_model,
        inputs=[input_images, text_prompt, num_faces, export_format, enable_modular],
        outputs=[model_preview, download_file, status_output],
        api_name="generate",
    )

    # ----- API-only endpoint for n8n / external callers -----
    # Accepts images as JSON text: URLs, data URIs, or raw base64
    def generate_3d_model_api(
        images_json: str,
        text_prompt: str,
        num_faces: float = DEFAULT_FACES,
        export_format: str = "glb",
        enable_modular: bool = False,
    ):
        """
        API wrapper that accepts images as a JSON array string.
        Each element can be:
          - An HTTP(S) URL (downloaded server-side)
          - A data URI "data:image/png;base64,iVBOR..."
          - A raw base64 string
        """
        import tempfile, base64, json as _json
        import urllib.request

        temp_paths = []
        try:
            images_list = _json.loads(images_json) if images_json else []
        except Exception as e:
            logger.warning(f"Failed to parse images_json: {e}")
            images_list = []

        logger.info(f"API called with {len(images_list)} images, prompt={text_prompt!r}")

        for i, img_src in enumerate(images_list):
            if not img_src:
                continue
            try:
                if img_src.startswith("http://") or img_src.startswith("https://"):
                    # Download image from URL using stdlib
                    logger.info(f"Downloading image {i} from URL: {img_src[:80]}...")
                    req = urllib.request.Request(img_src, headers={"User-Agent": "Mozilla/5.0"})
                    with urllib.request.urlopen(req, timeout=30) as resp:
                        img_bytes = resp.read()
                elif img_src.startswith("data:"):
                    # Data URI: "data:image/png;base64,iVBOR..."
                    b64_part = img_src.split(",", 1)[1]
                    img_bytes = base64.b64decode(b64_part)
                else:
                    # Raw base64
                    img_bytes = base64.b64decode(img_src)

                tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
                tmp.write(img_bytes)
                tmp.close()
                temp_paths.append(tmp.name)
                logger.info(f"Image {i} saved to {tmp.name} ({len(img_bytes)} bytes)")
            except Exception as e:
                logger.error(f"Failed to process image {i}: {e}")

        if not temp_paths and not text_prompt:
            raise gr.Error("No valid images provided and no text prompt. Please provide at least one image URL or base64 string.")

        logger.info(f"Calling generate_3d_model with {len(temp_paths)} images")
        return generate_3d_model(
            input_images=temp_paths if temp_paths else None,
            text_prompt=text_prompt,
            num_faces=int(num_faces),
            export_format=export_format,
            enable_modular=enable_modular,
        )

    generate_btn_api = gr.Button("Generate (API)", visible=False)
    generate_btn_api.click(
        fn=generate_3d_model_api,
        inputs=[api_image_input, api_prompt_input, api_faces_input,
                api_fmt_input, api_modular_input],
        outputs=[api_model_output, api_download_output, api_status_output],
        api_name="generate_api",
    )


# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
demo.queue(max_size=5)
demo.launch(
    show_api=True,  # Expose /api/ endpoints for n8n
)