import os
import io
import tempfile

import gradio as gr
import requests
from PIL import Image

import fal_client

APP_NAME = os.getenv("APP_NAME", "Scenora")
MODEL_ID = os.getenv("MODEL_ID", "fal-ai/flux-pro/kontext")
APP_PASSWORD = os.getenv("APP_PASSWORD")

PROMPT_SUFFIX = (
    "commercial product photography, natural light, sharp focus on the product, "
    "realistic shadows and reflections, 24mm lens"
)

PRESETS = {
    "Studio neutre": "on a seamless light grey studio backdrop, soft diffused lighting",
    "Table en bois": "placed on a warm oak table, blurred kitchen in the background, morning light",
    "Beton mineral": "on a raw concrete surface, minimalist industrial setting, side light",
    "Exterieur nature": "on a mossy stone outdoors, soft green foliage behind, golden hour",
    "Marbre luxe": "on a white marble slab, elegant premium setting, soft reflections",
}

NEGATIVE_DEFAULT = "text, watermark, logo, people, clutter, distorted proportions"


def _image_to_url(image: Image.Image) -> str:
    """Envoie l'image sur le stockage temporaire de fal et renvoie son URL."""
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        path = tmp.name
    image.convert("RGB").save(path, format="PNG")
    try:
        return fal_client.upload_file(path)
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass


def _download(url: str) -> Image.Image:
    resp = requests.get(url, timeout=120)
    resp.raise_for_status()
    return Image.open(io.BytesIO(resp.content)).convert("RGB")


def generate(image, preset, prompt, negative, variations, progress=gr.Progress()):
    if not os.getenv("FAL_KEY"):
        raise gr.Error(
            "Cle API absente. Ajoutez FAL_KEY dans les Secrets du Space, "
            "puis redemarrez-le."
        )
    if image is None:
        raise gr.Error("Chargez d'abord une photo produit.")

    scene = prompt.strip() or PRESETS.get(preset, "")
    if not scene:
        raise gr.Error("Decrivez la scene souhaitee ou choisissez un preset.")

    instruction = (
        f"Replace the background only. Keep the product exactly as it is: "
        f"same shape, same colors, same proportions, same details. "
        f"New background: {scene}. Style: {PROMPT_SUFFIX}."
    )
    if negative.strip():
        instruction += f" Avoid: {negative.strip()}."

    progress(0, desc="Envoi de l'image...")
    image_url = _image_to_url(image)

    results = []
    n = int(variations)
    for i in range(n):
        progress((i + 0.2) / n, desc=f"Generation {i + 1}/{n}...")
        try:
            out = fal_client.subscribe(
                MODEL_ID,
                arguments={"prompt": instruction, "image_url": image_url},
            )
        except Exception as exc:
            raise gr.Error(f"Erreur API ({MODEL_ID}) : {exc}")

        images = out.get("images") or []
        if not images:
            raise gr.Error(f"Reponse vide du modele : {out}")

        progress((i + 0.8) / n, desc=f"Recuperation {i + 1}/{n}...")
        results.append(_download(images[0]["url"]))

    return results, instruction


CSS = """
html, body, .gradio-container {
    width: 100% !important;
    max-width: 100% !important;
}
.gradio-container {
    margin: 0 !important;
    padding: 1rem 1.5rem 2rem !important;
}
.main, .wrap, .contain, .app {
    max-width: 100% !important;
    width: 100% !important;
}
footer {display: none !important;}
"""

with gr.Blocks(css=CSS, title=APP_NAME, fill_width=True, fill_height=True) as demo:
    gr.Markdown(f"# {APP_NAME}")
    gr.Markdown(
        "Chargez une photo produit, choisissez ou decrivez le decor, generez. "
        "Le produit est preserve, seul l'arriere-plan change."
    )

    with gr.Row(equal_height=True):
        with gr.Column(scale=1, min_width=320):
            image_in = gr.Image(label="Photo produit", type="pil", height=420)
            preset = gr.Radio(
                label="Preset de decor",
                choices=list(PRESETS.keys()),
                value="Studio neutre",
            )
            prompt = gr.Textbox(
                label="Decor sur mesure (prioritaire sur le preset)",
                placeholder="sur un plan de travail en marbre, cuisine lumineuse en arriere-plan",
                lines=2,
            )
            negative = gr.Textbox(
                label="A eviter",
                value=NEGATIVE_DEFAULT,
                lines=2,
            )
            variations = gr.Slider(
                label="Nombre de variantes",
                minimum=1, maximum=4, step=1, value=2,
            )
            run = gr.Button("Generer", variant="primary")

        with gr.Column(scale=1, min_width=320):
            gallery = gr.Gallery(
                label="Resultats", columns=2, height=520,
                object_fit="contain", show_download_button=True,
            )
            debug = gr.Textbox(label="Instruction envoyee au modele", lines=4)

    run.click(
        fn=generate,
        inputs=[image_in, preset, prompt, negative, variations],
        outputs=[gallery, debug],
    )

    gr.Markdown(f"Modele : `{MODEL_ID}` — facture a l'image par le fournisseur.")

if __name__ == "__main__":
    auth = ("admin", APP_PASSWORD) if APP_PASSWORD else None
    demo.queue(max_size=8).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
        auth=auth,
        show_api=False,
    )