import os
import sys
import traceback
from io import BytesIO
from PIL import Image
import gradio as gr
from huggingface_hub import InferenceClient


# ============================
#   TOKEN / CLIENTE
# ============================

hf_token = os.environ.get("HF_TOKEN")
token_missing = hf_token is None

client = None if token_missing else InferenceClient(
    provider="fal-ai",
    api_key=hf_token,
)


# ============================
#   FUNÇÕES PRINCIPAIS
# ============================

def aspect_ratio_preview(img, max_dim=256):
    img_copy = img.copy()
    img_copy.thumbnail((max_dim, max_dim))
    return img_copy


def generate_images(img1, img2, img3, prompt="", n_outputs=1, aspect_ratio="4:3"):
    if token_missing:
        return []

    images = [img for img in [img1, img2, img3] if img is not None]
    if not images:
        return [Image.new("RGB", (256, 256), "gray")]

    n_outputs = min(int(n_outputs), 3)
    outputs = []

    user_prompt = f"{prompt.strip()} | {aspect_ratio}" if prompt.strip() else aspect_ratio

    for _ in range(n_outputs):
        try:
            image = client.image_to_image(
                images[0],
                prompt=user_prompt,
                model="Qwen/Qwen-Image-Edit"
            )
        except Exception:
            traceback.print_exc()
            image = Image.new("RGB", (256, 256), "gray")

        outputs.append(aspect_ratio_preview(image))

    return outputs


# ============================
#   INTERFACE GRADIO
# ============================

with gr.Blocks(title="Qwen Image Multi-Edit By Chalenger") as demo:

    if token_missing:
        gr.Markdown("⚠️ **HF_TOKEN not set!** Adicione o token em *Settings → Secrets*.")

    gr.Markdown("## Multi-Image Editor for Qwen Image Edit — Max 3 Images")

    with gr.Row():
        img1 = gr.Image(label="Image Upload 1", type="pil")
        img2 = gr.Image(label="Image Upload 2", type="pil")
        img3 = gr.Image(label="Image Upload 3", type="pil")

    prompt = gr.Textbox(
        label="Prompt (opcional)",
        placeholder="Description / Descrição",
        lines=2
    )

    nout = gr.Dropdown(
        label="Outputs / Saída",
        choices=[1, 2, 3],
        value=1
    )

    aspect_ratio_state = gr.State("4:3")

    with gr.Row():
        btn_4_3 = gr.Button("📺 4:3 CRT")
        btn_16_9 = gr.Button("🖥️ 16:9 LED")
        btn_9_16 = gr.Button("📱 9:16 Phone")
        btn_1_1 = gr.Button("⬛ 1:1 Square")

    btn_4_3.click(lambda: "4:3", None, aspect_ratio_state)
    btn_16_9.click(lambda: "16:9", None, aspect_ratio_state)
    btn_9_16.click(lambda: "9:16", None, aspect_ratio_state)
    btn_1_1.click(lambda: "1:1", None, aspect_ratio_state)

    with gr.Row():
        play_btn = gr.Button("▶️ Start Engine")
        stop_btn = gr.Button("⏹️ Kill Tasks")

    gallery = gr.Gallery(
        label="Outputs / Saídas",
        columns=2,
        height="500"
    )

    play_btn.click(
        fn=generate_images,
        inputs=[img1, img2, img3, prompt, nout, aspect_ratio_state],
        outputs=gallery
    )

    stop_btn.click(lambda: None, None, gallery)


# ============================
#   LAUNCH COMPATÍVEL
# ============================

if __name__ == "__main__":
    try:
        demo.queue()  # ativa fila sem depender do launch
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            debug=True
        )
    except Exception:
        traceback.print_exc()
        sys.exit(1)