import gradio as gr
import os
import urllib.parse # 新增：用於處理網址編碼

# 獲取當前絕對路徑
BASE_DIR = os.path.abspath(os.getcwd())

SCENES = [
    {
        "name": "Office 0", 
        "thumb": "assets/office0.PNG",            
        "model": "assets/office0.ply",
        "tracking_img": "assets/office0_tracking.PNG" 
    },
    {
        "name": "Office 2",
        "thumb": "assets/office2.PNG",
        "model": "assets/office2.ply",
        "tracking_img": "assets/office2_tracking.PNG"
    },
    {
        "name": "Room 0",
        "thumb": "assets/room0.PNG",
        "model": "assets/room0.ply",
        "tracking_img": "assets/room0_tracking.PNG"
    }
]

gallery_items = [(item["thumb"], item["name"]) for item in SCENES]

import urllib.parse # 確保有這一行

def load_scene(evt: gr.SelectData, request: gr.Request):
    index = evt.index
    scene = SCENES[index]
    
    # 取得 Hugging Face Space 的 ID (例如: yfyangd/rgs-slam-demo)
    space_id = os.environ.get("SPACE_ID")
    
    # --- 關鍵修正邏輯 ---
    
    if space_id:
        # 【情境 A：在 Hugging Face Spaces 線上執行】
        # 我們改用 Hugging Face 的 "resolve" 網址。
        # 這是直接指向 Git 倉庫的原始檔案，速度快且通常不會被 CORS 擋住。
        # 格式: https://huggingface.co/spaces/{SPACE_ID}/resolve/main/{相對路徑}
        
        # 注意：這裡直接用 scene["model"] (例如 assets/office0.ply)，不需要絕對路徑
        file_url = f"https://huggingface.co/spaces/{space_id}/resolve/main/{scene['model']}"
        
        # 有些 viewer 需要加上 download=true 才能正確觸發 header，但 antimatter 通常不用
        print(f"🌍 Using HF Repo URL: {file_url}")
        
    else:
        # 【情境 B：在 本機 (Localhost) 執行】
        # 保持原本的邏輯，使用 localhost 的連結
        full_model_path = os.path.join(BASE_DIR, scene["model"])
        
        if not os.path.exists(full_model_path):
            return (
                "<div style='text-align:center; padding:50px; color:red'>⚠️ File not found</div>",
                None,
                f"**⚠️ Error:** File not found: {full_model_path}"
            )

        # 處理 Windows 路徑
        clean_path = full_model_path.replace("\\", "/")
        
        base_url = "http://127.0.0.1:7860"
        if request:
             host = request.headers.get("host")
             if "gradio.live" in host:
                 base_url = f"https://{host}"
             else:
                 base_url = f"http://{host}"
        
        file_url = f"{base_url}/file={clean_path}"
        print(f"🏠 Using Local URL: {file_url}")

    # --- 通用部分：編碼並生成 iframe ---
    
    # 編碼 file_url (處理特殊字元)
    encoded_file_url = urllib.parse.quote(file_url, safe=':/')
    
    # 組合外部 Viewer 連結
    viewer_url = f"https://antimatter15.com/splat/?url={encoded_file_url}"
    
    html = f"""
    <iframe 
        src="{viewer_url}" 
        width="100%" 
        height="600px" 
        frameborder="0"
        allow="camera; display-capture"
        style="border-radius: 8px; background: #000;"
    ></iframe>
    """
    
    status_msg = f"**{scene['name']}** - Loaded via {'HF CDN' if space_id else 'Local Server'}"

    return (
        html,
        scene["tracking_img"],
        status_msg
    )

# --- UI (保持不變) ---
with gr.Blocks(title="RGS-SLAM Demo") as demo:
    gr.Markdown("# 🚀 RGS-SLAM Gaussian Splatting Viewer")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📂 Select Scene")
            scene_gallery = gr.Gallery(
                value=gallery_items, 
                label="Available Scenes",
                columns=1,
                height=400,
                object_fit="contain",
                allow_preview=False
            )
        with gr.Column(scale=3):
            status = gr.Markdown("**👈 Select a scene to auto-load**")
            with gr.Tab("🎯 3D Viewer"):
                viewer_frame = gr.HTML(
                    value="<div style='text-align: center; padding: 100px; color: #666;'>Select a scene</div>"
                )
            with gr.Tab("📍 Camera Trajectory"):
                trajectory_img = gr.Image(label="Camera Path", height=600)
    
    scene_gallery.select(
        fn=load_scene,
        inputs=None,
        outputs=[viewer_frame, trajectory_img, status]
    )

if __name__ == "__main__":
    # 判斷是否在 Hugging Face Spaces 上執行
    # HF Spaces 會有名為 "SPACE_ID" 的環境變數
    is_on_hf_spaces = "SPACE_ID" in os.environ

    demo.launch(
        share=not is_on_hf_spaces, # 如果在 HF Spaces 就 False，本地就 True
        allowed_paths=[BASE_DIR]
    )