import gradio as gr
import fal_client
import os

# API Key Check
if not os.environ.get("FAL_KEY"):
    print("WARNUNG: 'FAL_KEY' fehlt in den Secrets!")

MODELS = [
    "fal-ai/qwen-image-2512",
    "fal-ai/flux-2",
    "fal-ai/flux-2/klein/9b",
    "fal-ai/flux/schnell", 
]

def generate_image(prompt, model_name, steps, cfg_scale, width, height):
    
    print(f"Generiere: {model_name} | {width}x{height} | Steps: {steps} | CFG: {cfg_scale}")

    # Argumente vorbereiten
    args = {
        "prompt": prompt,
        # Bei fal.ai kann man oft benutzerdefinierte Größen als Dictionary übergeben
        # oder direkt als width/height, je nach Modell.
        # Die sicherste Variante für Flux/Qwen ist oft "image_size" als Dict oder String.
        "image_size": {
            "width": int(width),
            "height": int(height)
        },
        "num_inference_steps": int(steps),
        "guidance_scale": float(cfg_scale), # Das ist der CFG Scale
        "enable_safety_checker": False
    }

    # Ausnahme für Flux Pro/v1.1 ("Flux-2"): 
    # Diese Modelle akzeptieren oft keine manuellen Steps oder CFG.
    # Wir filtern das, damit kein Fehler kommt.
    if "flux-pro" in model_name or "flux-2" in model_name:
        if "num_inference_steps" in args: del args["num_inference_steps"]
        if "guidance_scale" in args: del args["guidance_scale"]

    try:
        handler = fal_client.submit(model_name, arguments=args)
        result = handler.get()
        
        if "images" in result: return result["images"][0]["url"]
        elif "image" in result: return result["image"]["url"]
        else: raise gr.Error("Keine Bild-URL erhalten.")
            
    except Exception as e:
        raise gr.Error(f"Fehler bei '{model_name}': {str(e)}")

# --- UI Aufbau ---
with gr.Blocks(title="Custom Model Gen") as demo:
    gr.Markdown("## ⚡ Custom Model Generator")
    
    with gr.Row():
        # LINKER BEREICH
        with gr.Column(scale=1):
            model_selector = gr.Dropdown(
                choices=MODELS, 
                value=MODELS[1], 
                label="Modell ID"
            )
            
            prompt_input = gr.Textbox(
                label="Prompt", 
                placeholder="Beschreibe dein Bild...",
                lines=3
            )
            
            # --- HIER IST DEIN NEUES DESIGN ---
            with gr.Accordion("⚙️ Erweiterte Parameter", open=True):
                
                # 1. Steps Slider
                steps_slider = gr.Slider(
                    minimum=10, maximum=50, value=25, step=1,
                    label="Steps (Qualität)",
                    info="Mehr Schritte = Bessere Qualität, aber langsamer."
                )
                
                # 2. CFG Scale Slider
                cfg_slider = gr.Slider(
                    minimum=1, maximum=20, value=7.5, step=0.5,
                    label="CFG Scale (Text-Treue)",
                    info="Höher = Hält sich strikter an den Text."
                )
                
                # 3. Breite und Höhe nebeneinander
                with gr.Row():
                    width_slider = gr.Slider(
                        minimum=256, maximum=1024, value=1024, step=8,
                        label="Breite"
                    )
                    height_slider = gr.Slider(
                        minimum=256, maximum=1024, value=768, step=8,
                        label="Höhe"
                    )
            # ----------------------------------

            run_btn = gr.Button("🚀 Generieren", variant="primary")
        
        # RECHTER BEREICH
        with gr.Column(scale=2):
            result_output = gr.Image(label="Ergebnis")

    # Inputs bündeln
    inputs = [prompt_input, model_selector, steps_slider, cfg_slider, width_slider, height_slider]
    
    run_btn.click(fn=generate_image, inputs=inputs, outputs=result_output)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Base())