import os
import time
import uuid

import spaces
import torch
import gradio as gr
from diffusers import LTXPipeline, LTXImageToVideoPipeline
from diffusers.utils import export_to_video
from huggingface_hub import InferenceClient

# Model used to expand a short idea (in any language) into a detailed,
# English, LTX-style prompt. This runs on HF's free serverless Inference
# API, NOT on your ZeroGPU quota — separate budget entirely.
ENHANCER_MODEL = "Qwen/Qwen2.5-7B-Instruct"

# Requires a Space secret named HF_TOKEN (Settings -> Variables and
# secrets -> New secret). Create the token at
# https://huggingface.co/settings/tokens with "read" access.
_hf_token = os.environ.get("HF_TOKEN")
_hf_client = InferenceClient(token=_hf_token)

ENHANCER_SYSTEM_PROMPT = (
    "You turn a short video idea, in any language, into a single detailed "
    "English prompt for the LTX-Video text-to-video AI model. Always "
    "translate to English first. Describe: the subject's appearance, the "
    "action/motion happening, the environment, camera framing, and "
    "lighting/style. Keep it to 2-4 sentences, vivid and concrete, no "
    "bullet points, no preamble, no quotation marks — output ONLY the "
    "final prompt text."
)

# Available checkpoints: fast/light vs. higher quality/slower.
# Pipelines are loaded lazily (only when first selected) and cached,
# so startup stays quick and you're not holding multiple models in
# memory unless you actually use them.
MODEL_OPTIONS = {
    "Fast (LTX-Video, original) — lower quality, cheapest on quota": "Lightricks/LTX-Video",
    "Higher quality (LTX-Video-0.9.5) — slower, costs more quota": "Lightricks/LTX-Video-0.9.5",
}
DEFAULT_MODEL_LABEL = "Fast (LTX-Video, original) — lower quality, cheapest on quota"

# Presets: (width, height, num_frames, steps, guidance)
PRESETS = {
    "🟢 Draft (cheap & fast)": (512, 320, 49, 20, 3.0),
    "🔵 Quality (slower, better look)": (704, 480, 65, 32, 3.0),
    "🟣 Long (more frames, same res as Draft)": (512, 320, 97, 20, 3.0),
}

HISTORY_DIR = "/tmp/history"
os.makedirs(HISTORY_DIR, exist_ok=True)
MAX_HISTORY = 8

# Global, in-memory history — fine for a single-user personal Space.
# Resets if the Space restarts (ephemeral storage), but persists across
# generations within a running session.
_history = []

_t2v_cache = {}
_i2v_cache = {}


def get_t2v_pipeline(model_label):
    model_id = MODEL_OPTIONS[model_label]
    if model_id not in _t2v_cache:
        pipe = LTXPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
        if torch.cuda.is_available():
            pipe.to("cuda")
        _t2v_cache[model_id] = pipe
    return _t2v_cache[model_id]


def get_i2v_pipeline(model_label):
    model_id = MODEL_OPTIONS[model_label]
    if model_id not in _i2v_cache:
        pipe = LTXImageToVideoPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
        if torch.cuda.is_available():
            pipe.to("cuda")
        _i2v_cache[model_id] = pipe
    return _i2v_cache[model_id]


# Warm up the default (fast) text-to-video model at startup so the first
# request doesn't also pay for a cold model load. Image-to-video and the
# quality model load lazily on first use instead, to keep startup light.
get_t2v_pipeline(DEFAULT_MODEL_LABEL)

DEFAULT_NEGATIVE = (
    "worst quality, inconsistent motion, blurry, jittery, distorted, "
    "low resolution, deformed"
)


def _save_to_history(prompt, out_path):
    """Copy a generated clip into the history folder and record it."""
    ext = os.path.splitext(out_path)[1] or ".mp4"
    dest = os.path.join(HISTORY_DIR, f"{uuid.uuid4().hex}{ext}")
    with open(out_path, "rb") as src, open(dest, "wb") as dst:
        dst.write(src.read())
    _history.insert(0, {"path": dest, "prompt": prompt, "time": time.time()})
    del _history[MAX_HISTORY:]


def _history_gallery_items():
    return [(h["path"], h["prompt"]) for h in _history]


def apply_preset(preset_name):
    w, h, f, s, g = PRESETS[preset_name]
    return w, h, f, s, g


def enhance_prompt(short_text):
    if not short_text or not short_text.strip():
        raise gr.Error("Type a short idea first, then click Enhance.")

    try:
        completion = _hf_client.chat.completions.create(
            model=ENHANCER_MODEL,
            messages=[
                {"role": "system", "content": ENHANCER_SYSTEM_PROMPT},
                {"role": "user", "content": short_text.strip()},
            ],
            max_tokens=200,
            temperature=0.7,
        )
        return completion.choices[0].message.content.strip()
    except Exception as e:
        raise gr.Error(
            f"Prompt enhancer failed ({e}). You can still type your own "
            "detailed English prompt directly."
        )


@spaces.GPU(duration=60)  # ignored/no-op on non-ZeroGPU hardware; keep this
# as low as your typical generation allows — ZeroGPU reserves this many
# seconds from your daily quota on every call, whether or not you use it all.
def generate_t2v(model_label, prompt, negative_prompt, width, height, num_frames, steps, guidance, seed):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt describing the video you want.")

    pipe = get_t2v_pipeline(model_label)

    generator = None
    if seed is not None and int(seed) >= 0:
        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
        generator.manual_seed(int(seed))

    video = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or DEFAULT_NEGATIVE,
        width=int(width),
        height=int(height),
        num_frames=int(num_frames),
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        generator=generator,
    ).frames[0]

    out_path = f"/tmp/t2v_{uuid.uuid4().hex}.mp4"
    export_to_video(video, out_path, fps=24)
    _save_to_history(prompt, out_path)
    return out_path, gr.update(value=_history_gallery_items())


@spaces.GPU(duration=60)  # same quota note as generate_t2v above
def generate_i2v(model_label, image, prompt, negative_prompt, width, height, num_frames, steps, guidance, seed):
    if image is None:
        raise gr.Error("Please upload an image to animate.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt describing the motion/scene.")

    pipe = get_i2v_pipeline(model_label)

    generator = None
    if seed is not None and int(seed) >= 0:
        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
        generator.manual_seed(int(seed))

    video = pipe(
        image=image,
        prompt=prompt,
        negative_prompt=negative_prompt or DEFAULT_NEGATIVE,
        width=int(width),
        height=int(height),
        num_frames=int(num_frames),
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        generator=generator,
    ).frames[0]

    out_path = f"/tmp/i2v_{uuid.uuid4().hex}.mp4"
    export_to_video(video, out_path, fps=24)
    _save_to_history(f"[image-to-video] {prompt}", out_path)
    return out_path, gr.update(value=_history_gallery_items())


def _settings_block():
    """Shared preset + slider block, used by both tabs."""
    with gr.Row():
        preset = gr.Radio(
            choices=list(PRESETS.keys()),
            value="🟢 Draft (cheap & fast)",
            label="Preset (click to apply, then tweak below if you want)",
        )
    with gr.Row():
        width = gr.Slider(256, 768, value=512, step=32, label="Width")
        height = gr.Slider(256, 768, value=320, step=32, label="Height")
    with gr.Row():
        num_frames = gr.Slider(9, 97, value=49, step=8, label="Number of frames")
        steps = gr.Slider(8, 40, value=20, step=1, label="Inference steps")
    with gr.Row():
        guidance = gr.Slider(1.0, 10.0, value=3.0, step=0.1, label="Guidance scale")
        seed = gr.Number(value=-1, label="Seed (-1 = random)")

    preset.change(
        fn=apply_preset,
        inputs=preset,
        outputs=[width, height, num_frames, steps, guidance],
    )
    return width, height, num_frames, steps, guidance, seed


CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@500;600;700&family=Inter:wght@400;500;600&display=swap');

:root {
    --bg: #14131a;
    --panel: #1f1d27;
    --panel-border: #322f3d;
    --text: #f3efe7;
    --text-muted: #a29cb0;
    --gold: #e4b44b;
    --teal: #4fd1c5;
}

.gradio-container {
    background: var(--bg) !important;
    font-family: 'Inter', sans-serif !important;
    color: var(--text) !important;
}

#app-header {
    padding: 8px 4px 4px 4px;
    border-bottom: 1px solid var(--panel-border);
    margin-bottom: 12px;
    background:
        radial-gradient(circle, var(--panel-border) 1.5px, transparent 1.5px) 0 0 / 14px 14px,
        radial-gradient(circle, var(--panel-border) 1.5px, transparent 1.5px) 0 100% / 14px 14px;
    background-repeat: repeat-x;
    background-position: top left, bottom left;
    padding-top: 14px;
    padding-bottom: 14px;
}

#app-header h1 {
    font-family: 'Space Grotesk', sans-serif !important;
    font-weight: 700 !important;
    letter-spacing: -0.01em;
    background: linear-gradient(90deg, var(--gold), var(--teal));
    -webkit-background-clip: text;
    background-clip: text;
    color: transparent !important;
    display: inline-block;
}

#quota-tip {
    background: var(--panel) !important;
    border: 1px solid var(--panel-border) !important;
    border-radius: 12px !important;
    padding: 10px 14px !important;
    color: var(--text-muted) !important;
    font-size: 0.9em;
}

.tab-panel {
    background: var(--panel) !important;
    border: 1px solid var(--panel-border) !important;
    border-radius: 16px !important;
    padding: 18px !important;
}

button.primary {
    background: linear-gradient(90deg, var(--gold), #c98f2e) !important;
    border: none !important;
    color: #1a1408 !important;
    font-weight: 600 !important;
}

#enhance-btn-t2v, #enhance-btn-i2v {
    background: transparent !important;
    border: 1px solid var(--teal) !important;
    color: var(--teal) !important;
    font-weight: 500 !important;
}
#enhance-btn-t2v:hover, #enhance-btn-i2v:hover {
    background: rgba(79, 209, 197, 0.12) !important;
}

#history-panel {
    background: var(--panel) !important;
    border: 1px solid var(--panel-border) !important;
    border-radius: 16px !important;
    padding: 18px !important;
    margin-top: 8px;
}

#history-panel h2 {
    font-family: 'Space Grotesk', sans-serif !important;
    font-size: 1.1em !important;
    color: var(--text) !important;
    margin: 0 0 12px 0 !important;
}

#history-gallery {
    background: transparent !important;
    border: none !important;
}
#history-gallery .thumbnail-item,
#history-gallery .grid-wrap,
#history-gallery [data-testid="thumbnail item"] {
    border-radius: 12px !important;
    border: 1px solid var(--panel-border) !important;
    overflow: hidden;
    transition: transform 0.15s ease, border-color 0.15s ease;
}
#history-gallery .thumbnail-item:hover {
    transform: translateY(-3px);
    border-color: var(--gold) !important;
}
#history-gallery .caption,
#history-gallery .caption-label {
    background: rgba(20, 19, 26, 0.85) !important;
    color: var(--text) !important;
    font-size: 0.78em !important;
    white-space: nowrap;
    overflow: hidden;
    text-overflow: ellipsis;
}

#footer-tip {
    color: var(--text-muted) !important;
    font-size: 0.85em;
    text-align: center;
    margin-top: 6px;
}
"""


with gr.Blocks(title="Free LTX Video Generator", css=CUSTOM_CSS) as demo:
    with gr.Column(elem_id="app-header"):
        gr.Markdown(
            """
            # 🎬 LTX Video Generator
            Personal video generator powered by [LTX-Video](https://huggingface.co/Lightricks/LTX-Video).
            """
        )
    gr.Markdown(
        "💡 **Quota tip:** ZeroGPU gives you a small daily budget of GPU "
        "seconds. Lower resolution/frames/steps use less of it per "
        "generation — the **Draft** preset is the cheapest, **Quality** and "
        "**Long** cost more.",
        elem_id="quota-tip",
    )

    with gr.Tabs():
        with gr.Tab("Text → Video"):
            with gr.Row():
                with gr.Column(scale=1, elem_classes="tab-panel"):
                    model_choice_t2v = gr.Dropdown(
                        choices=list(MODEL_OPTIONS.keys()),
                        value=DEFAULT_MODEL_LABEL,
                        label="Model",
                    )
                    prompt_t2v = gr.Textbox(
                        label="Prompt (Hebrew or English — short idea is fine)",
                        placeholder="A golden retriever running through a field of sunflowers at sunset, cinematic lighting, realistic style",
                        lines=4,
                    )
                    enhance_btn_t2v = gr.Button("✨ Enhance & Translate Prompt", elem_id="enhance-btn-t2v")
                    negative_prompt_t2v = gr.Textbox(
                        label="Negative prompt (optional)",
                        value=DEFAULT_NEGATIVE,
                        lines=2,
                    )
                    (width_t2v, height_t2v, frames_t2v, steps_t2v,
                     guidance_t2v, seed_t2v) = _settings_block()
                    run_btn_t2v = gr.Button("Generate Video", variant="primary")

                with gr.Column(scale=1, elem_classes="tab-panel"):
                    output_video_t2v = gr.Video(label="Result")

        with gr.Tab("Image → Video"):
            with gr.Row():
                with gr.Column(scale=1, elem_classes="tab-panel"):
                    model_choice_i2v = gr.Dropdown(
                        choices=list(MODEL_OPTIONS.keys()),
                        value=DEFAULT_MODEL_LABEL,
                        label="Model",
                    )
                    input_image = gr.Image(label="Starting image", type="pil")
                    prompt_i2v = gr.Textbox(
                        label="Prompt (Hebrew or English — short idea is fine)",
                        placeholder="The cat slowly turns its head and blinks, gentle breeze moving its fur, cinematic lighting",
                        lines=4,
                    )
                    enhance_btn_i2v = gr.Button("✨ Enhance & Translate Prompt", elem_id="enhance-btn-i2v")
                    negative_prompt_i2v = gr.Textbox(
                        label="Negative prompt (optional)",
                        value=DEFAULT_NEGATIVE,
                        lines=2,
                    )
                    (width_i2v, height_i2v, frames_i2v, steps_i2v,
                     guidance_i2v, seed_i2v) = _settings_block()
                    run_btn_i2v = gr.Button("Animate Image", variant="primary")

                with gr.Column(scale=1, elem_classes="tab-panel"):
                    output_video_i2v = gr.Video(label="Result")

    with gr.Column(elem_id="history-panel"):
        gr.Markdown("## 🕘 History — last 8 generations")
        history_gallery = gr.Gallery(
            label="History",
            show_label=False,
            value=_history_gallery_items(),
            columns=4,
            rows=2,
            object_fit="cover",
            height=340,
            elem_id="history-gallery",
        )

    gr.Markdown(
        "Tip: keep width/height multiples of 32 and frames as `8n+1` "
        "(e.g. 49, 65, 97) — these match LTX-Video's training constraints "
        "and avoid shape errors.",
        elem_id="footer-tip",
    )

    enhance_btn_t2v.click(
        fn=enhance_prompt,
        inputs=prompt_t2v,
        outputs=prompt_t2v,
    )

    enhance_btn_i2v.click(
        fn=enhance_prompt,
        inputs=prompt_i2v,
        outputs=prompt_i2v,
    )

    run_btn_t2v.click(
        fn=generate_t2v,
        inputs=[model_choice_t2v, prompt_t2v, negative_prompt_t2v, width_t2v,
                height_t2v, frames_t2v, steps_t2v, guidance_t2v, seed_t2v],
        outputs=[output_video_t2v, history_gallery],
    )

    run_btn_i2v.click(
        fn=generate_i2v,
        inputs=[model_choice_i2v, input_image, prompt_i2v, negative_prompt_i2v,
                width_i2v, height_i2v, frames_i2v, steps_i2v, guidance_i2v, seed_i2v],
        outputs=[output_video_i2v, history_gallery],
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch()