import gradio as gr
import spaces
import torch
from model import ModelHandler
from generator import Generator
# --- IMPORT CONFIG ---
from config import Config

# 1. Initialize Models Globally (in RAM)
# ZeroGPU will move them to VRAM inside the @spaces.GPU function
print("Initializing Application...")
handler = ModelHandler()
handler.load_models()
gen = Generator(handler)

# 2. Define GPU-enabled Inference Function
@spaces.GPU(duration=20) # <-- MODIFIED
def process_img(
    image, 
    prompt,
    negative_prompt,
    cfg_scale, 
    steps, 
    img_strength, 
    face_strength,
    depth_strength, 
    edge_strength,
    # tile_strength,     # <-- REMOVED
    seed
):
    if image is None:
        raise gr.Error("Please upload an image first.")
    
    try:
        print("--- Starting Generation ---")
        # Pass all parameters to the generator
        result = gen.predict(
            image, 
            prompt,
            negative_prompt=negative_prompt,
            guidance_scale=cfg_scale,
            num_inference_steps=steps,
            img2img_strength=img_strength,
            face_strength=face_strength,
            depth_strength=depth_strength,
            lineart_strength=edge_strength,
            # tile_strength=tile_strength,      # <-- REMOVED
            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="Face To Pixel Art", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎮 Face to Pixel Art
        Upload any image. If there is a face, we'll keep the identity. If not, we'll pixelate the scene!
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            input_img = gr.Image(type="pil", label="Input Image")
            prompt = gr.Textbox(
                label="Prompt (Optional)", 
                placeholder="Leave empty for auto-captioning...",
                info="The trigger words 'p1x3l4rt, pixel art' are added automatically."
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt (Optional)", 
                placeholder="e.g., blurry, text, watermark, bad art...",
                value=Config.DEFAULT_NEGATIVE_PROMPT # <-- MODIFIED
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                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 Number"
                )
                img_strength = gr.Slider(
                    elem_id="img_strength",
                    minimum=0.1, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.IMG_STRENGTH, 
                    label="Image Strength (Img2Img)"
                )
                face_strength = gr.Slider(
                    elem_id="face_strength",
                    minimum=0.0, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.FACE_STRENGTH, 
                    label="Face Strength"
                )
                depth_strength = gr.Slider(
                    elem_id="depth_strength",
                    minimum=0.0, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.DEPTH_STRENGTH, 
                    label="DepthMap Strength"
                )
                edge_strength = gr.Slider(
                    elem_id="edge_strength",
                    minimum=0.0, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.EDGE_STRENGTH, 
                    label="EdgeMap Strength (LineArt)"
                )
                # --- MODIFIED: Renamed slider ---
                # tile_strength = gr.Slider(...) # <-- REMOVED
            
            run_btn = gr.Button("Generate Pixel Art", variant="primary")
            
        with gr.Column(scale=1):
            output_img = gr.Image(label="Pixel Art Result")
                
    # Event Handler
    all_inputs = [
        input_img, 
        prompt, 
        negative_prompt,
        cfg_scale, 
        steps, 
        img_strength,
        face_strength,
        depth_strength, 
        edge_strength,
        # tile_strength,     # <-- REMOVED
        seed
    ]
    
    run_btn.click(
        fn=process_img, 
        inputs=all_inputs, 
        outputs=[output_img]
    )


# 4. Launch the App
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 # share=True is not needed on Spaces
    )