import os
import tempfile

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

from composer import compose_frames, crop_reserved_region
from pipeline import load_pipeline, run_inference
from video_utils import (
    compute_target_size,
    extract_audio,
    frames_for_duration,
    load_video_frames,
    resize_frames,
    save_video,
)

DEFAULT_RESOLUTION = 768
REGION_SIZE = 256
_pipeline_state = None


def make_temp_file(suffix: str) -> str:
    fd, path = tempfile.mkstemp(suffix=suffix)
    os.close(fd)
    return path


def generate(
    guide_video_path: str | None,
    face_image: Image.Image | None,
    prompt: str,
    duration: float,
    fps: float,
    lora_strength: float,
    seed: int,
    hf_token: str = "",
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> tuple[str | None, str]:
    global _pipeline_state

    if not guide_video_path:
        return None, "Please upload a guide video."
    if face_image is None:
        return None, "Please upload a reference face image."
    if not prompt or not prompt.strip():
        return None, "Please enter a text prompt."

    if not os.path.isfile(guide_video_path):
        return None, f"Guide video path is not a real file: {guide_video_path}"

    if _pipeline_state is None:
        progress(0, desc="Loading models (first run only — ~5 min)…")
        _pipeline_state = load_pipeline(
            token=hf_token.strip() or None,
            progress_cb=lambda msg: progress(0, desc=msg),
        )

    progress(0.05, desc="Loading guide video…")
    frames, source_fps = load_video_frames(guide_video_path)
    if len(frames) == 0:
        return None, "Could not read frames from the guide video."

    audio_tmp = make_temp_file(".wav")
    has_audio = extract_audio(guide_video_path, audio_tmp)

    progress(0.10, desc="Resizing frames…")
    orig_h, orig_w = frames.shape[1], frames.shape[2]
    target_w, target_h = compute_target_size(orig_w, orig_h, DEFAULT_RESOLUTION)
    frames = resize_frames(frames, target_w, target_h)

    n_frames = frames_for_duration(fps, duration)
    if len(frames) >= n_frames:
        frames = frames[:n_frames]
    else:
        pad = np.stack([frames[-1]] * (n_frames - len(frames)))
        frames = np.concatenate([frames, pad], axis=0)

    progress(0.15, desc="Compositing reference face strip…")
    composed = compose_frames(
        frames,
        face_image,
        region_position="left",
        region_size_px=REGION_SIZE,
    )

    progress(0.20, desc="Running LTX-2.3 diffusion…")
    generated = run_inference(
        _pipeline_state,
        composed,
        prompt=prompt,
        fps=fps,
        lora_strength=lora_strength,
        seed=int(seed),
        progress_cb=lambda msg: progress(0.20, desc=msg),
    )

    progress(0.90, desc="Cropping reserved region…")
    cropped = crop_reserved_region(
        generated,
        region_position="left",
        region_size_px=REGION_SIZE,
        output_size=(target_w, target_h),
    )

    progress(0.95, desc="Encoding output video…")
    out_path = make_temp_file(".mp4")
    save_video(
        cropped,
        fps=fps,
        output_path=out_path,
        audio_path=audio_tmp if has_audio else None,
        audio_duration=duration,
    )

    if not os.path.isfile(out_path):
        return None, f"Output file was not created: {out_path}"

    progress(1.0, desc="Done.")
    return out_path, "Genera