"""A simple Gradio Space for LVSPM prefill, camera export, and novel views."""

import os
import shutil
from pathlib import Path
from urllib.parse import unquote

os.environ.setdefault("JAXTYPING_DISABLE", "1")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import spaces  # Import before torch for Hugging Face ZeroGPU.

import gradio as gr
from starlette.middleware import Middleware
from starlette.responses import PlainTextResponse
from demo_core import (EXPORT_ROOT, MAX_VIEWS, ORIENTATIONS, cleanup_scene, load_model, prepare_uploads,
                       prefill_scene, render_scene)

EXAMPLES = [("car", "Historic car"), ("library", "Library"), ("gallery", "Gallery")]


class ArtifactPathGuard:
    """Reject hidden files and parent paths before Gradio resolves downloads."""

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] == "http":
            path = scope.get("path", "")
            for _ in range(4):
                decoded = unquote(path)
                if decoded == path:
                    break
                path = decoded
            if "/file=" in path:
                filename = path.split("/file=", 1)[1]
                if any(part.startswith(".") for part in Path(filename).parts):
                    response = PlainTextResponse("Forbidden", status_code=403)
                    await response(scope, receive, send)
                    return
        await self.app(scope, receive, send)


@spaces.GPU(duration=120)
def prefill(video: str | None, images: list[str] | None, views: int, previous_scene: dict | None,
            orientation: str = "Auto (landscape)"):
    """Predict cameras from a video or ordered images, and cache LVSPM scene state."""
    try:
        arrays, names = prepare_uploads(video, images, int(views), orientation)
        scene, glb, nerfstudio, colmap, status = prefill_scene(arrays, names)
        cleanup_scene(previous_scene)
        return scene, glb, None, gr.update(value=nerfstudio, interactive=True), gr.update(value=colmap, interactive=True), status, gr.update(interactive=True)
    except (ValueError, RuntimeError, OSError) as error:
        raise gr.Error(str(error)) from error


@spaces.GPU(duration=180)
def render(scene: dict, method: str, frames: int):
    """Render a linear or spline path using the existing prefill state."""
    try:
        return render_scene(scene, method, int(frames))
    except (ValueError, RuntimeError, OSError) as error:
        raise gr.Error(str(error)) from error


def reset(scene: dict | None):
    """Discard the current scene when its input changes."""
    cleanup_scene(scene)
    return None, None, None, gr.update(value=None, interactive=False), gr.update(value=None, interactive=False), "Upload one scene, then click Run.", gr.update(interactive=False)


def example_asset(name: str, suffix: str) -> str:
    """Copy a bundled example asset into the dedicated served directory."""
    if name not in {key for key, _ in EXAMPLES} or suffix != "jpg":
        raise ValueError("Unknown example")
    example_dir = EXPORT_ROOT / "examples"
    example_dir.mkdir(exist_ok=True)
    example = example_dir / f"{name}.{suffix}"
    shutil.copyfile(Path(__file__).parent / f"examples/{name}.{suffix}", example)
    return str(example)


def load_example(evt: gr.SelectData):
    """Select a real scene from its visible example card."""
    index = int(evt.index)
    if not 0 <= index < len(EXAMPLES):
        raise gr.Error("Choose one of the example scenes")
    name = EXAMPLES[index][0]
    source = Path(__file__).parent / "examples" / name
    target = EXPORT_ROOT / "examples" / name
    target.mkdir(parents=True, exist_ok=True)
    paths = []
    for image in sorted(source.glob("*.jpg")):
        destination = target / image.name
        shutil.copyfile(image, destination)
        paths.append(str(destination))
    if len(paths) != 64:
        raise gr.Error("Example image collection is incomplete")
    return None, paths, 64, "Auto (landscape)", gr.update(selected="images")


def preview_inputs(video, images, views, orientation):
    """Preview the exact oriented and cropped images used by the model."""
    if not video and not images:
        return [], "Preview appears here after upload."
    try:
        arrays, names = prepare_uploads(video, images, int(views), orientation)
        gallery = [(array, f"{index + 1}/{len(arrays)} · {name}")
                   for index, (array, name) in enumerate(zip(arrays, names))]
        return gr.update(value=gallery, selected_index=0), f"{len(arrays)} input views · click a thumbnail or use the arrows."
    except (ValueError, RuntimeError, OSError) as error:
        return [], str(error)


def refresh_inputs(video, images, views, orientation, scene):
    return (*reset(scene), *preview_inputs(video, images, views, orientation))


def upload_video(video, views, orientation, scene):
    return (None, *refresh_inputs(video, None, views, orientation, scene))


def upload_images(images, views, orientation, scene):
    return (None, *refresh_inputs(None, images, views, orientation, scene))


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="LVSPM") as demo:
        scene = gr.State(None, time_to_live=1800, delete_callback=cleanup_scene)
        gr.HTML('<div class="demo-header"><div><h1>LVSPM: Long Sequence View Synthesis and Pose Estimation Model</h1><p>Turn unposed images collections into cameras and render novel views.</p></div><div class="demo-links"><a href="https://burningdust21.github.io/Projects/LVSPM/" target="_blank" rel="noopener">Project ↗</a><a href="https://huggingface.co/LittleFrog/lvspm" target="_blank" rel="noopener">Weights ↗</a></div></div>')
        with gr.Row(equal_height=False, elem_id="workspace"):
            with gr.Column(scale=4, min_width=330, elem_id="input-panel"):
                with gr.Column(min_width=0, elem_classes="panel-card"):
                    gr.Markdown("### Input")
                    with gr.Tabs() as input_tabs:
                        with gr.Tab("Video", id="video"):
                            video = gr.File(label="Upload a video", file_types=[".mp4", ".mov", ".mkv", ".webm", ".avi"], type="filepath", height=110)
                        with gr.Tab("Images", id="images"):
                            images = gr.File(label="Upload ordered images", file_count="multiple", file_types=["image"], type="filepath", height=110)
                    with gr.Row():
                        views = gr.Slider(2, MAX_VIEWS, value=64, step=1, label="Input views", min_width=140)
                        orientation = gr.Dropdown(ORIENTATIONS, value="Auto (landscape)", label="Orientation", min_width=170)
                    gr.Markdown("Auto rotates portrait inputs clockwise to landscape before cropping. Images follow filename order.", elem_classes="hint")
                with gr.Column(min_width=0, elem_classes="panel-card"):
                    preview = gr.Gallery(label="Input preview", columns=4, rows=1, height=180,
                                         preview=True, selected_index=0, object_fit="contain",
                                         interactive=False, buttons=["fullscreen"], elem_id="input-preview")
                    preview_status = gr.Markdown("Preview appears here after upload.", elem_classes="hint")
                with gr.Column(min_width=0, elem_classes="panel-card"):
                    gr.Markdown("### Examples")
                    examples = gr.Gallery(
                        value=[(example_asset(key, "jpg"), label) for key, label in EXAMPLES],
                        columns=3, rows=1, height=140, show_label=False,
                        object_fit="cover", allow_preview=False, interactive=False,
                        buttons=[], elem_id="example-cards",
                    )
                    gr.Markdown("Click a scene to run it with your render settings.", elem_classes="hint")
            with gr.Column(scale=7, min_width=520, elem_id="results-panel"):
                with gr.Column(min_width=0, elem_classes="panel-card"):
                    with gr.Row():
                        method = gr.Radio(["Linear", "Spline"], value="Spline", label="Camera path", scale=1, min_width=200)
                        frames = gr.Slider(24, 240, value=120, step=24, label="Frames · 24 fps", scale=2, min_width=180)
                        with gr.Column(scale=1, min_width=220):
                            with gr.Row():
                                run_button = gr.Button("Run", variant="primary", min_width=95)
                                render_button = gr.Button("Re-render", interactive=False, min_width=95)
                    status = gr.Markdown("Upload a scene and click Run, or click an example below.", elem_id="run-status")
                with gr.Row(equal_height=True):
                    cameras = gr.Model3D(label="Camera poses", height=320, clear_color=[0.97, 0.98, 0.99, 1.0], min_width=240)
                    novel_view = gr.Video(label="Novel views", height=320, interactive=False, format="mp4", min_width=240)
                gr.Markdown("Drag to orbit the cameras. **Blue:** input views · **Orange:** render path.", elem_classes="hint")
                with gr.Column(min_width=0, elem_id="downloads", elem_classes="panel-card"):
                    gr.Markdown("### Download cameras + images")
                    with gr.Row():
                        nerfstudio = gr.DownloadButton(label="Nerfstudio · OpenGL ↓", interactive=False, size="sm", min_width=180)
                        colmap = gr.DownloadButton(label="COLMAP · OpenCV ↓", interactive=False, size="sm", min_width=180)
                    gr.Markdown("Each ZIP includes the oriented, cropped images as JPEGs.", elem_classes="hint")
                gr.Markdown("**Run** prepares the scene and renders your selected path. Change render settings and **Re-render** to reuse the scene.", elem_classes="hint")
        outputs = [scene, cameras, novel_view, nerfstudio, colmap, status, render_button]
        run_button.click(prefill, [video, images, views, scene, orientation], outputs, api_name="prefill", concurrency_id="gpu", concurrency_limit=1).success(
            render, [scene, method, frames], [cameras, novel_view, status], api_name=False,
            concurrency_id="gpu", concurrency_limit=1,
        )
        render_button.click(render, [scene, method, frames], [cameras, novel_view, status], api_name="render", concurrency_id="gpu", concurrency_limit=1)
        input_outputs = [*outputs, preview, preview_status]
        input_values = [video, images, views, orientation, scene]
        video.upload(upload_video, [video, views, orientation, scene], [images, *input_outputs],
                     api_name=False, concurrency_id="gpu", concurrency_limit=1)
        images.upload(upload_images, [images, views, orientation, scene], [video, *input_outputs],
                      api_name=False, concurrency_id="gpu", concurrency_limit=1)
        examples.select(load_example, outputs=[video, images, views, orientation, input_tabs], api_name=False).success(
            refresh_inputs, input_values, input_outputs, api_name=False,
            concurrency_id="gpu", concurrency_limit=1,
        ).success(
            prefill, [video, images, views, scene, orientation], outputs, api_name=False,
            concurrency_id="gpu", concurrency_limit=1,
        ).success(
            render, [scene, method, frames], [cameras, novel_view, status], api_name=False,
            concurrency_id="gpu", concurrency_limit=1,
        )
        for component in (video, images):
            component.clear(refresh_inputs, input_values, input_outputs, api_name=False,
                            concurrency_id="gpu", concurrency_limit=1)
        for component in (views, orientation):
            component.input(refresh_inputs, input_values, input_outputs, api_name=False,
                             concurrency_id="gpu", concurrency_limit=1)
    return demo


DEMO_CSS = """
.gradio-container {width: 100% !important; box-sizing: border-box !important; max-width: 1480px !important; margin: auto; padding: clamp(14px, 2vw, 24px) !important;}
.gradio-container .main {padding: 0 !important; width: 100% !important; min-width: 0 !important;}
.gradio-container .contain {width: 100% !important; min-width: 0 !important;}
body {background: #f7f8fc; color-scheme: light;}
.gradio-container {color-scheme: light;}
.demo-header {display: flex; flex-direction: column; align-items: center; text-align: center; gap: 10px; padding: 4px 0 14px;}
.demo-header h1 {font-size: clamp(24px, 2.4vw, 32px); line-height: 1.25; letter-spacing: -0.8px; font-weight: 750; margin: 0 auto; max-width: 1100px; color: #172038;}
.demo-header p {margin: 6px 0 0; color: #667085; font-size: 15px;}
.demo-links {display: flex; justify-content: center; gap: 12px; align-items: center; font-size: 13px;}
.demo-links span {background: #eef0ff; color: #5451b9; border-radius: 24px; padding: 7px 12px;}
.demo-links a {color: #5451b9; text-decoration: none; background: #eef0ff; border: 1px solid #e0e2ff; border-radius: 24px; padding: 7px 15px;}
.demo-links a:hover {background: #e3e5ff;}
.hint {font-size: 12px !important; color: #7b8193 !important; line-height: 1.6;}
.hint p {font-size: 12px !important; margin: 0 !important; color: #7b8193 !important;}
.panel-card {background: #fff !important; border: 1px solid #e5e8f0 !important; border-radius: 16px !important; padding: 12px !important; gap: 8px !important; box-shadow: 0 2px 8px #17203804;}
#input-panel > div, #downloads {border-radius: 14px;}
#input-panel h3, #downloads h3 {font-size: 15px; margin: 4px 0 8px;}
#run-status {font-size: 13px; min-height: 26px;}
#results-panel button {white-space: nowrap;}
#downloads {margin-top: 0;}
#downloads h3 {margin: 0 0 4px;}
#downloads button, #downloads a {min-height: 36px;}
@media (max-width: 700px) { .demo-header {gap: 12px;} #workspace {flex-direction: column !important;} #input-panel, #results-panel {min-width: 0 !important; width: 100% !important;} }
"""


def bright_theme():
    """Keep one readable bright palette regardless of browser theme preference."""
    theme = gr.themes.Soft(primary_hue="indigo", neutral_hue="slate").set(
        body_background_fill="#f7f8fc", block_label_background_fill="white",
        block_label_text_color="#475467", block_label_text_size="12px",
        block_label_border_width="0px", block_title_background_fill="white",
        block_title_text_color="#475467", block_radius="12px",
        button_primary_background_fill="#6366f1",
        button_primary_background_fill_hover="#5457df",
    )
    styles = theme.to_dict()["theme"]
    return theme.set(**{key: styles[key[:-5]] for key in styles
                        if key.endswith("_dark") and key[:-5] in styles})


def main():
    # ZeroGPU packs module-level CUDA model weights before GPU calls.
    load_model()
    demo = build_demo().queue(max_size=8)
    kwargs = {"server_name": "0.0.0.0", "server_port": int(os.environ.get("PORT", "18761")),
              "allowed_paths": [str(EXPORT_ROOT)], "max_file_size": "512mb",
              "app_kwargs": {"middleware": [Middleware(ArtifactPathGuard)]},
              "show_error": False, "mcp_server": bool(os.environ.get("SPACE_ID")),
              "theme": bright_theme(),
              "css": DEMO_CSS}
    if os.environ.get("SPACE_ID"):
        kwargs["server_port"] = int(os.environ.get("PORT", "7860"))
    else:
        user, password = os.environ.get("LVSPM_AUTH_USER"), os.environ.get("LVSPM_AUTH_PASSWORD")
        cert, key = os.environ.get("LVSPM_SSL_CERT"), os.environ.get("LVSPM_SSL_KEY")
        if not all((user, password, cert, key)):
            raise RuntimeError("Local network serving requires LVSPM_AUTH_USER, LVSPM_AUTH_PASSWORD, LVSPM_SSL_CERT and LVSPM_SSL_KEY")
        kwargs.update(auth=(user, password), ssl_certfile=cert, ssl_keyfile=key, ssl_verify=False)
    demo.launch(**kwargs)


if __name__ == "__main__":
    main()