import os
import gc
import gradio as gr
from gradio import Server
from fastapi.responses import HTMLResponse
import numpy as np
import spaces
import torch
import random
import base64
import json
from io import BytesIO
from PIL import Image

from diffusers import FlowMatchEulerDiscreteScheduler
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

MAX_SEED = np.iinfo(np.int32).max
LANCZOS = getattr(Image, "Resampling", Image).LANCZOS

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16

print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("torch.__version__ =", torch.__version__)
print("Using device:", device)

print("Loading FLUX.2 Klein 9B model base...")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
        torch_dtype=dtype,
        device_map="cuda",
    ),
    torch_dtype=dtype,
).to(device)

try:
    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
    print("Flash Attention 3 Processor set successfully.")
except Exception as e:
    print(f"Warning: Could not set FA3 processor: {e}")

# ── LoRA adapter registry ──────────────────────────────────────────────────────
ADAPTER_SPECS = {
    "Qwen-Image-Edit-2511-Object-Adder": {
        "repo": "prithivMLmods/Qwen-Image-Edit-2511-Object-Adder",
        "weights": "Qwen-Image-Edit-2511-Object-Adder.safetensors",
        "adapter_name": "object-adder"
    },
    "Qwen-Image-Edit-2511-Object-Remover": {
        "repo": "prithivMLmods/Qwen-Image-Edit-2511-Object-Remover",
        "weights": "Qwen-Image-Edit-2511-Object-Remover.safetensors",
        "adapter_name": "object-remover"
    },
    "QIE-2511-Object-Remover-v2": {
        "repo": "prithivMLmods/QIE-2511-Object-Remover-v2",
        "weights": "Qwen-Image-Edit-2511-Object-Remover-v2-9200.safetensors",
        "adapter_name": "object-remover"
    },
    "Zoom-Master": {
        "repo": "prithivMLmods/QIE-2511-Zoom-Master",
        "weights": "Qwen-Image-Edit-2511-Zoom-Master-8800.safetensors",
        "adapter_name": "zoom-master"
    },
    "Extract-Outfit": {
        "repo": "prithivMLmods/QIE-2511-Extract-Outfit",
        "weights": "QIE-2511-Extract-Outfit-4200.safetensors",
        "adapter_name": "extract-outfit"
    },    
    "Outfit-Design-Layout": {
        "repo": "prithivMLmods/QIE-2511-Outfit-Design-Layout",
        "weights": "QIE-2511-Outfit-Design-Layout-3000.safetensors",
        "adapter_name": "layout-outfit"
    }, 
}

LOADED_ADAPTERS: set = set()
ADAPTER_NAMES = list(ADAPTER_SPECS.keys())

EXAMPLES_CONFIG = [
    {"images": ["examples/D.jpg"], "prompt": "Add the batman logo to the image while preserving the background lighting and surrounding elements maintaining realism and original details.", "lora": "Qwen-Image-Edit-2511-Object-Adder"},
    {"images": ["examples/A.jpg"], "prompt": "Add the slim rectangular transparent frame sunglasses to the image while preserving the background lighting and surrounding elements maintaining realism and original details.", "lora": "Qwen-Image-Edit-2511-Object-Adder"},
    {"images": ["examples/B.jpeg"], "prompt": "Remove the necklace and goggles from the image while preserving the background and remaining elements, maintaining realism and original details.", "lora": "Qwen-Image-Edit-2511-Object-Remover"},
    {"images": ["examples/DL2.jpg"], "prompt": "add the nike tick design inside the red marked area.", "lora": "Outfit-Design-Layout"},
    {"images": ["examples/DL1.jpg"], "prompt": "add the akatsuki cloud design inside the red marked area.", "lora": "Outfit-Design-Layout"},
    {"images": ["examples/C.png"], "prompt": "Add the leather cowboy cap to the image while preserving the background lighting and surrounding elements maintaining realism and original details.", "lora": "Qwen-Image-Edit-2511-Object-Adder"},
    {"images": ["examples/ZM.jpg"], "prompt": "Zoom into the red highlighted area.", "lora": "Zoom-Master"},
    {"images": ["examples/OBJ1.jpg"], "prompt": "Remove the red highlighted object from the scene.", "lora": "QIE-2511-Object-Remover-v2"},
    {"images": ["examples/OBJ2.jpg"], "prompt": "Remove the red highlighted object from the scene.", "lora": "QIE-2511-Object-Remover-v2"},
    {"images": ["examples/OE.jpg"], "prompt": "Extract the clothing and create a flat mockup.", "lora": "Extract-Outfit"},
]

def make_thumb_b64(path, max_dim=220):
    if not os.path.exists(path):
        return ""
    try:
        img = Image.open(path).convert("RGB")
        img.thumbnail((max_dim, max_dim), LANCZOS)
        buf = BytesIO()
        img.save(buf, format="JPEG", quality=65)
        return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
    except Exception as e:
        return ""

def encode_full_image(path):
    if not os.path.exists(path):
        return ""
    try:
        with open(path, "rb") as f:
            data = f.read()
        ext = path.rsplit(".", 1)[-1].lower()
        mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
        return f"data:{mime};base64,{base64.b64encode(data).decode()}"
    except Exception as e:
        return ""

def build_client_config():
    examples = []
    for i, ex in enumerate(EXAMPLES_CONFIG):
        examples.append({
            "idx": i,
            "thumbs": [make_thumb_b64(p) for p in ex["images"]],
            "n_images": len(ex["images"]),
            "lora": ex["lora"],
            "prompt": ex["prompt"],
        })
    return {
        "loras": ADAPTER_NAMES,
        "default_lora": "Qwen-Image-Edit-2511-Object-Adder",
        "examples": examples,
    }

print("Building client config (example thumbnails)…")
CLIENT_CONFIG = build_client_config()
print(f"Built config with {len(EXAMPLES_CONFIG)} examples and {len(ADAPTER_NAMES)} LoRAs.")

def b64_to_pil_list(b64_json_str):
    if not b64_json_str or b64_json_str.strip() in ("", "[]"):
        return []
    try:
        b64_list = json.loads(b64_json_str)
    except Exception:
        return []
    pil_images = []
    for b64_str in b64_list:
        if not b64_str or not isinstance(b64_str, str):
            continue
        try:
            if b64_str.startswith("data:image"):
                _, data = b64_str.split(",", 1)
            else:
                data = b64_str
            image_data = base64.b64decode(data)
            pil_images.append(Image.open(BytesIO(image_data)).convert("RGB"))
        except Exception as e:
            print(f"Error decoding image: {e}")
    return pil_images

def pil_to_b64_png(image: Image.Image) -> str:
    buf = BytesIO()
    image.save(buf, format="PNG")
    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"

def update_dimensions_on_upload(image):
    if image is None:
        return 1024, 1024
    w, h = image.size
    if w > h:
        nw = 1024
        nh = int(nw * h / w)
    else:
        nh = 1024
        nw = int(nh * w / h)
    return (nw // 8) * 8, (nh // 8) * 8

# ── Gradio Server (Server mode): FastAPI + Gradio queue/API engine ────────────
app = Server(title="Qwen-Image-Edit-Object-Manipulator")

@app.mcp.tool(name="edit_image")
@app.api(name="edit_image")
@spaces.GPU(size="xlarge")
def infer(
    images_b64_json: str,
    prompt: str,
    lora_adapter: str,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float,
    steps: int,
) -> dict:
    """Edit one or more images with Qwen-Image-Edit + a lazily-loaded LoRA."""
    gc.collect()
    torch.cuda.empty_cache()

    pil_images = b64_to_pil_list(images_b64_json)
    if not pil_images:
        raise gr.Error("Please upload at least one image to edit.")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter an edit prompt.")

    spec = ADAPTER_SPECS.get(lora_adapter)
    if not spec:
        raise gr.Error(f"Configuration not found for: {lora_adapter}")

    adapter_name = spec["adapter_name"]
    if adapter_name not in LOADED_ADAPTERS:
        print(f"--- Downloading and Loading Adapter: {lora_adapter} ---")
        try:
            pipe.load_lora_weights(spec["repo"], weight_name=spec["weights"], adapter_name=adapter_name)
            LOADED_ADAPTERS.add(adapter_name)
        except Exception as e:
            raise gr.Error(f"Failed to load adapter {lora_adapter}: {e}")
    else:
        print(f"--- Adapter {lora_adapter} already loaded. ---")

    pipe.set_adapters([adapter_name], adapter_weights=[1.0])

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    negative_prompt = (
        "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, "
        "extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
    )
    width, height = update_dimensions_on_upload(pil_images[0])

    try:
        result_image = pipe(
            image=pil_images,
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            generator=generator,
            true_cfg_scale=guidance_scale,
        ).images[0]
        return {"image": pil_to_b64_png(result_image), "seed": seed}
    except Exception as e:
        raise e
    finally:
        gc.collect()
        torch.cuda.empty_cache()


@app.api(name="load_example", queue=False)
def load_example(idx: float) -> dict:
    """Return base64-encoded example images + prompt + LoRA for a given example index."""
    try:
        i = int(idx)
    except (ValueError, TypeError):
        i = -1
    if i < 0 or i >= len(EXAMPLES_CONFIG):
        return {"images": [], "prompt": "", "lora": "", "names": [], "status": "error"}
    ex = EXAMPLES_CONFIG[i]
    b64_list, names = [], []
    for path in ex["images"]:
        b64 = encode_full_image(path)
        if b64:
            b64_list.append(b64)
            names.append(os.path.basename(path))
    return {"images": b64_list, "prompt": ex["prompt"], "lora": ex["lora"], "names": names, "status": "ok"}


@app.get("/api/config")
def client_config():
    """Plain FastAPI route: LoRA choices + example card data for the frontend."""
    return CLIENT_CONFIG


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


if __name__ == "__main__":
    app.launch(show_error=True, mcp_server=True)