"""
Specs Remover — single-purpose app.

Upload a photo -> glasses are automatically removed -> download the result.
"""

import os
import random
import time

import cv2
import gradio as gr
import mediapipe as mp
import numpy as np
import spaces
import torch
from PIL import Image, ImageFilter

from diffusers import QwenImageEditPlusPipeline
from diffusers.models import QwenImageTransformer2DModel


GLASSES_REMOVAL_PROMPT = (
    "Remove the eyeglasses from the image while preserving the background "
    "and remaining elements, maintaining realism and original details. "
    "Keep the eyes, eyebrows, eyelashes, nose, skin tone, lighting, "
    "hairstyle, and facial identity exactly the same and fully natural."
)
NEGATIVE_PROMPT = " "

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

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

print("Loading transformer...")
transformer = QwenImageTransformer2DModel.from_pretrained(
    "Qwen/Qwen-Image-Edit-2511",
    subfolder="transformer",
    torch_dtype=dtype,
)

print("Loading pipeline...")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2511",
    transformer=transformer,
    torch_dtype=dtype,
).to(device)

print("Loading Lightning LoRA...")
pipe.load_lora_weights(
    "lightx2v/Qwen-Image-Edit-2511-Lightning",
    weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors",
    adapter_name="lightning",
)

print("Loading Object-Remover LoRA...")
pipe.load_lora_weights(
    "prithivMLmods/Qwen-Image-Edit-2511-Object-Remover",
    adapter_name="object_remover",
)

pipe.set_adapters(["lightning", "object_remover"], adapter_weights=[1.0, 1.0])
pipe.fuse_lora(adapter_names=["lightning", "object_remover"])

print("Pipeline ready.")


_FACE_MESH = mp.solutions.face_mesh.FaceMesh(
    static_image_mode=True,
    max_num_faces=1,
    refine_landmarks=True,
    min_detection_confidence=0.5,
)

_GLASSES_REGION_INDICES = [
    70, 63, 105, 66, 107, 55, 65, 52, 53, 46,
    300, 293, 334, 296, 336, 285, 295, 282, 283, 276,
    33, 133, 160, 159, 158, 157, 173, 246,
    362, 263, 387, 386, 385, 384, 398, 466,
    6, 197, 195, 5, 4,
    127, 234, 93, 356, 454, 323,
]


def build_glasses_region_mask(image_pil):
    image_bgr = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    h, w = image_bgr.shape[:2]

    results = _FACE_MESH.process(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
    if not results.multi_face_landmarks:
        return None

    landmarks = results.multi_face_landmarks[0].landmark
    points = np.array(
        [[landmarks[i].x * w, landmarks[i].y * h] for i in _GLASSES_REGION_INDICES],
        dtype=np.float32,
    )

    hull = cv2.convexHull(points.astype(np.int32))
    mask = np.zeros((h, w), dtype=np.uint8)
    cv2.fillConvexPoly(mask, hull, 255)

    pad = max(8, int(0.03 * min(h, w)))
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (pad * 2, pad * 2))
    mask = cv2.dilate(mask, kernel)

    mask_pil = Image.fromarray(mask, mode="L")
    feather_radius = max(4, int(0.015 * min(h, w)))
    mask_pil = mask_pil.filter(ImageFilter.GaussianBlur(radius=feather_radius))

    return mask_pil


def blend_with_original(original_pil, ai_output_pil):
    if ai_output_pil.size != original_pil.size:
        ai_output_pil = ai_output_pil.resize(original_pil.size, Image.LANCZOS)

    mask = build_glasses_region_mask(original_pil)
    if mask is None:
        return ai_output_pil

    return Image.composite(ai_output_pil, original_pil, mask)


@spaces.GPU(duration=300)
def remove_glasses(
    input_image,
    seed=42,
    randomize_seed=True,
    num_inference_steps=4,
    progress=gr.Progress(track_tqdm=True),
):
    print("remove_glasses called")
    if input_image is None:
        raise gr.Error("Please upload a photo first.")

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

        generator = torch.Generator(device=device).manual_seed(seed)

        if isinstance(input_image, str):
            pil_image = Image.open(input_image).convert("RGB")
        else:
            pil_image = input_image.convert("RGB")

        progress(0.1, desc="Removing glasses...")

        result = pipe(
            image=[pil_image],
            prompt=GLASSES_REMOVAL_PROMPT,
            negative_prompt=NEGATIVE_PROMPT,
            num_inference_steps=num_inference_steps,
            generator=generator,
            true_cfg_scale=1.0,
            num_images_per_prompt=1,
        ).images

        output_image = result[0]

        progress(0.9, desc="Preserving original details...")
        output_image = blend_with_original(pil_image, output_image)

        os.makedirs("outputs", exist_ok=True)
        output_path = f"outputs/result_{seed}_{int(time.time() * 1000)}.png"
        output_image.save(output_path)

        print("remove_glasses done")
        return output_path

    except Exception as e:
        import traceback
        traceback.print_exc()
        raise gr.Error(f"Generation failed: {e}")


css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
#title {
    text-align: center;
}
"""

with gr.Blocks(css=css, title="Specs Remover") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 👓 Specs Remover
            Upload a photo. Glasses are automatically removed —
            everything else about the face stays natural.
            """,
            elem_id="title",
        )

        input_image = gr.Image(
            label="Upload a photo",
            type="filepath",
            interactive=True,
        )

        run_button = gr.Button("Remove Glasses", variant="primary")

        output_image = gr.Image(label="Result", type="filepath")

        with gr.Accordion("Advanced (optional)", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            num_inference_steps = gr.Slider(
                label="Inference steps",
                minimum=1,
                maximum=40,
                step=1,
                value=4,
            )

    run_button.click(
        fn=remove_glasses,
        inputs=[input_image, seed, randomize_seed, num_inference_steps],
        outputs=[output_image],
    )


if __name__ == "__main__":
    demo.launch()