"""
Specs Remover — single-purpose app.

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

Built on the same qwenimage/ pipeline as the duplicated FireRed Space
(fast, FA3-accelerated). The edit prompt is fixed internally — the user
never sees or edits it.
"""

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 qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3


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(
    "FireRedTeam/FireRed-Image-Edit-1.1",
    subfolder="transformer",
    torch_dtype=dtype,
)

print("Loading pipeline...")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "FireRedTeam/FireRed-Image-Edit-1.1",
    transformer=transformer,
    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}")

print("Pipeline ready.")


# --- Face-region masking for identity-preserving blend ---
_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,        # left eyebrow
    300, 293, 334, 296, 336, 285, 295, 282, 283, 276,  # right eyebrow
    33, 133, 160, 159, 158, 157, 173, 246,            # left eye
    362, 263, 387, 386, 385, 384, 398, 466,           # right eye
    6, 197, 195, 5, 4,                                 # nose bridge
    127, 234, 93, 356, 454, 323,                       # temples
]


def build_glasses_region_mask(image_pil: Image.Image):
    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: Image.Image, ai_output_pil: Image.Image):
    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=60)
def remove_glasses(
    input_image,
    seed=42,
    randomize_seed=True,
    num_inference_steps=8,
    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=8,
            )

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


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