import gradio as gr
import os
import time
import requests
import spaces
import sys
import shutil
import tempfile
import torch
import cv2
import subprocess
import numpy as np
import trimesh
import open3d as o3d
from huggingface_hub import hf_hub_download
import html
import base64
import inspect

# ==========================================
# 1. Core Paths and Environment Config
# ==========================================
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
STATIC_DIR = os.path.join(BASE_DIR, "static")
EXAMPLES_DIR = os.path.join(BASE_DIR, "examples")

# ==========================================
# 2. Prepare Static Assets & Environment
# ==========================================
def prepare_local_assets():
    os.makedirs(STATIC_DIR, exist_ok=True)
    base_url = "https://registry.npmmirror.com/three/0.160.0/files"
    assets = {
        "three.module.js": f"{base_url}/build/three.module.js",
        "OrbitControls.js": f"{base_url}/examples/jsm/controls/OrbitControls.js",
        "GLTFLoader.js": f"{base_url}/examples/jsm/loaders/GLTFLoader.js",
        "BufferGeometryUtils.js": f"{base_url}/examples/jsm/utils/BufferGeometryUtils.js"
    }
    
    for name, url in assets.items():
        path = os.path.join(STATIC_DIR, name)
        if not os.path.exists(path):
            try:
                r = requests.get(url, verify=False, timeout=10)
                if r.status_code == 200:
                    with open(path, "wb") as f:
                        f.write(r.content)
            except Exception as e:
                print(f"Error downloading {name}: {e}")

prepare_local_assets()

def install_pytorch3d():
    try:
        import pytorch3d
        print("✅ PyTorch3D already installed.")
        return
    except ImportError:
        print("⏳ PyTorch3D not found. Starting dynamic installation...")

    pyt_version_str = torch.__version__.split("+")[0].replace(".", "")
    version_str = "".join([
        f"py3{sys.version_info.minor}_",
        f"cu{torch.version.cuda.replace('.', '')}_",
        f"pyt{pyt_version_str}"
    ])
    
    whl_url = f"https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/pytorch3d-0.7.8-cp3{sys.version_info.minor}-cp3{sys.version_info.minor}-linux_x86_64.whl"
    
    print(f"🔍 Detected Env: {version_str}")
    print(f"⬇️  Attempting to install from Wheel: {whl_url}")

    try:
        subprocess.run([sys.executable, "-m", "pip", "install", whl_url], check=True)
        print("✅ PyTorch3D installed via Wheel!")
    except subprocess.CalledProcessError:
        print("⚠️  Wheel installation failed. Falling back to source compilation...")
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "--no-build-isolation", "git+https://github.com/facebookresearch/pytorch3d.git@stable"], 
            check=True
        )
        print("✅ PyTorch3D installed via Source Build!")

install_pytorch3d()

def install_mmcv():
    try:
        import mmcv
        print(f"✅ mmcv {mmcv.__version__} is already installed.")
        return
    except ImportError:
        print("⏳ mmcv not found. Starting dynamic installation...")
        cuda_ver = torch.version.cuda.replace(".", "")
        torch_ver = ".".join(torch.__version__.split(".")[:2])
        find_links_url = f"https://download.openmmlab.com/mmcv/dist/cu{cuda_ver}/torch{torch_ver}/index.html"
        
        print(f"🔍 Detected Env: CUDA={cuda_ver}, Torch={torch_ver}")
        print(f"⬇️  Installing mmcv==2.2.0 from: {find_links_url}")

        try:
            subprocess.run([
                sys.executable, "-m", "pip", "install", 
                "mmcv==2.2.0", 
                "--find-links", find_links_url
            ], check=True)
            print("✅ mmcv installed successfully.")
        except subprocess.CalledProcessError:
            print("⚠️  Installation failed. Attempting fallback using openmim...")
            subprocess.run([sys.executable, "-m", "pip", "install", "openmim"], check=True)
            subprocess.run(["mim", "install", "mmcv==2.2.0"], check=True)

install_mmcv()

def install_sam2():
    try:
        import sam2
    except ImportError:
        print("Installing SAM 2 with patch...")
        subprocess.run(["git", "clone", "https://github.com/facebookresearch/segment-anything-2.git", "_tmp_sam2"], check=True)
        setup_path = "_tmp_sam2/setup.py"
        with open(setup_path, "r") as f:
            content = f.read()
        content = content.replace("torch>=2.5.1", "torch>=2.4.1")
        with open(setup_path, "w") as f:
            f.write(content)
        subprocess.run(["pip", "install", "--no-build-isolation", "--no-deps", "-v", "."], cwd="_tmp_sam2", check=True)
        shutil.rmtree("_tmp_sam2")

install_sam2()

def install_smpl_dependencies():
    try:
        import chumpy
        print("✅ Chumpy already installed.")
    except ImportError:
        print("⏳ Installing Chumpy...")
        subprocess.run([sys.executable, "-m", "pip", "install", "chumpy", "--no-build-isolation"], check=True)

    try:
        import smplx
        print("✅ smplx already installed.")
    except ImportError:
        print("⏳ Installing smplx...")
        subprocess.run([sys.executable, "-m", "pip", "install", "smplx"], check=True)

install_smpl_dependencies()

sys.path.append(BASE_DIR)

# Import core model code
from unish.utils.inference_utils import (
    load_model, process_video, run_inference,
    generate_mixed_geometries_in_memory,
    save_smpl_meshes_per_frame
)

MODEL = None
BODY_MODELS_PATH = "body_models/"

# ==========================================
# 4. Helper Functions
# ==========================================
def download_smpl_assets(body_models_path):
    if 'smpl' not in body_models_path:
        model_path = os.path.join(body_models_path, 'smpl')
    else:
        model_path = body_models_path
    
    target_dir = os.path.join(model_path, 'smpl')
    os.makedirs(target_dir, exist_ok=True)
    
    files = ["SMPL_NEUTRAL.pkl", "SMPL_MALE.pkl", "SMPL_FEMALE.pkl"]
    repo_id = "Murphyyyy/UniSH-Private-Assets"
    
    token = os.environ.get("SMPL_DOWNLOAD_TOKEN")
    if not token:
        print("❌ CRITICAL ERROR: 'SMPL_DOWNLOAD_TOKEN' not found in environment variables!")
        print("👉 Since 'UniSH-Private-Assets' is likely private, inference WILL fail without a token.")
    
    for filename in files:
        file_path = os.path.join(target_dir, filename)
        if not os.path.exists(file_path):
            try:
                print(f"📥 Downloading {filename} from {repo_id}...")
                hf_hub_download(
                    repo_id=repo_id, 
                    filename=filename,
                    token=token, 
                    local_dir=target_dir, 
                    local_dir_use_symlinks=False
                )
                print(f"✅ Downloaded to: {file_path}")
            except Exception as e:
                print(f"❌ Failed to download {filename}: {e}")
                print(f"   (Check if your HF Token has access to {repo_id})")

def pack_sequence_to_glb(base_dir, output_path, start_frame, end_frame, scene_rate=1.0):
    scene = trimesh.Scene()
    
    scene_cloud_dir = os.path.join(base_dir, "scene_clouds_per_frame")
    smpl_mesh_dir = os.path.join(base_dir, "smpl_meshes_per_frame")
    
    # 60,000 points per frame
    MAX_POINTS_PER_FRAME = 60000 

    for i in range(start_frame, end_frame):
        candidates = [
            os.path.join(smpl_mesh_dir, f"combined_smpl_mesh_frame_{i:04d}.ply"),
            os.path.join(smpl_mesh_dir, f"smpl_mesh_frame_{i:04d}.ply")
        ]
        target_human_path = None
        for p in candidates:
            if os.path.exists(p):
                target_human_path = p
                break
            
        if target_human_path:
            try:
                human_mesh = trimesh.load(target_human_path)
                node_name = f"frame_{i}_human"
                scene.add_geometry(human_mesh, node_name=node_name, geom_name=node_name)
            except Exception:
                pass

        scene_pcd_path = os.path.join(scene_cloud_dir, f"scene_frame_{i:04d}.ply")
        if os.path.exists(scene_pcd_path):
            try:
                scene_pc = trimesh.load(scene_pcd_path)
                if hasattr(scene_pc, 'vertices') and len(scene_pc.vertices) > 0:
                    num_points = len(scene_pc.vertices)
                    if num_points > MAX_POINTS_PER_FRAME:
                        choice = np.random.choice(num_points, MAX_POINTS_PER_FRAME, replace=False)
                        scene_pc.vertices = scene_pc.vertices[choice]
                        if hasattr(scene_pc, 'colors') and len(scene_pc.colors) > 0:
                            scene_pc.colors = scene_pc.colors[choice]
                    node_name = f"frame_{i}_scene"
                    scene.add_geometry(scene_pc, node_name=node_name, geom_name=node_name)
            except Exception:
                pass

    if len(scene.geometry) == 0:
        dummy = trimesh.creation.box(extents=[0.01, 0.01, 0.01])
        scene.add_geometry(dummy, node_name='dummy')

    scene.export(output_path)
    if not os.path.exists(output_path):
         raise FileNotFoundError(f"Export failed: {output_path}")

def get_video_duration(video_path):
    if not video_path: return 10.0
    try:
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened(): return 10.0
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        duration = frame_count / fps if fps > 0 else 10.0
        cap.release()
        return duration
    except:
        return 10.0

# ==========================================
# 5. Frontend Logic: Player & Loading HTML
# ==========================================

def get_loading_html(message="Processing..."):
    # Updated to use CSS variables for Dark Mode compatibility
    return f"""
    <div style="height: 600px; width: 100%; background: var(--background-fill-secondary); border-radius: 12px; border: 1px solid var(--border-color-primary); display: flex; flex-direction: column; align-items: center; justify-content: center; font-family: sans-serif; color: var(--body-text-color);">
        <div class="loader-ring"></div>
        <p style="margin-top: 20px; font-weight: 500; font-size: 1.1em; animation: pulse 2s infinite;">{message}</p>
        <style>
            .loader-ring {{ display: inline-block; width: 64px; height: 64px; }}
            .loader-ring:after {{ content: " "; display: block; width: 46px; height: 46px; margin: 8px; border-radius: 50%; border: 5px solid #4f46e5; border-color: #4f46e5 transparent #4f46e5 transparent; animation: ring-spin 1.2s linear infinite; }}
            @keyframes ring-spin {{ 0% {{ transform: rotate(0deg); }} 100% {{ transform: rotate(360deg); }} }}
            @keyframes pulse {{ 0% {{ opacity: 0.6; }} 50% {{ opacity: 1; }} 100% {{ opacity: 0.6; }} }}
        </style>
    </div>
    """

def get_player_html(abs_glb_path):
    def read_and_patch(filename):
        path = os.path.join(STATIC_DIR, filename)
        if not os.path.exists(path): return ""
        with open(path, "r", encoding="utf-8") as f: content = f.read()
        content = content.replace('../utils/BufferGeometryUtils.js', 'BufferGeometryUtils')
        content = content.replace('./BufferGeometryUtils.js', 'BufferGeometryUtils')
        content = content.replace('../../build/three.module.js', 'three')
        content = content.replace('../build/three.module.js', 'three')
        return content

    js_three = read_and_patch("three.module.js")
    js_orbit = read_and_patch("OrbitControls.js")
    js_loader = read_and_patch("GLTFLoader.js")
    js_buffer = read_and_patch("BufferGeometryUtils.js")

    def to_data_url(content, mime="text/javascript"):
        b64 = base64.b64encode(content.encode('utf-8')).decode('utf-8')
        return f"data:{mime};base64,{b64}"

    blob_three = to_data_url(js_three)
    blob_orbit = to_data_url(js_orbit)
    blob_loader = to_data_url(js_loader)
    blob_buffer = to_data_url(js_buffer)
    
    if not os.path.exists(abs_glb_path): return '<div style="color:var(--error-text-color); padding:20px;">Error: Output file not found.</div>'
    with open(abs_glb_path, "rb") as f: glb_data = f.read()
    model_data_url = f"data:model/gltf-binary;base64,{base64.b64encode(glb_data).decode('utf-8')}"

    # Note: 3D scene background is kept light gray (0xf9fafb) for model visibility
    raw_html = f"""
    <!DOCTYPE html>
    <html>
    <head>
      <meta charset="utf-8">
      <style>
        body {{ margin: 0; background: transparent; height: 100vh; overflow: hidden; font-family: sans-serif; }}
        #container {{ width: 100%; height: 100%; }}
        .controls {{ position: absolute; bottom: 20px; left: 50%; transform: translateX(-50%); display: flex; gap: 12px; padding: 12px 20px; background: rgba(255, 255, 255, 0.9); border-radius: 30px; z-index: 100; box-shadow: 0 4px 15px rgba(0,0,0,0.15); backdrop-filter: blur(5px); align-items: center; }}
        button {{ padding: 8px 20px; cursor: pointer; background: #4f46e5; color: white; border: none; border-radius: 20px; font-weight: 600; font-size: 14px; transition: background 0.2s; }}
        button:hover {{ background: #4338ca; }}
        input[type=range] {{ width: 200px; cursor: pointer; accent-color: #4f46e5; }}
      </style>
      <script type="importmap">
      {{ "imports": {{ "three": "{blob_three}", "three/addons/controls/OrbitControls.js": "{blob_orbit}", "three/addons/loaders/GLTFLoader.js": "{blob_loader}", "BufferGeometryUtils": "{blob_buffer}" }} }}
      </script>
    </head>
    <body>
      <div id="container"></div>
      <div class="controls"><button id="btn-play">Play</button><input type="range" id="slider" min="0" max="0" value="0" step="1"></div>
      <script type="module">
        import * as THREE from 'three';
        import {{ OrbitControls }} from 'three/addons/controls/OrbitControls.js';
        import {{ GLTFLoader }} from 'three/addons/loaders/GLTFLoader.js';
        const container = document.getElementById('container');
        const scene = new THREE.Scene(); scene.background = new THREE.Color(0xf9fafb);
        
        const camera = new THREE.PerspectiveCamera(45, window.innerWidth / window.innerHeight, 0.1, 100); 
        camera.position.set(0, 1.5, 4); 
        
        const renderer = new THREE.WebGLRenderer({{ antialias: true, alpha: true }}); renderer.setSize(window.innerWidth, window.innerHeight); renderer.setPixelRatio(window.devicePixelRatio); container.appendChild(renderer.domElement);
        const controls = new OrbitControls(camera, renderer.domElement); controls.enableDamping = true;
        scene.add(new THREE.AmbientLight(0xffffff, 0.8)); const dirLight = new THREE.DirectionalLight(0xffffff, 1.2); dirLight.position.set(5, 10, 7); scene.add(dirLight);
        let frames = []; let currentFrame = 0; let isPlaying = false; let timer = null;
        new GLTFLoader().load("{model_data_url}", (gltf) => {{
            gltf.scene.rotation.x = Math.PI; scene.add(gltf.scene);
            gltf.scene.traverse(n => {{
                if(n.name && n.name.startsWith('frame_')) {{
                    let parts = n.name.split('_'); let idx = parseInt(parts[1]);
                    if(!isNaN(idx)) {{ if(!frames[idx]) frames[idx] = []; frames[idx].push(n); n.visible = false; }}
                }}
                if(n.isMesh) {{ n.geometry.computeVertexNormals(); n.material = new THREE.MeshStandardMaterial({{ color: 0x6366f1, roughness: 0.4, metalness: 0.1 }}); }}
                if(n.isPoints) {{ 
                    let size = n.name.includes('scene') ? 0.05 : 0.005; 
                    n.material.size = size; 
                }}
            }});
            if(frames.length > 0) {{ document.getElementById('slider').max = frames.length - 1; showFrame(0); animate(); }}
        }}, undefined, (e) => console.error(e));
        function showFrame(idx) {{
            if(frames[currentFrame]) frames[currentFrame].forEach(o => o.visible = false);
            if(frames[idx]) frames[idx].forEach(o => o.visible = true);
            currentFrame = idx;
            const slider = document.getElementById('slider');
            if(slider) slider.value = idx;
        }}
        function animate() {{ requestAnimationFrame(animate); controls.update(); renderer.render(scene, camera); }}
        document.getElementById('btn-play').onclick = () => {{
            isPlaying = !isPlaying; const btn = document.getElementById('btn-play');
            if(isPlaying) {{
                btn.innerText = "Pause"; btn.style.background = "#ef4444";
                timer = setInterval(() => {{ if(frames.length > 0) {{ let next = (currentFrame + 1) % frames.length; showFrame(next); }} }}, 166);
            }} else {{ clearInterval(timer); btn.innerText = "Play"; btn.style.background = "#4f46e5"; }}
        }};
        document.getElementById('slider').oninput = (e) => {{ if(isPlaying) document.getElementById('btn-play').click(); showFrame(parseInt(e.target.value)); }};
        window.onresize = () => {{ camera.aspect = window.innerWidth / window.innerHeight; camera.updateProjectionMatrix(); renderer.setSize(window.innerWidth, window.innerHeight); }};
      </script>
    </body>
    </html>
    """
    return f'<iframe srcdoc="{html.escape(raw_html)}" width="100%" height="600px" style="border:none; border-radius: 12px;"></iframe>'

# ==========================================
# 6. Inference Logic Wrapper
# ==========================================
if os.environ.get("SPACE_ID"):
    from spaces import GPU
    gpu_decorator = GPU(duration=120)
else:
    def gpu_decorator(func):
        def wrapper(*args, **kwargs):
            if inspect.isgeneratorfunction(func):
                yield from func(*args, **kwargs)
            else:
                return func(*args, **kwargs)
        return wrapper

@gpu_decorator
def predict(video_path, start_time=0.0, end_time=10.0):
    # 1. Limit Check
    duration_input = end_time - start_time
    if duration_input > 10.0:
        raise gr.Error(f"Video limit exceeded ({duration_input:.1f}s). Please keep it under 10 seconds.")

    if start_time >= end_time:
        raise gr.Error("Error: End time must be greater than Start time.")
    
    # 2. Normal Flow
    yield get_loading_html("Processing...")
    
    output_dir = tempfile.mkdtemp()
    trimmed_video_path = os.path.join(output_dir, "input_trim.mp4")
    duration = end_time - start_time
    
    cmd = ["ffmpeg", "-y", "-ss", str(start_time), "-i", video_path, "-t", str(duration), "-c:v", "libx264", "-c:a", "aac", trimmed_video_path]
    subprocess.run(cmd, check=True)
    
    global MODEL
    if MODEL is None: MODEL = load_model()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    MODEL.to(device)
    MODEL.eval()
    
    download_smpl_assets(BODY_MODELS_PATH)
    
    data_dict = process_video(trimmed_video_path, 6.0, 0, 518, bbox_scale=1.0)
    results = run_inference(MODEL, data_dict, device, chunk_size=300)
    
    seq_name = results['seq_name']
    viz_scene, viz_smpl, viz_scene_only, _ = generate_mixed_geometries_in_memory(
        results, BODY_MODELS_PATH, fps=6.0, conf_thres=0.1
    )
    
    save_smpl_meshes_per_frame(results, output_dir, BODY_MODELS_PATH)
    
    base_dir = os.path.join(output_dir, seq_name)
    scene_cloud_dir = os.path.join(base_dir, "scene_clouds_per_frame")
    os.makedirs(scene_cloud_dir, exist_ok=True)
    for i, pcd in enumerate(viz_scene_only):
        if len(pcd.points) > 0:
            o3d.io.write_point_cloud(os.path.join(scene_cloud_dir, f"scene_frame_{i:04d}.ply"), pcd)
    
    # Use temp file directly for thread safety
    tmp_glb_path = os.path.join(output_dir, "output.glb")
    pack_sequence_to_glb(base_dir, tmp_glb_path, 0, len(viz_scene), 0.5)
        
    yield get_player_html(tmp_glb_path)

# ==========================================
# 7. UI Construction
# ==========================================

examples_list = []
if os.path.exists(EXAMPLES_DIR):
    examples_list = [[os.path.join("examples", f)] for f in os.listdir(EXAMPLES_DIR) if f.endswith(".mp4")]

js_reset_video = """(args) => {
    setTimeout(() => {
        var video = document.querySelector('#input-video video');
        if (video) {
            video.currentTime = 0;
            video.pause();
        }
    }, 100);
    return args; 
}"""

js_scrub = """(val) => {
    var video = document.querySelector('#input-video video');
    if (video) {
        video.currentTime = val;
    }
    return val;
}"""

custom_css = """
footer {visibility: hidden}
h1.header-title { text-align: center; font-family: 'Segoe UI', sans-serif; font-weight: 700; color: var(--body-text-color); margin-bottom: 2rem; }

.alert-box { padding: 1rem; border-radius: 0.5rem; margin-bottom: 1rem; font-size: 0.9rem; line-height: 1.5; }

.viewer-container { box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1); border-radius: 12px; overflow: hidden; border: 1px solid var(--border-color-primary); background: var(--background-fill-secondary); }
#run-btn { background: linear-gradient(135deg, #6366f1 0%, #a855f7 100%); border: none; color: white; font-weight: bold; transition: all 0.2s; }
#run-btn:hover { transform: translateY(-2px); box-shadow: 0 10px 20px rgba(99, 102, 241, 0.4); }
#input-video { margin-bottom: 15px !important; }
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", radius_size="md"), css=custom_css, title="UniSH Demo") as demo:
    
    with gr.Column(elem_classes=["header-container"]):
        gr.Markdown("# UniSH: Unifying Scene and Human Reconstruction in a Feed-Forward Pass", elem_classes=["header-title"])

    with gr.Row(equal_height=False):
        
        with gr.Column(scale=3, variant="panel"):
            gr.Markdown("### 🛠️ Configuration")
            
            with gr.Group():
                input_video = gr.Video(label="Upload Video", format="mp4", height=260, interactive=True, elem_id="input-video")
                
                with gr.Row():
                    start_time = gr.Slider(minimum=0, maximum=10, value=0, step=0.01, label="Start Time (s)")
                    end_time = gr.Slider(minimum=0, maximum=10, value=3, step=0.01, label="End Time (s)")
            
            gr.HTML("""
            <div style="background-color: #f0fdf4; border-left: 4px solid #22c55e; padding: 1rem; border-radius: 0.5rem; margin-bottom: 1rem; color: #15803d !important;">
                <div style="font-weight: bold; margin-bottom: 8px; color: #15803d !important;">💡 Use Tips:</div>
                <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 10px; font-size: 0.9em; line-height: 1.4; color: #15803d !important;">
                    <div style="color: #15803d !important;">1. Contain only a <strong style="color: #15803d !important;">single person</strong>.</div>
                    <div style="color: #15803d !important;">2. <strong style="color: #15803d !important;">No occlusion</strong> (self-occlusion is fine).</div>
                    <div style="color: #15803d !important;">3. Keep the <strong style="color: #15803d !important;">full body</strong> mostly visible.</div>
                    <div style="color: #15803d !important;">4. Subject should <strong style="color: #15803d !important;">not be too small</strong>.</div>
                </div>
            </div>
            """)
            
            submit_btn = gr.Button("🚀 Start Reconstruction", variant="primary", elem_id="run-btn", size="lg")
            
            if examples_list:
                gr.Markdown("### 🎥 Examples")
                gr.Examples(
                    examples=examples_list,
                    inputs=[input_video], 
                    label="Click to try:",
                    cache_examples=False,
                    run_on_click=False 
                )

        with gr.Column(scale=7):
            gr.Markdown("### ▶️ Interactive Results")
            with gr.Group(elem_classes=["viewer-container"]):
                output_html = gr.HTML(
                    label="3D Viewer", 
                    min_height=600,
                    value='<div style="height:600px; display:flex; align-items:center; justify-content:center; color:var(--body-text-color); font-family:sans-serif; background:var(--background-fill-secondary);">Upload a video and click Start to view result.</div>'
                )
            
            gr.HTML("""
            <div style="background-color: #fefce8; border-left: 4px solid #eab308; padding: 1rem; border-radius: 0.5rem; margin-bottom: 1rem; color: #854d0e !important;">
                <strong style="color: #854d0e !important;">⚡ Performance Note:</strong><br>
                <span style="color: #854d0e !important;">Inference (feed-forward) is very fast, but generating visualization assets (meshing & rendering) takes up most of the processing time.</span>
            </div>
            
            <div style="background-color: #eff6ff; border-left: 4px solid #3b82f6; padding: 1rem; border-radius: 0.5rem; margin-bottom: 1rem; color: #1e40af !important;">
                <strong style="color: #1e40af !important;">👁️ Visual Quality:</strong><br>
                <span style="color: #1e40af !important;">The displayed results are downsampled for better web rendering performance.</span>
            </div>
            """)

    def update_slider_range(video_path, *args):
        if not video_path:
            return gr.update(value=0), gr.update(value=0)
        dur = get_video_duration(video_path)
        dur = round(dur, 2)
        return gr.update(maximum=dur, value=0), gr.update(maximum=dur, value=dur)

    input_video.upload(fn=None, js="() => { setTimeout(() => { var v = document.querySelector('#input-video video'); if(v) {v.currentTime=0; v.pause();} }, 100); }")
    
    input_video.change(
        fn=update_slider_range, 
        inputs=[input_video], 
        outputs=[start_time, end_time]
    )

    start_time.change(fn=None, inputs=[start_time], outputs=None, js=js_scrub)
    end_time.change(fn=None, inputs=[end_time], outputs=None, js=js_scrub)

    submit_btn.click(
        fn=predict, 
        inputs=[input_video, start_time, end_time], 
        outputs=[output_html]
    )

demo.queue()
demo.launch(
    server_name="0.0.0.0", 
    server_port=7860, 
    allowed_paths=[BASE_DIR, "/tmp", EXAMPLES_DIR]
)