from __future__ import annotations

import os
from typing import List, Tuple

import gradio as gr

from cinegen import CharacterDesigner, StoryGenerator, VideoDirector
from cinegen.models import Storyboard

STYLE_CHOICES = [
    "Cinematic Realism",
    "Neo-Noir Animation",
    "Analog Horror",
    "Retro-Futuristic",
    "Dreamlike Documentary",
]

VIDEO_MODEL_CHOICES = [
    ("Wan 2.1 (fal-ai)", "Wan-AI/Wan2.1-T2V-14B"),
    ("LTX Video 0.9.7", "Lightricks/LTX-Video-0.9.7-distilled"),
    ("Hunyuan Video 1.5", "tencent/HunyuanVideo-1.5"),
    ("CogVideoX 5B", "THUDM/CogVideoX-5b"),
]

SCENE_COLUMNS = ["Scene", "Title", "Action", "Visuals", "Characters", "Duration (s)"]
CHARACTER_COLUMNS = ["ID", "Name", "Role", "Traits"]


def _ensure_storyboard(board: Storyboard | None) -> Storyboard:
    if not board:
        raise gr.Error("Create a storyboard first.")
    return board


def _validate_inputs(idea: str | None, image_path: str | None):
    if not idea and not image_path:
        raise gr.Error("Provide either a story idea or upload a reference image.")


def handle_storyboard(
    idea: str,
    inspiration_image: str | None,
    style: str,
    scene_count: int,
    google_api_key: str,
) -> Tuple[str, List[List[str]], List[List[str]], Storyboard]:
    _validate_inputs(idea, inspiration_image)
    generator = StoryGenerator(api_key=google_api_key or None)
    storyboard = generator.generate(
        idea=idea,
        style=style,
        scene_count=scene_count,
        inspiration_path=inspiration_image,
    )
    summary_md = f"### {storyboard.title}\n{storyboard.synopsis}"
    scene_rows = storyboard.scenes_table()
    character_rows = storyboard.characters_table()
    return (
        summary_md,
        [[row[col] for col in SCENE_COLUMNS] for row in scene_rows],
        [[row[col] for col in CHARACTER_COLUMNS] for row in character_rows],
        storyboard,
    )


def handle_character_design(
    storyboard: Storyboard | None,
    google_api_key: str,
):
    board = _ensure_storyboard(storyboard)
    designer = CharacterDesigner(api_key=google_api_key or None)
    gallery, updated_board = designer.design(board)
    if not gallery:
        raise gr.Error("Failed to design characters.")
    return gallery, updated_board


def handle_video_render(
    storyboard: Storyboard | None,
    hf_token: str,
    model_choice: str,
):
    board = _ensure_storyboard(storyboard)
    prioritized_models = [model_choice] + [
        model for _, model in VIDEO_MODEL_CHOICES if model != model_choice
    ]
    director = VideoDirector(token=hf_token or None, models=prioritized_models)
    final_cut, logs = director.render(board)
    log_md = "\n".join(f"- {line}" for line in logs)
    return final_cut, log_md


css = """
#cinegen-app {
    max-width: 1080px;
    margin: 0 auto;
}
"""


with gr.Blocks(css=css, fill_height=True, theme=gr.themes.Soft(), elem_id="cinegen-app") as demo:
    gr.Markdown(
        "## 🎬 CineGen AI Director\n"
        "Drop an idea or inspiration image and let CineGen produce a storyboard, character boards, "
        "and a compiled short film using Hugging Face video models."
    )

    story_state = gr.State()

    with gr.Row():
        idea_box = gr.Textbox(
            label="Movie Idea",
            placeholder="E.g. A time loop love story set in a neon bazaar.",
            lines=3,
        )
        inspiration = gr.Image(label="Reference Image (optional)", type="filepath")

    with gr.Row():
        style_dropdown = gr.Dropdown(
            label="Visual Style",
            choices=STYLE_CHOICES,
            value=STYLE_CHOICES[0],
        )
        scene_slider = gr.Slider(
            label="Scene Count",
            minimum=3,
            maximum=8,
            value=4,
            step=1,
        )
        video_model_dropdown = gr.Dropdown(
            label="Preferred Video Model",
            choices=[choice for choice, _ in VIDEO_MODEL_CHOICES],
            value=VIDEO_MODEL_CHOICES[0][0],
        )

    with gr.Accordion("API Keys", open=False):
        google_key_input = gr.Textbox(
            label="Google API Key (Gemini)",
            type="password",
            placeholder="Required for live Gemini calls. Leave blank to use offline stubs.",
            value=os.environ.get("GOOGLE_API_KEY", ""),
        )
        hf_token_input = gr.Textbox(
            label="Hugging Face Token",
            type="password",
            placeholder="Needed for Wan/LTX/Hunyuan video generation.",
            value=os.environ.get("HF_TOKEN", ""),
        )

    storyboard_btn = gr.Button("Create Storyboard", variant="primary")
    summary_md = gr.Markdown("Storyboard output will appear here.")
    scenes_df = gr.Dataframe(headers=SCENE_COLUMNS, wrap=True)
    characters_df = gr.Dataframe(headers=CHARACTER_COLUMNS, wrap=True)

    storyboard_btn.click(
        fn=handle_storyboard,
        inputs=[idea_box, inspiration, style_dropdown, scene_slider, google_key_input],
        outputs=[summary_md, scenes_df, characters_df, story_state],
    )

    with gr.Row():
        design_btn = gr.Button("Design Characters", variant="secondary")
        render_btn = gr.Button("Render Short Film", variant="primary")

    gallery = gr.Gallery(label="Character References", columns=4, height=320)
    render_logs = gr.Markdown(label="Render Log")
    final_video = gr.Video(label="CineGen Short Film", interactive=False)

    design_btn.click(
        fn=handle_character_design,
        inputs=[story_state, google_key_input],
        outputs=[gallery, story_state],
    )

    def _model_value(label: str) -> str:
        lookup = dict(VIDEO_MODEL_CHOICES)
        return lookup.get(label, VIDEO_MODEL_CHOICES[0][1])

    def render_wrapper(board, token, label):
        return handle_video_render(board, token, _model_value(label))

    render_btn.click(
        fn=render_wrapper,
        inputs=[story_state, hf_token_input, video_model_dropdown],
        outputs=[final_video, render_logs],
    )

if __name__ == "__main__":
    demo.launch()