import gradio as gr
import spaces
from model import ModelHandler
from generator import Generator
from config import Config

# 1. Initialize Models
print("Initializing Application...")
handler = ModelHandler()
handler.load_models()
gen = Generator(handler)

# 2. Define GPU Inference Function
@spaces.GPU(duration=20)
def process_text(
    prompt,
    negative_prompt,
    aspect_ratio,
    cfg_scale, 
    steps, 
    seed
):
    try:
        print("--- Starting Generation ---")
        result = gen.predict(
            user_prompt=prompt,
            negative_prompt=negative_prompt,
            aspect_ratio=aspect_ratio,
            guidance_scale=cfg_scale,
            num_inference_steps=steps,
            seed=seed
        )
        print("--- Generation Complete ---")
        return result
        
    except Exception as e:
        print(f"Error during generation: {e}")
        raise gr.Error(f"An error occurred: {str(e)}")

# 3. Build Gradio Interface
with gr.Blocks(title="Pixel Art Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎮 Text to Pixel Art
        Type a prompt to generate high-quality pixel art scenes.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Prompt", 
                placeholder="e.g. cyberpunk city street at night, rain",
                info="The trigger words 'p1x3l4rt, pixel art' are added automatically."
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt", 
                placeholder="e.g., blurry, text, watermark, bad art...",
                value=Config.DEFAULT_NEGATIVE_PROMPT
            )
            
            with gr.Accordion("Settings", open=True):
                aspect_ratio = gr.Dropdown(
                    label="Aspect Ratio",
                    choices=list(Config.ASPECT_RATIOS.keys()),
                    value=Config.DEFAULT_ASPECT_RATIO,
                    info="Image dimensions (all ~1MP resolution)"
                )
                
                seed = gr.Number(
                    label="Seed", 
                    value=-1, 
                    info="-1 for random", 
                    precision=0
                )
                
                cfg_scale = gr.Slider(
                    elem_id="cfg_scale",
                    minimum=1.0, 
                    maximum=5.0, 
                    step=0.1, 
                    value=Config.CGF_SCALE, 
                    label="CFG Scale"
                )
                steps = gr.Slider(
                    elem_id="steps",
                    minimum=4, 
                    maximum=20, 
                    step=1, 
                    value=Config.STEPS_NUMBER, 
                    label="Steps"
                )
            
            run_btn = gr.Button("Generate", variant="primary")
            
        with gr.Column(scale=1):
            output_img = gr.Image(label="Result")
                
    # Event Handler
    all_inputs = [
        prompt, 
        negative_prompt,
        aspect_ratio,
        cfg_scale, 
        steps, 
        seed
    ]
    
    run_btn.click(
        fn=process_text, 
        inputs=all_inputs, 
        outputs=[output_img]
    )

# 4. Launch
if __name__ == "__main__":
    demo.queue(max_size=20, api_open=True)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=True
    )