import gradio as gr
import json
from deforum_engine import DeforumRunner

# Initialize Engine
runner = DeforumRunner(device="cpu")

def process(
    prompts_json, neg_prompt, max_frames, 
    width, height,
    zoom, angle, tx, ty,
    strength, noise,
    fps, steps
):
    # Validate JSON
    try:
        # Convert single quotes to double if user messed up
        prompts_json = prompts_json.replace("'", '"')
        # Ensure keys are integers
        prompts_dict = json.loads(prompts_json)
        prompts_fixed = {int(k): v for k, v in prompts_dict.items()}
    except:
        # Fallback manual parse
        try:
            prompts_fixed = eval(prompts_json)
        except:
            return None, None, None, "Error parsing prompts JSON"

    yield from runner.render(
        prompts_fixed, neg_prompt, int(max_frames),
        int(width), int(height),
        zoom, angle, tx, ty,
        strength, noise,
        int(fps), int(steps)
    )

css = """
#col-container {max_width: 900px; margin-left: auto; margin-right: auto;}
"""

with gr.Blocks() as demo:
    gr.Markdown("# 🌀 Authentic Deforum (CPU/SDXS)")
    gr.Markdown("""
    **Core Config:**
    * **Model:** SDXS-512-0.9 (Fast 1-step capable).
    * **Authenticity:** Full 2D Warping, Noise Schedules, and Img2Img Loopback.
    * **Tip:** Keep `Steps` around 3-4. If you use 1 step, 'Strength' (fade) will not work and frames will look disjointed.
    """)
    
    with gr.Row():
        with gr.Column():
            # Prompts
            prompts_input = gr.Code(
                label="Prompts (JSON format)", 
                value='{\n "0": "tiny cute swamp bunny, highly detailed, intricate, 8k",\n "30": "anthro swamp bunny, cyberpunk city background, neon lights"\n}',
                language="json"
            )
            neg_prompt = gr.Textbox(label="Negative Prompt", value="watermark, text, blurry")
            
            with gr.Row():
                max_frames = gr.Number(label="Max Frames", value=120, precision=0)
                fps = gr.Number(label="FPS", value=12, precision=0)
                steps = gr.Slider(1, 10, value=3, step=1, label="Steps (Rec: 3+)")
            
            with gr.Accordion("Motion Settings", open=True):
                gr.Markdown("Supports Math: `sin(t/10)`, `t*0.5`, etc. `t` = frame index.")
                angle = gr.Textbox(label="Angle (deg)", value="0:(0)")
                zoom = gr.Textbox(label="Zoom (1.0 = static)", value="0:(1.0 + 0.01*sin(t/20))")
                tx = gr.Textbox(label="Translation X", value="0:(2*sin(t/10))")
                ty = gr.Textbox(label="Translation Y", value="0:(0)")

            with gr.Accordion("Coherence (The 'Glue')", open=True):
                strength = gr.Textbox(label="Strength Schedule (0.0-1.0)", value="0:(0.65)")
                noise = gr.Textbox(label="Noise Schedule (0.0-1.0)", value="0:(0.04)")
                
            run_btn = gr.Button("Generate Animation", variant="primary")

        with gr.Column():
            out_image = gr.Image(label="Last Frame", type="pil")
            out_video = gr.Video(label="Video Output")
            out_zip = gr.File(label="Download Frames")
    
    run_btn.click(
        fn=process,
        inputs=[
            prompts_input, neg_prompt, max_frames,
            gr.State(256), gr.State(256), # Fixed resolution for CPU speed
            zoom, angle, tx, ty,
            strength, noise,
            fps, steps
        ],
        outputs=[out_image, out_video, out_zip]
    )

if __name__ == "__main__":
    demo.launch(css=css, theme=gr.themes.Glass())