import os

os.environ.setdefault("HF_HOME", "/tmp/huggingface")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import tempfile
import time
from pathlib import Path

import spaces  # must be before torch
import gradio as gr
import torch
from diffusers import LTX2InContextPipeline

PROMPT_PREFIX = "3DREAL. Make it photorealistic."
EXAMPLE_DIR = Path("examples")
ASSET_DIR = Path("assets")

BASE_MODEL_ID = "diffusers/LTX-2.3-Diffusers"
LORA_MODEL_ID = "fal/LTX-2.3-3DREAL-LoRA"

DEFAULT_NEGATIVE_PROMPT = (
    "color distortion, overexposure, static, blurry details, subtitles, style, artwork, "
    "painting, frame, still, dim overall tone, worst quality, low quality, JPEG compression "
    "artifacts, ugly, mutilated, extra fingers, poorly drawn hands, poorly drawn face, "
    "deformed, disfigured, malformed limbs, fused fingers, motionless frame, cluttered "
    "background, three legs, crowded background, walking backwards"
)

RESOLUTION_MAP = {
    "720p": (1280, 704),
    "480p": (832, 480),
}

EXAMPLES = {
    "harbor": {
        "title": "Harbor blockout",
        "render": EXAMPLE_DIR / "harbor_render.mp4",
        "reference": EXAMPLE_DIR / "harbor_reference.png",
        "reference_video": EXAMPLE_DIR / "harbor_reference.mp4",
        "result": EXAMPLE_DIR / "harbor_photoreal.mp4",
        "source_gif": ASSET_DIR / "example_harbor.gif",
        "prompt": (
            "3DREAL. Make it photorealistic. A cinematic harbor scene with fishing boats, "
            "cranes, wet dock surfaces, morning light, and realistic industrial details."
        ),
        "intensity": "strong",
        "resolution": "720p",
    },
    "cargo": {
        "title": "Cargo ship",
        "render": EXAMPLE_DIR / "cargo_render.mp4",
        "reference": EXAMPLE_DIR / "cargo_reference.png",
        "reference_video": EXAMPLE_DIR / "cargo_reference.mp4",
        "result": EXAMPLE_DIR / "cargo_photoreal.mp4",
        "source_gif": ASSET_DIR / "example_cargo.gif",
        "prompt": (
            "3DREAL. Make it photorealistic. A large cargo ship at sea with cinematic lighting, "
            "realistic ocean spray, metal surfaces, and natural atmosphere."
        ),
        "intensity": "strong",
        "resolution": "720p",
    },
}


CSS = """  /* responsive v2 - 1782728030 */

.gradio-container {
  max-width: 1180px !important;
  margin: 0 auto;
}

.app-header {
  padding: 18px 0 8px;
}

.app-header h1 {
  font-size: clamp(1.75rem, 4vw, 3.25rem);
  line-height: 1;
  letter-spacing: 0;
  margin: 0;
}

.app-header p {
  max-width: 760px;
  margin: 10px 0 0;
  color: var(--body-text-color-subdued);
  font-size: 1rem;
}

.gradio-container video,
.gradio-container .video-container,
.gradio-container .video-wrapper,
.gradio-container .media-container {
  background: transparent !important;
  border: 0 !important;
  box-shadow: none !important;
}

.generate-button button {
  min-height: 46px;
  font-weight: 700;
}

/* Responsive: stack columns on narrow screens v2 */
@media (max-width: 768px) {
  /* Force all gr.Row children to stack vertically */
  .gradio-container .gr-row {
    flex-direction: column !important;
  }

  /* Reset column min-widths so they don't force horizontal overflow */
  .gradio-container .gr-column {
    min-width: 0 !important;
    width: 100% !important;
  }

  /* Ensure video/media fills the column width */
  .gradio-container video,
  .gradio-container .video-container,
  .gradio-container .video-wrapper,
  .gradio-container .media-container {
    width: 100% !important;
    max-width: 100% !important;
  }

  .app-header {
    padding-top: 10px;
  }

  .app-header h1 {
    font-size: 1.75rem;
  }

  .app-header p {
    font-size: 0.9rem;
  }
}

/* Extra small screens */
@media (max-width: 480px) {
  .gradio-container {
    padding: 8px !important;
  }

  .app-header {
    padding: 10px 0 4px;
  }

  .app-header h1 {
    font-size: 1.5rem;
  }

  /* Stack the inner rows in advanced settings too */
  .gradio-container .gr-row .gr-row {
    flex-direction: column !important;
  }
}
"""

_pipe = None
_pipe_lora_variant = None


def _load_pipeline():
    """Load the pipeline and both LoRA variants at module level on CPU.

    Downloads happen during APP_STARTING (1h timeout). The first @spaces.GPU
    call moves the pipeline to CUDA; subsequent calls reuse the warmed model.
    """
    global _pipe, _pipe_lora_variant

    if _pipe is not None:
        return

    _pipe = LTX2InContextPipeline.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.bfloat16,
    )
    _pipe.load_lora_weights(
        LORA_MODEL_ID,
        adapter_name="3dreal-light",
        weight_name="3DREAL-light.safetensors",
    )
    _pipe.load_lora_weights(
        LORA_MODEL_ID,
        adapter_name="3dreal-strong",
        weight_name="3DREAL-strong.safetensors",
    )
    _pipe_lora_variant = None


def _select_lora(intensity: str):
    """Switch the active LoRA adapter based on intensity."""
    global _pipe, _pipe_lora_variant

    adapter = "3dreal-light" if intensity == "light" else "3dreal-strong"
    if _pipe_lora_variant != adapter:
        _pipe.set_adapters(adapter, 1.0)
        _pipe_lora_variant = adapter
    return _pipe


def _component_path(value) -> str | None:
    if value is None:
        return None
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        return value.get("path") or value.get("name")
    return getattr(value, "path", None) or str(value)


def _normalize_prompt(prompt: str | None) -> str:
    prompt = (prompt or "").strip()
    if not prompt:
        return PROMPT_PREFIX

    lower_prompt = prompt.lower()
    lower_prefix = PROMPT_PREFIX.lower()
    if lower_prompt.startswith(lower_prefix):
        return prompt

    if lower_prompt.startswith("3dreal."):
        tail = prompt[len("3DREAL.") :].strip()
    elif lower_prompt.startswith("3dreal"):
        tail = prompt[len("3DREAL") :].strip(" .")
    else:
        tail = prompt

    if tail.lower().startswith("make it photorealistic."):
        return f"3DREAL. {tail}"
    return f"{PROMPT_PREFIX} {tail}".strip()


def _load_video_frames(video_path: str, num_frames: int = 121):
    """Load video frames as a list of PIL Images, padded to exactly num_frames.

    The LTX-2 VAE temporal compression ratio is 8, so num_frames must be 8k+1.
    If the source video has fewer frames, we repeat the last frame to pad.
    """
    from diffusers.utils import load_video

    frames = load_video(video_path)
    if len(frames) > num_frames:
        stride = max(1, len(frames) // num_frames)
        frames = frames[::stride][:num_frames]
    if len(frames) < num_frames:
        last = frames[-1]
        frames = frames + [last] * (num_frames - len(frames))
    return frames


def _load_image_as_pil(image_path: str):
    from PIL import Image

    return Image.open(image_path).convert("RGB")


def _round_num_frames(n: int) -> int:
    """Round num_frames to nearest valid value (8k+1)."""
    n = max(9, int(n))
    return ((n - 1) // 8) * 8 + 1


def _estimate_duration(*args, **kwargs):
    if len(args) > 7:
        num_frames = args[5]
        num_inference_steps = args[7]
    else:
        num_frames = kwargs.get("num_frames", 121)
        num_inference_steps = kwargs.get("num_inference_steps", 15)
    try:
        num_frames = int(num_frames)
        num_inference_steps = int(num_inference_steps)
    except (TypeError, ValueError):
        return 240
    # 22B model on xlarge (96GB): pipe.to("cuda") ~20s cold start,
    # ~3s per denoising step, VAE decode + video encoding ~15s.
    base = 40 + num_inference_steps * 4
    frame_factor = max(1.0, num_frames / 41)
    return min(240, int(base * frame_factor))


_load_pipeline()


@spaces.GPU(duration=_estimate_duration, size="xlarge")
def generate_video(
    render_video,
    reference_image,
    prompt,
    intensity,
    resolution,
    num_frames,
    frames_per_second,
    num_inference_steps,
    guidance_scale,
    generate_audio,
    enable_prompt_expansion,
    seed,
    video_quality,
    video_write_mode,
    progress=gr.Progress(track_tqdm=False),
):
    from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
    from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
    from diffusers.utils import encode_video

    render_path = _component_path(render_video)
    if not render_path:
        raise gr.Error("Add a 3D or CG render video before generating.")

    started = time.perf_counter()
    prompt = _normalize_prompt(prompt)

    pipe = _select_lora(intensity)
    pipe.to("cuda")

    width, height = RESOLUTION_MAP.get(resolution, RESOLUTION_MAP["720p"])
    num_frames = _round_num_frames(num_frames)

    progress(0.10, desc="Reading render video")
    render_frames = _load_video_frames(render_path, num_frames=num_frames)

    progress(0.15, desc="Preparing conditions")
    reference_conditions = [LTX2ReferenceCondition(frames=render_frames, strength=1.0)]

    conditions = None
    reference_path = _component_path(reference_image)
    if reference_path:
        ref_pil = _load_image_as_pil(reference_path)
        conditions = [LTX2VideoCondition(frames=ref_pil, index=0, strength=1.0)]

    generator = None
    gen_seed = None
    if seed is not None and str(seed).strip():
        gen_seed = int(seed)
        generator = torch.Generator(device="cuda").manual_seed(gen_seed)

    progress(0.20, desc="Running inference")
    video, audio = pipe(
        prompt=prompt,
        negative_prompt=DEFAULT_NEGATIVE_PROMPT,
        reference_conditions=reference_conditions,
        conditions=conditions,
        height=height,
        width=width,
        num_frames=num_frames,
        frame_rate=float(frames_per_second),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        generator=generator,
        output_type="np",
        return_dict=False,
    )

    progress(0.85, desc="Encoding video")
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    tmp.close()

    video_array = video[0]
    audio_tensor = None
    audio_sample_rate = None
    if generate_audio and audio is not None:
        audio_tensor = audio[0].float().cpu()
        audio_sample_rate = pipe.vocoder.config.output_sampling_rate

    encode_video(
        video_array,
        fps=int(frames_per_second),
        output_path=tmp.name,
        audio=audio_tensor,
        audio_sample_rate=audio_sample_rate,
    )

    elapsed = time.perf_counter() - started
    used_seed = gen_seed if gen_seed is not None else "random"
    status = (
        f"**Complete in {elapsed:.1f}s.** Seed: `{used_seed}`\n\n"
        f"Prompt used: `{prompt}`"
    )
    return tmp.name, status


def load_cached_example(render_video, reference_image, prompt, intensity, resolution):
    render_path = _component_path(render_video) or ""
    name = Path(render_path).stem.replace("_render", "")
    example = EXAMPLES.get(name)
    if not example:
        raise gr.Error("Could not match this cached example.")
    status = (
        f"**Cached example loaded.** Variant: `{example['intensity']}`. "
        f"Resolution: `{example['resolution']}`.\n\n"
        f"Prompt: `{example['prompt']}`"
    )
    return str(example["result"]), status


@spaces.GPU(duration=1)
def _zerogpu_probe():
    return "ready"


theme = gr.themes.Soft(
    primary_hue="cyan",
    secondary_hue="amber",
    neutral_hue="slate",
    radius_size="sm",
).set(
    body_text_color="*neutral_900",
    body_text_color_dark="*neutral_100",
    block_border_width="1px",
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
)

with gr.Blocks(
    title="LTX 2.3 3DREAL Render-to-Real",
) as demo:
    gr.HTML(
        """
        <header class="app-header">
          <h1>LTX 2.3 3DREAL Render-to-Real</h1>
          <p>Turn rough 3D, CG, and game renders into photorealistic video with the LTX-2.3 3DREAL IC-LoRA running locally on ZeroGPU.</p>
        </header>
        """
    )

    with gr.Row(elem_classes=["main-layout"]):
        with gr.Column(scale=6, min_width=0):
            render_video = gr.Video(
                label="3D / CG render video",
                sources=["upload"],
                format="mp4",
            )
            reference_image = gr.Image(
                label="Optional photoreal first-frame reference",
                sources=["upload"],
                type="filepath",
                height=250,
            )

        with gr.Column(scale=5, min_width=0):
            prompt = gr.Textbox(
                label="Prompt",
                value=PROMPT_PREFIX,
                lines=4,
                max_lines=8,
            )
            with gr.Row():
                intensity = gr.Radio(
                    choices=["light", "strong"],
                    value="light",
                    label="Variant",
                    elem_classes=["variant-toggle"],
                )
                resolution = gr.Dropdown(
                    choices=["720p", "480p"],
                    value="720p",
                    label="Resolution",
                )

            with gr.Accordion("Advanced", open=False):
                with gr.Row():
                    num_frames = gr.Slider(25, 361, value=121, step=1, label="Frames")
                    frames_per_second = gr.Slider(12, 30, value=24, step=1, label="FPS")
                with gr.Row():
                    num_inference_steps = gr.Slider(8, 30, value=15, step=1, label="Steps")
                    guidance_scale = gr.Slider(0.1, 3.0, value=1.0, step=0.1, label="Guidance")
                with gr.Row():
                    video_quality = gr.Dropdown(
                        choices=["low", "medium", "high", "maximum"],
                        value="high",
                        label="Video quality",
                    )
                    video_write_mode = gr.Dropdown(
                        choices=["fast", "balanced", "small"],
                        value="balanced",
                        label="MP4 mode",
                    )
                with gr.Row():
                    generate_audio = gr.Checkbox(value=False, label="Generate audio")
                    enable_prompt_expansion = gr.Checkbox(value=False, label="Prompt expansion")
                seed = gr.Number(value=None, precision=0, label="Seed")

            generate = gr.Button(
                "Generate video",
                variant="primary",
                elem_classes=["generate-button"],
            )

    result_video = gr.Video(label="Result video")
    status = gr.Markdown("Generate a new video.")

    gr.Examples(
        examples=[
            [
                str(EXAMPLES["harbor"]["render"]),
                str(EXAMPLES["harbor"]["reference"]),
                EXAMPLES["harbor"]["prompt"],
                EXAMPLES["harbor"]["intensity"],
                EXAMPLES["harbor"]["resolution"],
            ],
            [
                str(EXAMPLES["cargo"]["render"]),
                str(EXAMPLES["cargo"]["reference"]),
                EXAMPLES["cargo"]["prompt"],
                EXAMPLES["cargo"]["intensity"],
                EXAMPLES["cargo"]["resolution"],
            ],
        ],
        inputs=[render_video, reference_image, prompt, intensity, resolution],
        outputs=[result_video, status],
        fn=load_cached_example,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples",
        examples_per_page=2,
    )

    generate.click(
        fn=generate_video,
        inputs=[
            render_video,
            reference_image,
            prompt,
            intensity,
            resolution,
            num_frames,
            frames_per_second,
            num_inference_steps,
            guidance_scale,
            generate_audio,
            enable_prompt_expansion,
            seed,
            video_quality,
            video_write_mode,
        ],
        outputs=[result_video, status],
    )

demo.queue(max_size=12, default_concurrency_limit=2)

if __name__ == "__main__":
    demo.launch(
        theme=theme,
        css=CSS,
        allowed_paths=[str(EXAMPLE_DIR), str(ASSET_DIR)],
    )