import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from PIL import Image


SPACE_ROOT = Path(__file__).resolve().parent
GEN_MODELS_ROOT = Path("/opt/generative-models")
CHECKPOINTS_DIR = GEN_MODELS_ROOT / "checkpoints"
WORK_DIR = Path("/tmp/sv3d-space")
MODEL_REPO_ID = "stabilityai/sv3d"
MODEL_FILES = {
    "sv3d_u": "sv3d_u.safetensors",
    "sv3d_p": "sv3d_p.safetensors",
}


def ensure_checkpoint(version: str) -> Path:
    token = os.getenv("HF_TOKEN")
    if not token:
        raise gr.Error(
            "Missing HF_TOKEN secret. Add a Hugging Face user access token to the Space "
            "settings, and make sure that token belongs to an account that already accepted "
            "the stabilityai/sv3d license."
        )

    CHECKPOINTS_DIR.mkdir(parents=True, exist_ok=True)
    filename = MODEL_FILES[version]
    target_path = CHECKPOINTS_DIR / filename
    if target_path.exists():
        return target_path

    hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=filename,
        local_dir=CHECKPOINTS_DIR,
        token=token,
    )
    return target_path


def newest_output_file(output_dir: Path) -> Path:
    candidates = sorted(
        output_dir.rglob("*.mp4"),
        key=lambda path: path.stat().st_mtime,
        reverse=True,
    )
    if candidates:
        return candidates[0]

    fallback = sorted(
        output_dir.rglob("*"),
        key=lambda path: path.stat().st_mtime,
        reverse=True,
    )
    for path in fallback:
        if path.is_file():
            return path

    raise gr.Error("The sampler finished without producing an output file.")


def run_sv3d(
    image: Image.Image,
    version: str,
    elevation_deg: float,
    num_steps: int,
    seed: int,
    remove_bg: bool,
    progress=gr.Progress(track_tqdm=False),
):
    if image is None:
        raise gr.Error("Upload an input image first.")
    if not torch.cuda.is_available():
        raise gr.Error("This Space requires GPU hardware. Assign a GPU in Space settings and try again.")

    progress(0.05, desc="Checking model weights")
    ensure_checkpoint(version)
    seed = int(seed)

    job_id = time.strftime("%Y%m%d-%H%M%S")
    job_dir = WORK_DIR / job_id
    input_path = job_dir / "input.png"
    output_dir = job_dir / "outputs"
    output_dir.mkdir(parents=True, exist_ok=True)

    image.convert("RGBA").save(input_path)

    command = [
        sys.executable,
        "scripts/sampling/simple_video_sample.py",
        "--input_path",
        str(input_path),
        "--version",
        version,
        "--output_folder",
        str(output_dir),
        "--num_steps",
        str(num_steps),
        "--seed",
        str(seed),
        "--encoding_t",
        "1",
        "--decoding_t",
        "1",
    ]

    if version == "sv3d_p":
        command.extend(["--elevations_deg", str(elevation_deg)])
    if remove_bg:
        command.extend(["--remove_bg", "True"])

    progress(0.15, desc="Running SV3D")
    result = subprocess.run(
        command,
        cwd=GEN_MODELS_ROOT,
        capture_output=True,
        text=True,
        env={**os.environ, "PYTHONUNBUFFERED": "1"},
    )

    logs = "\n".join(
        part for part in [result.stdout.strip(), result.stderr.strip()] if part
    ).strip()

    if result.returncode != 0:
        raise gr.Error(f"SV3D failed.\n\n{logs[-4000:]}")

    progress(0.95, desc="Collecting output")
    output_path = newest_output_file(output_dir)
    persisted_path = SPACE_ROOT / "latest_output" / output_path.name
    persisted_path.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(output_path, persisted_path)
    video_path = str(persisted_path) if persisted_path.suffix.lower() in {".mp4", ".webm"} else None

    progress(1.0, desc="Done")
    return video_path, str(persisted_path), logs or "Generation completed."


with gr.Blocks(title="SV3D Space") as demo:
    gr.Markdown(
        """
        # Stable Video 3D

        This Space runs `stabilityai/sv3d` with Stability AI's `generative-models` sampler.

        Requirements:
        - GPU hardware must be attached to the Space.
        - Add an `HF_TOKEN` Space secret from an account that has already accepted the
          `stabilityai/sv3d` gated model license.

        Input tips:
        - Best results come from a single centered object on a simple or white background.
        - `sv3d_u` generates an unconstrained orbital video from one image.
        - `sv3d_p` lets you choose the orbit elevation.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(type="pil", label="Input image")
            version = gr.Radio(
                choices=["sv3d_u", "sv3d_p"],
                value="sv3d_u",
                label="Model variant",
            )
            elevation_deg = gr.Slider(
                minimum=-60,
                maximum=60,
                value=10,
                step=1,
                label="Elevation (sv3d_p only)",
            )
            num_steps = gr.Slider(
                minimum=10,
                maximum=50,
                value=30,
                step=1,
                label="Sampling steps",
            )
            seed = gr.Number(value=23, precision=0, label="Seed")
            remove_bg = gr.Checkbox(
                value=False,
                label="Remove background before sampling",
            )
            generate = gr.Button("Generate orbital video", variant="primary")

        with gr.Column(scale=1):
            output_video = gr.Video(label="Result")
            output_file = gr.File(label="Download")
            logs = gr.Textbox(label="Logs", lines=18)

    generate.click(
        fn=run_sv3d,
        inputs=[image, version, elevation_deg, num_steps, seed, remove_bg],
        outputs=[output_video, output_file, logs],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=7860)