# app_wan.py (v1.5.0 - com yield para atualizações em tempo real)

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

# === Constantes da UI ===
MAX_SEED = np.iinfo(np.int32).max
MIN_TOTAL_FRAMES = 9
MAX_TOTAL_FRAMES = 81
STEP_TOTAL_FRAMES = 4 # Garante 4n+1 (9, 13, 17...)

MIN_HANDLER_FRAME = 17 # Garante 8n+1 e buffer inicial (2*8+1)
MAX_HANDLER_FRAME = 73 # Garante buffer final (9*8+1)
STEP_HANDLER_FRAME = 8 # Garante 8n+1

# === Importa o serviço de geração (manager) ===
from aduc_framework.managers.wan_manager import WanManager
wan_manager = WanManager()

# === Wrapper da UI para o Serviço (agora é um GERADOR) ===
def ui_generate_video(
    # --- Inputs da Aba 1: Image-to-Video ---
    convergent_img_i2v,
    handler_img_i2v,
    handler_frame_slider,
    handler_weight_slider,
    causal_img_i2v,
    total_frames_slider,
    causal_weight_slider,
    # --- Inputs da Aba 2: Video Extender ---
    convergent_img_v2v,
    causal_video_upload,
    # --- Inputs Comuns ---
    prompt,
    negative_prompt,
    fps,
    resolution,
    steps,
    guidance_scale,
    guidance_scale_2,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    # Determina o modo de operação e a imagem convergente correta
    is_v2v_mode = causal_video_upload is not None
    convergent_img = convergent_img_v2v if is_v2v_mode else convergent_img_i2v
    
    # Validação de entradas essenciais
    if convergent_img is None:
        raise gr.Error("A 'Convergent Image (Start)' é obrigatória em ambos os modos.")
    if not is_v2v_mode and causal_img_i2v is None:
        raise gr.Error("A 'Causal Image (End)' é obrigatória no modo Image-to-Video.")

    # Itera sobre os resultados parciais produzidos pelo manager
    # A última iteração conterá o resultado final completo.
    for result in wan_manager.generate_video(
        convergent_img=convergent_img,
        causal_video_path=causal_video_upload,
        causal_img=causal_img_i2v,
        handler_img=handler_img_i2v,
        total_frames=int(total_frames_slider),
        handler_frame=int(handler_frame_slider),
        causal_weight=float(causal_weight_slider),
        handler_weight=float(handler_weight_slider),
        prompt=prompt,
        negative_prompt=negative_prompt,
        fps=int(fps),
        resolution=resolution,
        steps=int(steps),
        guidance_scale=float(guidance_scale),
        guidance_scale_2=float(guidance_scale_2),
        seed=int(seed),
        randomize_seed=bool(randomize_seed),
    ):
        # result é uma tupla: (video_path, seed, lista_de_videos_de_passo)
        final_video_path, current_seed, denoising_videos = result
        
        # Cria um dicionário de atualização para o Gradio.
        # A cada passo, atualizamos a galeria com a lista de vídeos.
        # No último passo, o vídeo principal e a seed também são atualizados.
        update_dict = {
            output_video: gr.update(value=final_video_path),
            seed_input: gr.update(value=current_seed),
            steps_video_gallery: gr.update(value=denoising_videos)
        }
        
        # Produz (yield) a atualização para a UI do Gradio
        yield update_dict

# === Interface Gradio ===
css = '''
.fillable{max-width: 1100px !important}
.dark .progress-text {color: white}
#general_items{margin-top: 2em}
'''

with gr.Blocks(theme=gr.themes.Glass(), css=css) as app:
    gr.Markdown("# Wan 2.2 Aduca-SDR")

    with gr.Row(elem_id="general_items"):
        with gr.Column(scale=2):
            with gr.Tabs() as tabs:
                with gr.TabItem("Image-to-Video", id="i2v_tab"):
                    with gr.Row():
                        with gr.Column():
                            convergent_image_i2v = gr.Image(type="pil", label="Convergent Img (Start)", sources=["upload", "clipboard"])
                        with gr.Column():
                            handler_image_i2v = gr.Image(type="pil", label="Handler Img (Middle)", sources=["upload", "clipboard"])
                            handler_frame_slider = gr.Slider(
                                minimum=MIN_HANDLER_FRAME, maximum=MAX_HANDLER_FRAME, step=STEP_HANDLER_FRAME,
                                value=25, label="Handler Frame"
                            )
                            handler_weight_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="Handler Weight")
                        with gr.Column():
                            causal_image_i2v = gr.Image(type="pil", label="Causal Img (End)", sources=["upload", "clipboard"])
                            total_frames_slider = gr.Slider(
                                minimum=MIN_TOTAL_FRAMES, maximum=MAX_TOTAL_FRAMES, step=STEP_TOTAL_FRAMES,
                                value=33, label="Total Frames"
                            )
                            causal_weight_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="Causal Weight")
                
                with gr.TabItem("Video Extender", id="v2v_tab"):
                    with gr.Row():
                        with gr.Column():
                            convergent_image_v2v = gr.Image(type="pil", label="Convergent Img (Start)", sources=["upload", "clipboard"])
                            gr.Markdown("<p style='font-size:0.8rem;color:gray'>A imagem para iniciar a nova seção do vídeo.</p>")
                        with gr.Column():
                            causal_video_upload = gr.Video(label="Causal Video", sources=["upload"])
                            gr.Markdown("<p style='font-size:0.8rem;color:gray'>O vídeo a ser estendido.</p>")

            gr.Markdown("---")
            prompt = gr.Textbox(
                label="Prompt",
                info="Descreva a nova cena ou transição. Ex: 'the scene changes to a futuristic city at night'."
            )
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    resolution_selector = gr.Dropdown(
                        choices=["480x832", "832x480", "240x416", "416x240", "720x1280", "1280x720", "512x512"],
                        value="480x832", label="Resolution (H x W)",
                        info="No modo Video Extender, o vídeo será conformado para esta resolução."
                    )
                    fps_selector = gr.Dropdown(choices=[8, 16, 24], value=16, label="FPS", info="O vídeo será conformado para este FPS.")
                
                negative_prompt_input = gr.Textbox(
                    label="Negative Prompt",
                    value=wan_manager.default_negative_prompt,
                    lines=3
                )
                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=8, label="Inference Steps")
                guidance_scale_input = gr.Slider(
                    minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale (High Noise)"
                )
                guidance_scale_2_input = gr.Slider(
                    minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale (Low Noise)"
                )
                with gr.Row():
                    seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                    randomize_seed_checkbox = gr.Checkbox(label="Randomize Seed", value=True)

            generate_button = gr.Button("Generate Video", variant="primary")

        with gr.Column(scale=1):
            output_video = gr.Video(label="Generated Video", autoplay=True)
            with gr.Accordion("Denoising Process Visuals", open=True):
                gr.Markdown("Evolução do vídeo completo em cada passo do Denoising.")
                steps_video_gallery = gr.Gallery(
                    label="Denoising Steps",
                    object_fit="contain",
                    height="auto",
                    columns=4,
                )

    ui_inputs = [
        convergent_image_i2v,
        handler_image_i2v, handler_frame_slider, handler_weight_slider,
        causal_image_i2v, total_frames_slider, causal_weight_slider,
        convergent_image_v2v,
        causal_video_upload,
        prompt, negative_prompt_input,
        fps_selector, resolution_selector,
        steps_slider, guidance_scale_input, guidance_scale_2_input,
        seed_input, randomize_seed_checkbox,
    ]
    ui_outputs = [output_video, seed_input, steps_video_gallery]

    generate_button.click(fn=ui_generate_video, inputs=ui_inputs, outputs=ui_outputs)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860, show_error=True)