"""MiniMax-H3 Prompt Rewriter LoRA 8B — Gradio Space.

Loads Qwen3-VL-8B-Instruct with the MiniMax-H3-Prompt-Rewriter LoRA adapter
(lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA-8B) and rewrites short user prompts
into production-ready, structured MiniMax-H3 audio-video prompts.

Supports four tasks:
  - T2VA  (text-only to audio-video)
  - I2VA  (first-frame image + text)
  - L2VA  (last-frame image + text)
  - FL2VA (first & last-frame images + text)

Reference: https://huggingface.co/lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA-8B
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / transformers / peft
import torch
import transformers
from PIL import Image, ImageOps
from peft import PeftModel
from transformers import AutoProcessor

import gradio as gr

from prompt_template import build_messages, expected_image_count, normalize_task


BASE_MODEL = "Qwen/Qwen3-VL-8B-Instruct"
ADAPTER_REPO = "lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA-8B"

TASKS = ["T2VA", "I2VA", "L2VA", "FL2VA"]
RESOLUTIONS = ["16:9", "adaptive", "21:9", "4:3", "1:1", "3:4", "9:16"]


def _get_model_class():
    """Prefer Qwen3-VL's concrete class, with portable AutoModel fallbacks."""
    candidates = (
        "Qwen3VLForConditionalGeneration",
        "AutoModelForImageTextToText",
        "AutoModelForVision2Seq",
        "AutoModelForMultimodalLM",
    )
    for name in candidates:
        model_class = getattr(transformers, name, None)
        if model_class is not None:
            return model_class
    raise RuntimeError(
        "This Transformers installation does not expose a Qwen3-VL-compatible "
        "conditional-generation class. Upgrade Transformers and retry."
    )


print(f"[boot] Loading processor from {BASE_MODEL}…", flush=True)
processor = AutoProcessor.from_pretrained(
    BASE_MODEL,
    trust_remote_code=True,
    min_pixels=256 * 256,
    max_pixels=1024 * 1024,
)

print(f"[boot] Loading base model {BASE_MODEL}…", flush=True)
model = _get_model_class().from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
    attn_implementation="sdpa",
)

print(f"[boot] Loading LoRA adapter {ADAPTER_REPO}…", flush=True)

# On ZeroGPU, safetensors' torch.load_file is intercepted by the spaces
# hijack which routes tensor creation through CUDA — this fails at module
# scope because there is no real GPU in the main process.  We work around
# this by loading adapter weights as numpy arrays (bypassing torch entirely)
# and converting to CPU tensors, then applying via set_peft_model_state_dict.
from safetensors import safe_open
from peft import LoraConfig, PeftModel, set_peft_model_state_dict
from huggingface_hub import hf_hub_download
import json as _json

adapter_file = hf_hub_download(ADAPTER_REPO, "adapter_model.safetensors")
adapter_state_dict = {}
with safe_open(adapter_file, framework="numpy", device="cpu") as f:
    for key in f.keys():
        adapter_state_dict[key] = torch.from_numpy(f.get_tensor(key))

# Build the PEFT config from the adapter_config.json
adapter_config_path = hf_hub_download(ADAPTER_REPO, "adapter_config.json")
with open(adapter_config_path) as f:
    adapter_config_dict = _json.load(f)
peft_config = LoraConfig(
    r=adapter_config_dict["r"],
    lora_alpha=adapter_config_dict["lora_alpha"],
    lora_dropout=adapter_config_dict["lora_dropout"],
    target_modules=adapter_config_dict["target_modules"],
    bias=adapter_config_dict["bias"],
    task_type=adapter_config_dict["task_type"],
)

model = PeftModel(model, peft_config)
set_peft_model_state_dict(model, adapter_state_dict)
model.eval()

# Move the full model (base + LoRA) to cuda via the ZeroGPU hijack
model = model.to("cuda")
print("[boot] Model ready.", flush=True)


def _load_pil(image):
    """Load a Gradio image input into a PIL RGB Image (or None)."""
    if image is None:
        return None
    if isinstance(image, Image.Image):
        return ImageOps.exif_transpose(image).convert("RGB")
    if isinstance(image, str):
        return ImageOps.exif_transpose(Image.open(image)).convert("RGB")
    return image


@spaces.GPU(duration=60)
def rewrite_prompt(
    prompt: str,
    task: str = "T2VA",
    duration: int = 10,
    resolution: str = "16:9",
    first_frame=None,
    last_frame=None,
    max_new_tokens: int = 4096,
    temperature: float = 0.7,
    top_p: float = 0.8,
    greedy: bool = True,
    seed: int = 42,
    progress=gr.Progress(track_tqdm=True),
):
    """Rewrite a short prompt into a structured MiniMax-H3 audio-video prompt.

    Args:
        prompt: The original short user prompt to enhance.
        task: One of T2VA (text-only), I2VA (first frame), L2VA (last frame), FL2VA (both frames).
        duration: Target video duration in seconds (4–15).
        resolution: Output resolution preset; defaults to 16:9 for T2VA, adaptive for image tasks.
        first_frame: First-frame reference image (for I2VA, FL2VA).
        last_frame: Last-frame reference image (for L2VA, FL2VA).
        max_new_tokens: Maximum tokens to generate.
        temperature: Sampling temperature (used when greedy=False).
        top_p: Nucleus sampling top-p (used when greedy=False).
        greedy: If True, use greedy decoding; otherwise sample with temperature/top_p.
        seed: RNG seed for reproducibility.
    """
    task_norm = normalize_task(task)

    # Validate image count
    required = expected_image_count(task_norm)
    images = []
    if task_norm in ("i2av", "fl2av"):
        img = _load_pil(first_frame)
        if img is not None:
            images.append(img)
    if task_norm in ("l2av", "fl2av"):
        img = _load_pil(last_frame)
        if img is not None:
            images.append(img)
    if len(images) != required:
        return f"Error: {task} requires {required} reference image(s), but {len(images)} were provided."

    # Resolve resolution default
    if not resolution or resolution == "auto":
        resolution = "16:9" if task_norm == "t2av" else "adaptive"
    if resolution not in RESOLUTIONS:
        resolution = "16:9" if task_norm == "t2av" else "adaptive"

    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    messages = build_messages(
        prompt.strip(),
        task=task_norm,
        resolution=resolution,
        duration=duration,
    )
    rendered = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    processor_kwargs = {
        "text": [rendered],
        "return_tensors": "pt",
        "padding": False,
        "return_mm_token_type_ids": True,
    }
    if images:
        processor_kwargs["images"] = images
    inputs = processor(**processor_kwargs)
    inputs = {
        key: value.to("cuda") if isinstance(value, torch.Tensor) else value
        for key, value in inputs.items()
    }

    generation_kwargs = {"max_new_tokens": max_new_tokens}
    if greedy:
        generation_kwargs["do_sample"] = False
    else:
        generation_kwargs.update(
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
        )

    with torch.inference_mode():
        output_ids = model.generate(**inputs, **generation_kwargs)
    generated_ids = output_ids[0, inputs["input_ids"].shape[1]:]
    rewritten = processor.decode(generated_ids, skip_special_tokens=True).strip()

    return rewritten


# --- Gradio UI ---

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        "# MiniMax-H3 Prompt Rewriter LoRA 8B\n"
        "Transform short prompts into production-ready, structured MiniMax-H3 audio-video prompts. "
        "Fine-tuned LoRA on [Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct) "
        "from [lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA-8B](https://huggingface.co/lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA-8B)."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            prompt = gr.Textbox(
                label="Original Prompt",
                placeholder="A corgi runs through a rainy neon-lit alley.",
                lines=3,
                scale=4,
            )
            run_btn = gr.Button("Rewrite", variant="primary", scale=1)

        with gr.Row():
            task = gr.Dropdown(
                choices=TASKS,
                value="T2VA",
                label="Task",
                info="T2VA: text-only | I2VA: first frame | L2VA: last frame | FL2VA: both frames",
            )
            duration = gr.Slider(
                minimum=4, maximum=15, value=10, step=1,
                label="Duration (seconds)",
            )
            resolution = gr.Dropdown(
                choices=RESOLUTIONS,
                value="16:9",
                label="Resolution",
                info="Defaults to 16:9 for T2VA, adaptive for image tasks.",
            )

        with gr.Row():
            first_frame = gr.Image(
                label="First Frame (I2VA / FL2VA)",
                type="pil",
                height=200,
            )
            last_frame = gr.Image(
                label="Last Frame (L2VA / FL2VA)",
                type="pil",
                height=200,
            )

        output = gr.Textbox(
            label="Rewritten MiniMax-H3 Prompt",
            lines=20,
            buttons=["copy"],
        )

        with gr.Accordion("Advanced settings", open=False):
            greedy = gr.Checkbox(value=True, label="Greedy decoding")
            max_new_tokens = gr.Slider(
                minimum=256, maximum=8192, value=4096, step=256,
                label="Max new tokens",
            )
            temperature = gr.Slider(
                minimum=0.0, maximum=2.0, value=0.7, step=0.1,
                label="Temperature (when not greedy)",
            )
            top_p = gr.Slider(
                minimum=0.0, maximum=1.0, value=0.8, step=0.05,
                label="Top-p (when not greedy)",
            )
            seed = gr.Number(value=42, precision=0, label="Seed")

    run_btn.click(
        fn=rewrite_prompt,
        inputs=[prompt, task, duration, resolution, first_frame, last_frame,
                max_new_tokens, temperature, top_p, greedy, seed],
        outputs=output,
        api_name="rewrite_prompt",
    )

    gr.Examples(
        examples=[
            ["A corgi runs through a rainy neon-lit alley.", "T2VA", 5, "16:9"],
            ["A lone astronaut plants a flag on a red, windswept planet surface.", "T2VA", 10, "16:9"],
            ["A chef in a bustling kitchen tosses vegetables in a flaming wok.", "T2VA", 8, "16:9"],
        ],
        inputs=[prompt, task, duration, resolution],
        outputs=output,
        fn=rewrite_prompt,
        cache_examples=True,
        cache_mode="lazy",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)