import gradio as gr
import numpy as np
import random
import torch
import spaces

from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler
from optimization import optimize_pipeline_
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

import math
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import os

# --------------------------
# 🔹 CONFIGURACIÓN DEL TOKEN 🔹
# --------------------------
ACCESS_TOKEN = os.environ.get("Token_Nuevo")

# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "linoyts/Qwen-Image-Edit-Rapid-AIO",
        subfolder="transformer",
        torch_dtype=dtype,
        device_map="cuda",
    ),
    torch_dtype=dtype,
).to(device)

pipe.load_lora_weights(
    "eigen-ai-labs/eigen-banana-qwen-image-edit",
    weight_name="eigen-banana-qwen-image-edit-fp16-lora.safetensors",
    adapter_name="eigen-banana",
)
pipe.set_adapters(["eigen-banana"], adapter_weights=[1.0])
pipe.fuse_lora(adapter_names=["eigen-banana"], lora_scale=1.0)
pipe.unload_lora_weights()

pipe.transformer.__class__ = QwenImageTransformer2DModel
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())

optimize_pipeline_(
    pipe,
    image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))],
    prompt="prompt",
)

MAX_SEED = np.iinfo(np.int32).max

@spaces.GPU
def convert_to_anime(
    image,
    prompt,
    seed,
    randomize_seed,
    true_guidance_scale,
    num_inference_steps,
    height,
    width,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt or prompt.strip() == "":
        prompt = "edit"

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)

    pil_images = []
    if image is not None:
        if isinstance(image, Image.Image):
            pil_images.append(image.convert("RGB"))
        elif hasattr(image, "name"):
            pil_images.append(Image.open(image.name).convert("RGB"))

    if len(pil_images) == 0:
        raise gr.Error("Please upload an image first.")

    result = pipe(
        image=pil_images,
        prompt=prompt,
        height=height if height != 0 else None,
        width=width if width != 0 else None,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=1,
    ).images[0]

    return result, seed

# --- UI ---
css = '''
#col-container { 
    max-width: 900px; 
    margin: 0 auto; 
    padding: 2rem;
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
}
.gradio-container.light {
    background: linear-gradient(to bottom, #f5f5f7, #ffffff);
}
.gradio-container.dark {
    background: linear-gradient(to bottom, #1a1a1a, #0d0d0d);
}
#title {
    text-align: center;
    font-size: 2.5rem;
    font-weight: 600;
    margin-bottom: 0.5rem;
}
'''

def update_dimensions_on_upload(image):
    if image is None:
        return 1024, 1024

    original_width, original_height = image.size

    if original_width > original_height:
        new_width = 1024
        new_height = int((original_height / original_width) * new_width)
    else:
        new_height = 1024
        new_width = int((original_width / original_height) * new_height)

    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8

    return new_width, new_height

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:

    # ---------- LOGIN PANEL ----------
    with gr.Column():
        gr.Markdown("## 🔒 Acceso restringido")
        token_input = gr.Textbox(
            label="Introduce tu token",
            type="password",
            placeholder="Token",
        )
        status_text = gr.Textbox(
            label="Estado",
            interactive=False,
            value="",
        )
        login_btn = gr.Button("Ingresar")

    # ---------- APP AREA (oculta hasta login) ----------
    with gr.Column(visible=False) as app_area:
        with gr.Column(elem_id="col-container"):
            gr.Markdown(
                "# 🍌 Eigen-Banana-Qwen-Image-Edit: Fast Image Editing with Qwen-Image-Edit LoRA",
                elem_id="title",
            )

            with gr.Row():
                with gr.Column(scale=1):
                    image = gr.Image(label="Upload Photo", type="pil")
                    prompt = gr.Textbox(label="Prompt", value="Edit")

                    with gr.Accordion("⚙️ Advanced Settings", open=False):
                        seed = gr.Slider(0, MAX_SEED, value=0)
                        randomize_seed = gr.Checkbox(value=True)
                        true_guidance_scale = gr.Slider(1.0, 10.0, value=1.0)
                        num_inference_steps = gr.Slider(1, 40, value=4)
                        height = gr.Slider(256, 2048, step=8, value=1024, visible=False)
                        width = gr.Slider(256, 2048, step=8, value=1024, visible=False)

                    convert_btn = gr.Button("Edit", variant="primary")

                with gr.Column(scale=1):
                    result = gr.Image(label="Result")

        convert_btn.click(
            fn=convert_to_anime,
            inputs=[
                image,
                prompt,
                seed,
                randomize_seed,
                true_guidance_scale,
                num_inference_steps,
                height,
                width,
            ],
            outputs=[result, seed],
        )

        image.upload(
            fn=update_dimensions_on_upload,
            inputs=[image],
            outputs=[width, height],
        )

    # ---------- LOGIN LOGIC ----------
    def check_token_func(token_value):
        if token_value == ACCESS_TOKEN:
            return gr.update(visible=True), "Token correcto. Acceso concedido."
        else:
            return gr.update(visible=False), "Token incorrecto. Acceso denegado."

    login_btn.click(
        fn=check_token_func,
        inputs=token_input,
        outputs=[app_area, status_text],
    )
    
demo.launch()