import os

# ZeroGPU must patch Torch/CUDA before any CUDA-related package is imported.
os.environ["TORCH_COMPILE_DISABLE"] = "1"
os.environ["TORCHDYNAMO_DISABLE"] = "1"
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces


@spaces.GPU
def _zerogpu_registration_sentinel():
    """ZeroGPU registration sentinel. Intentionally never called."""
    return None


from ltxc.assets import ASSETS, REMOTE_EXAMPLES
from ltxc.model_runtime import initialize_global_colorizer

# Load and place the fixed pipeline on CUDA before any Gradio UI is constructed.
_STARTUP_COLORIZER = initialize_global_colorizer(ASSETS)

import gradio as gr

from ltxc.config import (
    EXAMPLE_ROOT,
    FRAME_CHOICES,
    IS_ZEROGPU,
    JOB_ROOT,
    MAX_SEED,
    RABBIT_PROMPT,
    RES_PRESETS,
)
from ltxc.duration import (
    DEFAULT_CALIBRATION_MULTIPLIER,
    DEFAULT_MANUAL_SECONDS,
    DEFAULT_SAFE_MODE_ENABLED,
    DURATION_MODE_SEMI_AUTO,
    DURATION_MODES,
    format_duration_panel,
)
from ltxc.generation import execute_generation
from ltxc.long_runtime import execute_long_generation
from ltxc.long_video import (
    LONG_CHUNK_FRAMES,
    LONG_MAX_CHUNKS_PER_CALLBACK,
    LONG_MAX_SOURCE_FRAMES,
    LONG_MAX_TOTAL_CHUNKS,
    LONG_MIN_SOURCE_FRAMES,
    LONG_OVERLAP_FRAMES,
    LONG_STRIDE_FRAMES,
)
from ltxc.preparation import (
    estimate_prepared_gpu_seconds,
    prepare_long_product,
    prepare_product,
)


def _gpu_entrypoint(fn):
    """Apply the real ZeroGPU allocation wrapper only in ZeroGPU runtime."""
    if IS_ZEROGPU:
        return spaces.GPU(size="large", duration=estimate_prepared_gpu_seconds)(fn)
    return fn


@_gpu_entrypoint
def _run_long_product(job_id: str, progress=gr.Progress(track_tqdm=True)):
    if not str(job_id or "").strip():
        raise gr.Error("Prepare a long-video job first.")
    outcome = execute_long_generation(job_id, make_bundle=True, progress=progress)
    if not outcome["ok"]:
        return (
            outcome["output_path"],
            outcome["bundle_path"],
            outcome["summary"],
            outcome["status"] + " You may use Continue / resume prepared job.",
        )
    return (
        outcome["output_path"],
        outcome["bundle_path"],
        outcome["summary"],
        outcome["status"],
    )


@_gpu_entrypoint
def _run_product(job_id: str, progress=gr.Progress(track_tqdm=True)):
    outcome = execute_generation(job_id, progress=progress)
    if not outcome["ok"]:
        raise gr.Error(outcome["status"])
    return outcome["output_path"], outcome["seed"], outcome["status"]


def _build_colorize_ui():
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Input video (any clip — recolored as B&W)")
            prompt = gr.Textbox(
                label="Prompt — describe the colorized scene, plus any sounds",
                lines=3,
                placeholder=RABBIT_PROMPT,
            )
            with gr.Accordion("Settings", open=False):
                preset = gr.Dropdown(
                    list(RES_PRESETS),
                    value="960×544 (recommended)",
                    label="Resolution",
                )
                num_frames = gr.Dropdown(
                    FRAME_CHOICES,
                    value=121,
                    label="Frames (24fps)",
                )
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
                with gr.Accordion("ZeroGPU duration", open=False):
                    duration_mode = gr.Dropdown(
                        DURATION_MODES,
                        value=DURATION_MODE_SEMI_AUTO,
                        label="Duration mode",
                    )
                    manual_gpu_seconds = gr.Slider(
                        30,
                        300,
                        value=DEFAULT_MANUAL_SECONDS,
                        step=1,
                        label="Manual duration (seconds)",
                    )
                    calibration_multiplier = gr.Number(
                        value=DEFAULT_CALIBRATION_MULTIPLIER,
                        minimum=0.5,
                        maximum=2.0,
                        step=0.05,
                        label="Semi-auto calibration multiplier",
                    )
                    safe_mode = gr.Checkbox(
                        DEFAULT_SAFE_MODE_ENABLED,
                        label="Safe mode (+30% to Semi-auto)",
                    )
                    duration_estimate = gr.Markdown()
            run = gr.Button("Colorize", variant="primary")
            status = gr.Markdown(
                "Ready. Progress appears above while the queued GPU task runs."
            )
            job_state = gr.State("")
        with gr.Column():
            video_out = gr.Video(label="Colorized result")
            gr.Markdown(
                "`960×544 / 121` is the accepted stable quality baseline. "
                "Use the Long video tab for batched multi-window processing."
            )

    preparation = run.click(
        prepare_product,
        inputs=[
            video_in,
            prompt,
            preset,
            num_frames,
            seed,
            randomize,
            duration_mode,
            manual_gpu_seconds,
            calibration_multiplier,
            safe_mode,
        ],
        outputs=[job_state, status],
        queue=True,
    )
    preparation.success(
        _run_product,
        inputs=[job_state],
        outputs=[video_out, seed, status],
        concurrency_id="ltx23-gpu",
        concurrency_limit=1,
    )

    gr.Examples(
        examples=REMOTE_EXAMPLES,
        inputs=[video_in, prompt, preset, num_frames, seed, randomize],
        cache_examples=False,
    )

    duration_inputs = [
        preset,
        num_frames,
        duration_mode,
        manual_gpu_seconds,
        calibration_multiplier,
        safe_mode,
    ]
    for component in duration_inputs:
        component.change(
            format_duration_panel,
            inputs=duration_inputs,
            outputs=[duration_estimate],
            queue=False,
        )
    demo.load(
        format_duration_panel,
        inputs=duration_inputs,
        outputs=[duration_estimate],
        queue=False,
    )
    return {
        "duration_mode": duration_mode,
        "manual_gpu_seconds": manual_gpu_seconds,
        "calibration_multiplier": calibration_multiplier,
        "safe_mode": safe_mode,
    }


def _build_long_video_ui():
    gr.Markdown(
        "## Long video — batched multi-window route\n"
        f"Accepts {LONG_MIN_SOURCE_FRAMES}–{LONG_MAX_SOURCE_FRAMES} constant-frame-rate frames. "
        f"Uses {LONG_CHUNK_FRAMES}-frame chunks, {LONG_OVERLAP_FRAMES}-frame overlap and "
        f"{LONG_STRIDE_FRAMES}-frame stride. Up to {LONG_MAX_CHUNKS_PER_CALLBACK} missing chunks "
        "run per ZeroGPU callback; use Continue / resume until final assembly. Original input "
        "audio is remuxed at the end."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Long input video")
            prompt = gr.Textbox(
                label="Prompt — describe the natural colors",
                lines=3,
                placeholder=RABBIT_PROMPT,
            )
            preset = gr.Dropdown(
                list(RES_PRESETS),
                value="960×544 (recommended)",
                label="Resolution",
            )
            randomize = gr.Checkbox(False, label="Randomize seed")
            seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Prepare and run first batch", variant="primary")
            resume = gr.Button("Continue / resume prepared job", variant="secondary")
            status = gr.Markdown("Ready.")
            job_state = gr.State("")
        with gr.Column():
            video_out = gr.Video(label="Long colorized result")
            bundle = gr.File(label="Progress or final diagnostics bundle ZIP")
            with gr.Accordion("Runtime summary", open=False):
                summary = gr.Textbox(lines=16, interactive=False, label="Summary")
            gr.Markdown(
                f"This build supports up to {LONG_MAX_TOTAL_CHUNKS} planned chunks and "
                f"{LONG_MAX_CHUNKS_PER_CALLBACK} chunks per callback. It still uses independent "
                "chunk generation plus overlap alignment; generated latents are not fed forward."
            )

    prepared = run.click(
        prepare_long_product,
        inputs=[video_in, prompt, preset, seed, randomize],
        outputs=[job_state, status],
        queue=True,
    )
    prepared.success(
        _run_long_product,
        inputs=[job_state],
        outputs=[video_out, bundle, summary, status],
        concurrency_id="ltx23-gpu",
        concurrency_limit=1,
    )
    resume.click(
        _run_long_product,
        inputs=[job_state],
        outputs=[video_out, bundle, summary, status],
        concurrency_id="ltx23-gpu",
        concurrency_limit=1,
    )


with gr.Blocks(title="LTX-2.3 Colorize") as demo:
    gr.Markdown(
        "# 🎨 LTX-2.3 Video Colorization\n"
        "Restore natural color while preserving subject identity, framing, scene geometry, and motion."
    )

    with gr.Tabs():
        with gr.Tab("Colorize"):
            _build_colorize_ui()
        with gr.Tab("Long video (experimental)"):
            _build_long_video_ui()


if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=1,
        max_size=8,
        status_update_rate=0.5,
    ).launch(
        show_error=True,
        ssr_mode=False,
        allowed_paths=[str(JOB_ROOT), str(EXAMPLE_ROOT)],
    )