import gradio as gr
import json
from deforum_engine import DeforumRunner

runner = DeforumRunner(device="cpu")

def process(prompts_json, neg, frames, width, height,
            z, a, tx, ty, stre, noi, 
            fps, steps, cadence, 
            color, border, seed_beh, init_img, 
            model, lora, sched):
    try:
        p_dict = json.loads(prompts_json.replace("'", '"'))
        prompts = {int(k): v for k, v in p_dict.items()}
    except Exception as e:
        yield None, None, None, f"JSON Error: {str(e)}"
        return

    # Pass exactly 20 args + self implicitly
    yield from runner.render(
        prompts, neg, int(frames), int(width), int(height),
        z, a, tx, ty, stre, noi,
        int(fps), int(steps), int(cadence),
        color, border, seed_beh, init_img,
        model, lora, sched
    )

def stop_gen():
    runner.stop()
    return "Stopping..."

css = """
#col-container {max_width: 1000px; margin: 0 auto;}
"""

with gr.Blocks() as demo:
    gr.Markdown("# 🌀 Deforum CPU: Full Featured\nAuthentic implementation with Cadence, Seed Control, and proper Color Coherence.")
    
    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=1):
            
            with gr.Accordion("⚙️ Engine Settings", open=False):
                model = gr.Dropdown(label="Model", value="AlekseyCalvin/acs_model", 
                                    choices=["AlekseyCalvin/acs_model", "runwayml/stable-diffusion-v1-5", "IDKiro/sdxs-512-dreamshaper"])
                lora = gr.Dropdown(label="LoRA", value="latent-consistency/lcm-lora-sdv1-5", 
                                   choices=["latent-consistency/lcm-lora-sdv1-5", "None"])
                sched = gr.Dropdown(label="Sampler", value="LCM", 
                                    choices=["LCM", "Euler A", "DDIM", "DPM++ 2M"])
                seed_beh = gr.Dropdown(label="Seed Behavior", value="iter", choices=["iter", "fixed", "random"])
                init_img = gr.Image(label="Init Image", type="pil", height=200)

            prompts = gr.Code(label="Prompts (JSON)", language="json", 
                              value='{\n "0": "a beautiful forest, sun rays, 8k",\n "30": "forest fire, smoke, dramatic lighting"\n}')
            neg = gr.Textbox(label="Negative Prompt", value="lowres, text, error, cropped, worst quality, low quality")
            
            with gr.Row():
                frames = gr.Number(label="Max Frames", value=120)
                fps = gr.Number(label="FPS", value=15)
            
            with gr.Row():
                width = gr.Slider(256, 512, value=256, step=64, label="Width")
                height = gr.Slider(256, 512, value=256, step=64, label="Height")
                
            with gr.Row():
                steps = gr.Slider(1, 20, value=4, step=1, label="Steps")
                cadence = gr.Slider(1, 8, value=2, step=1, label="Cadence (Speed/Smoothness)")

            with gr.Accordion("🎬 Motion & Coherence", open=True):
                with gr.Row():
                    color = gr.Dropdown(label="Color Match", value="LAB", choices=["None", "LAB", "HSV", "RGB"])
                    border = gr.Dropdown(label="Border Mode", value="Reflect", choices=["Reflect", "Replicate", "Wrap", "Black"])
                
                z = gr.Textbox(label="Zoom", value="0:(1.01)")
                a = gr.Textbox(label="Angle", value="0:(0)")
                tx = gr.Textbox(label="Translation X", value="0:(0)")
                ty = gr.Textbox(label="Translation Y", value="0:(0)")
                stre = gr.Textbox(label="Strength (Decay)", value="0:(0.65)")
                noi = gr.Textbox(label="Noise (Grain)", value="0:(0.02)")
                
            with gr.Row():
                btn = gr.Button("GENERATE", variant="primary", scale=2)
                stop = gr.Button("STOP", variant="stop", scale=1)
            
        with gr.Column(scale=1):
            status = gr.Markdown("Ready")
            preview = gr.Image(label="Last Frame")
            video_out = gr.Video(label="Rendered Video")
            zip_out = gr.File(label="Frames ZIP")

    # Arguments: 20 inputs + self implicitly handled by click
    inputs = [
        prompts, neg, frames, width, height,
        z, a, tx, ty, stre, noi,
        fps, steps, cadence,
        color, border, seed_beh, init_img,
        model, lora, sched
    ]
    
    btn.click(process, inputs=inputs, outputs=[preview, video_out, zip_out, status])
    stop.click(stop_gen, outputs=status)

demo.queue().launch(css=css, theme=gr.themes.Glass())