from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path

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

from stereo_core import luminance_depth, process_video

VERSION = "0.1.0"
_depth_model = None


def depth_anything_fn(image: Image.Image):
    global _depth_model
    try:
        import torch
        from transformers import pipeline
        if _depth_model is None:
            device = 0 if torch.cuda.is_available() else -1
            _depth_model = pipeline("depth-estimation", model="depth-anything/Depth-Anything-V2-Small-hf", device=device)
        result = _depth_model(image)
        depth = np.asarray(result["depth"], dtype=np.float32)
        depth = (depth - depth.min()) / max(float(depth.max() - depth.min()), 1e-6)
        return 1.0 - depth
    except Exception as exc:
        raise RuntimeError(f"Depth Anything V2 unavailable: {exc}") from exc


def run(video, backend, strength, layout, max_frames):
    if not video:
        raise gr.Error("Upload a video first.")
    work = Path(tempfile.mkdtemp(prefix="simam_mono2stereo_"))
    output = work / "simam_mono2stereo.mp4"
    if backend.startswith("Depth Anything"):
        depth_fn = depth_anything_fn; learned = True
    else:
        depth_fn = luminance_depth; learned = False
    try:
        preview, depth, diagnostics = process_video(video, str(output), depth_fn, float(strength), layout, int(max_frames))
    except Exception as exc:
        raise gr.Error(str(exc))
    depth_img = Image.fromarray(np.uint8(np.clip(depth, 0, 1) * 255), mode="L")
    diagnostics.update({"project": "SimamMono2Stereo", "version": VERSION, "backend": backend, "learned_model_used": learned, "claim_boundary": "pseudo-stereo; unseen content is not verified"})
    return str(output), preview, depth_img, json.dumps(diagnostics, indent=2)


with gr.Blocks(title="SimamMono2Stereo") as demo:
    gr.Markdown("# SimamMono2Stereo\n### Turn ordinary monocular video into an immersive stereo experiment\n\nThis is a transparent DIBR prototype. It creates a left/right pair from estimated depth and reports occlusion holes; it cannot recover unseen content like a real stereo camera.")
    with gr.Row():
        with gr.Column():
            video = gr.Video(label="Monocular input video", sources=["upload", "webcam"], format="mp4")
            backend = gr.Radio(["Fast luminance depth (CPU baseline)", "Depth Anything V2 Small (GPU/model download)"], value="Fast luminance depth (CPU baseline)", label="Depth backend")
            strength = gr.Slider(0.005, 0.08, value=0.025, step=0.005, label="Stereo strength", info="Higher values increase separation and occlusion risk.")
            layout = gr.Radio(["Side-by-side (SBS)", "Over-under (OU)"], value="Side-by-side (SBS)", label="Output layout")
            max_frames = gr.Slider(8, 240, value=96, step=8, label="Maximum processed frames")
            button = gr.Button("Generate immersive stereo", variant="primary")
            gr.Examples([["examples/sample_mono.mp4"]], inputs=[video], label="Try the included synthetic fixture")
        with gr.Column():
            output = gr.Video(label="Stereo video")
            preview = gr.Image(label="First stereo frame")
            depth = gr.Image(label="Estimated depth", image_mode="L")
            diagnostics = gr.Code(label="Diagnostics", language="json")
    button.click(run, [video, backend, strength, layout, max_frames], [output, preview, depth, diagnostics], api_name="generate_stereo")
    gr.Markdown("**Research roadmap:** temporal depth → DIBR → occlusion masks → diffusion inpainting → StereoCrafter adapter → GPU H.265/WebRTC output. Pseudo-stereo quality depends on the input and is not equivalent to captured stereo.")


if __name__ == "__main__":
    demo.launch(server_name=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"), server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")))