"""
AI 3D Generation Pipeline — Hugging Face Space (ZeroGPU)
TripoSG + rembg (lighter background removal)
"""
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
    class spaces:
        @staticmethod
        def GPU(fn=None, duration=None):
            if fn is None:
                return lambda f: f
            return fn

# Mesh processing
import trimesh
from rembg import remove  # ← 新增

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_ID = "VAST-AI/TripoSG"
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
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
# ---------------------------------------------------------------------------
triposg_pipeline = None
device = "cuda" if torch.cuda.is_available() else "cpu"

def load_models():
    """Load only TripoSG model."""
    global triposg_pipeline
    if triposg_pipeline is not None:
        return

    logger.info("Loading TripoSG pipeline from %s ...", MODEL_ID)
    try:
        from huggingface_hub import snapshot_download
        from triposg.pipelines.pipeline_triposg import TripoSGPipeline
        
        model_path = snapshot_download(MODEL_ID)
        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

# ---------------------------------------------------------------------------
# Image Processing Utilities
# ---------------------------------------------------------------------------
def remove_background(image: Image.Image) -> Image.Image:
    """Remove background using rembg (lighter & more reliable than RMBG)."""
    try:
        output = remove(image)  # returns RGBA image
        logger.info("Background removed with rembg")
        return output
    except Exception as e:
        logger.warning("rembg failed: %s", e)
        return image  # fallback to original

# ... (其他函數 select_best_view, create_composite_view, extract_dimensions_from_prompt, repair_mesh, scale_mesh, segment_mesh_parts, export_mesh 保持不變) ...

# ---------------------------------------------------------------------------
# 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 (with rembg).
    """
    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]

    # --- 之後步驟同之前一樣 (select_best_view, composite, dimensions, TripoSG inference, post-process, export) ---
    # ... (保持你原有 code 唔變) ...

    elapsed = time.time() - start_time
    log(f"✅ Complete in {elapsed:.1f}s")
    status_text = "\n".join(status_log)
    return preview_path, download_path, status_text   # 你原有變數

# ... (其餘 API 部分同 UI 部分保持不變) ...

# Launch
demo.queue(max_size=5)
demo.launch(show_api=True)