"""P-Q-Image-2.1-Remove-Object Studio — gr.Server studio + gr.Workflow canvas.

  - "/"          Studio: custom UI (studio.html) over the `remove` API endpoint
  - "/workflow"  Workflow canvas: node graph from workflow.json (pq_remove)

Text-guided masks: with no painted mask, SAM 3 can find the object the instruction names
(text_mask.py) and the removal runs on its outline, in the same GPU call.

The removal LoRA runs as its model card says: the few-step variants on top of PrunaAI's distillation
LoRA (at 1.0) at strength 0.7, the base variant alone at 40 steps; no CFG in all cases.
"""

import base64
import io
import json
import os
import random
import time
from urllib.parse import unquote, urlparse

import gradio as gr
import numpy as np
import spaces
import torch
import uvicorn
from fastapi import HTTPException
from fastapi.responses import FileResponse, HTMLResponse
from gradio.utils import get_upload_folder
from gradio_client import utils as client_utils
from PIL import Image, ImageOps
from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline

from draw_outline import draw_outline, outline_prompt
from examples import EXAMPLES
from text_mask import TextMasker

APP_DIR = os.path.dirname(os.path.abspath(__file__))
EXAMPLES_DIR = os.path.join(APP_DIR, "examples")
MAX_SEED = 2**31 - 1

REMOVAL_REPO = "PrunaAI/P-Q-Image-2.1-Remove-Object"
DISTILL_REPO = "PrunaAI/P-Q-Image-2.1"
DEFAULT_VARIANT = "6-step"
DEFAULT_INSTRUCTION = "Remove the outlined object."

# From the model card's Variants table. `max_side` caps the working resolution (side of the
# equivalent square): ~2 MP like the card's examples for the few-step variants, ~1 MP for the
# 40-step base variant to keep it within a ZeroGPU slot. `gpu_seconds` is the ZeroGPU reservation.
VARIANTS = {
    "6-step": {
        "removal_weight": "p_q_image_2.1_remove_object_6_steps_v0.1.safetensors",
        "distill_weight": "p_q_image_2.1_6_steps_v0.2.safetensors",
        "strength": 0.7,
        # Unshifted sigmas: the scheduler shifts them by image size.
        "sigmas": [1.0, 0.925, 0.85, 0.625, 0.375, 0.125],
        "dynamic_shifting": True,
        "max_side": 1414,
        "gpu_seconds": 60,
    },
    "8-step": {
        "removal_weight": "p_q_image_2.1_remove_object_8_steps_v0.1.safetensors",
        "distill_weight": "p_q_image_2.1_8_steps_v0.1.safetensors",
        "strength": 0.7,
        # Sigmas used exactly as given.
        "sigmas": [1.0, 14 / 15, 6 / 7, 10 / 13, 2 / 3, 6 / 11, 0.4, 2 / 9],
        "dynamic_shifting": False,
        "max_side": 1414,
        "gpu_seconds": 60,
    },
    "base": {
        "removal_weight": "p_q_image_2.1_remove_object_base_v0.1.safetensors",
        "distill_weight": None,
        "strength": 1.0,
        "sigmas": None,  # 40 steps, the pipeline's own scheduler
        "dynamic_shifting": None,
        "max_side": 1024,
        "gpu_seconds": 150,
    },
}
VARIANT_CHOICES = list(VARIANTS)
VARIANT_ALIASES = {
    "6-step": "6-step", "6-steps": "6-step", "6step": "6-step", "6": "6-step",
    "8-step": "8-step", "8-steps": "8-step", "8step": "8-step", "8": "8-step",
    "base": "base", "40-step": "base", "40-steps": "base", "40": "base",
}

# ------------------------------------------------------------------ pipeline

pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16)
for name, spec in VARIANTS.items():
    spec["removal_adapter"] = f"removal_{name}"
    pipe.load_lora_weights(REMOVAL_REPO, weight_name=spec["removal_weight"], adapter_name=spec["removal_adapter"])
    if spec["distill_weight"]:
        spec["distill_adapter"] = f"distill_{name}"
        pipe.load_lora_weights(DISTILL_REPO, weight_name=spec["distill_weight"], adapter_name=spec["distill_adapter"])

# One scheduler per variant, built once: the pipeline's own for base, and for the few-step variants
# the schedule the distillation LoRA was trained with (no terminal stretch).
SCHEDULERS = {"base": pipe.scheduler}
for name, spec in VARIANTS.items():
    if spec["distill_weight"]:
        SCHEDULERS[name] = FlowMatchEulerDiscreteScheduler.from_config(
            pipe.scheduler.config, use_dynamic_shifting=spec["dynamic_shifting"], shift=1.0, shift_terminal=None
        )

_active_variant = None


def _activate(name: str) -> dict:
    """Switch adapters and scheduler to `name`; a no-op when it is already active."""
    global _active_variant
    spec = VARIANTS[name]
    if _active_variant != name:
        if spec["distill_weight"]:
            pipe.set_adapters([spec["distill_adapter"], spec["removal_adapter"]], [1.0, spec["strength"]])
        else:
            pipe.set_adapters([spec["removal_adapter"]], [spec["strength"]])
        pipe.scheduler = SCHEDULERS[name]
        _active_variant = name
    return spec


_activate(DEFAULT_VARIANT)
pipe = pipe.to("cuda")
# Decoding a ~2 MP result in one piece runs out of memory on a ZeroGPU slot; tiles of 1024 px leave
# results up to 1024 px a side decoded exactly as before.
pipe.vae.enable_tiling(
    tile_sample_min_height=1024, tile_sample_min_width=1024,
    tile_sample_stride_height=768, tile_sample_stride_width=768,
)

# SAM 3 for text-guided masks. facebook/sam3 is gated: without access (HF_TOKEN) the Space still
# runs, with painted masks only. It stays in CPU memory and visits the GPU only to find a mask, so
# the removal keeps the GPU memory it had without it.
try:
    masker = TextMasker(device="cpu")
except Exception as error:  # noqa: BLE001
    print(f"Text-guided masks disabled, SAM 3 did not load: {error}")
    masker = None

# SAM 3 works at 1008 px; segmenting a larger copy only costs time.
FIND_MAX_SIDE = 1536


def _resolve_variant(variant) -> str:
    if variant is None or variant == "":
        return DEFAULT_VARIANT
    key = "-".join(str(variant).strip().lower().replace("_", "-").split())
    if key not in VARIANT_ALIASES:
        raise gr.Error(f"Variant must be one of: {', '.join(VARIANT_CHOICES)}.")
    return VARIANT_ALIASES[key]


# ------------------------------------------------------------------ rendering


def _find(image: Image.Image, text: str):
    """SAM 3 mask (L, white = the object) of what `text` names in `image`, at its size, or None."""
    if masker is None or not text:
        return None
    small = image
    if max(image.size) > FIND_MAX_SIDE:
        scale = FIND_MAX_SIDE / max(image.size)
        small = image.resize((round(image.width * scale), round(image.height * scale)), Image.LANCZOS)
    try:
        masker.model.to("cuda")
        mask, _ = masker(small, text)
    finally:
        masker.model.to("cpu")
        torch.cuda.empty_cache()
    return None if mask is None else mask.resize(image.size, Image.NEAREST)


def _gpu_seconds(prompt, variant, seed, image, resolution, find=""):
    return VARIANTS[variant]["gpu_seconds"] + (10 if find else 0)


@spaces.GPU(duration=_gpu_seconds)
def _render(prompt: str, variant: str, seed: int, image: Image.Image, resolution: int, find: str = ""):
    """The removal. With `find`, SAM 3 first looks for the object it names; when found, the photo
    gets its red outline and the prompt the outline sentence. Returns (image, seconds, mask or None)."""
    spec = _activate(variant)
    mask = _find(image, find) if find else None
    if mask is not None:
        image = draw_outline(image, mask)
        prompt = outline_prompt(prompt)
    torch.cuda.synchronize()
    start = time.perf_counter()
    result = pipe(
        prompt=prompt,
        image=[image],
        sigmas=spec["sigmas"],
        num_inference_steps=len(spec["sigmas"]) if spec["sigmas"] else 40,
        true_cfg_scale=1.0,
        output_resolution=resolution,
        generator=torch.Generator("cuda").manual_seed(seed),
    ).images[0]
    torch.cuda.synchronize()
    return result, time.perf_counter() - start, mask


@spaces.GPU(duration=20)
def _find_gpu(image: Image.Image, text: str):
    return _find(image, text)


def _canvas(ratio: float, side: int) -> tuple[int, int]:
    """The pipeline's own sizing: area side², this aspect ratio, sides rounded to multiples of 32."""
    width = (side**2 * ratio) ** 0.5
    return round(width / 32) * 32, round(width / ratio / 32) * 32


def _remove(photo: Image.Image, mask, instruction: str, variant, seed: int, auto_mask: bool = False):
    """Object removal as the model card does it: work at the photo's area (sides rounded to 32,
    capped per variant), draw the red outline at that size with the matching sentence, and resize
    the result back to the photo's aspect ratio. Without a mask (or an empty one) it runs on the
    instruction alone, or, with `auto_mask`, on the outline of what SAM 3 finds from the
    instruction. Returns (image, gpu seconds, variant, mode, found mask at the photo's size or None)."""
    name = _resolve_variant(variant)
    instruction = (instruction or "").strip()
    area = photo.width * photo.height
    side = min(round(area**0.5), VARIANTS[name]["max_side"])
    work_size = _canvas(photo.width / photo.height, side)
    work = photo.resize(work_size, Image.LANCZOS)
    mode = "prompt only"
    if mask is not None:
        mask = mask.resize(work_size, Image.NEAREST)
        if (np.array(mask) > 127).any():
            work = draw_outline(work, mask)
            instruction = outline_prompt(instruction or DEFAULT_INSTRUCTION)
            mode = "outline"
    if not instruction:
        raise gr.Error("Paint over the object to remove, or describe it.")
    find = instruction if auto_mask and mode == "prompt only" and masker is not None else ""
    try:
        result, seconds, found = _render(instruction, name, seed, work, side, find)
    except (RuntimeError, ValueError) as error:
        raise gr.Error(str(error).split("\n")[0])
    if find:
        mode = "text mask (SAM 3)" if found is not None else "prompt only, SAM 3 found nothing"
    if found is not None:
        found = found.resize(photo.size, Image.NEAREST)
    scale = side / area**0.5
    out_size = (round(photo.width * scale), round(photo.height * scale)) if scale < 1 else photo.size
    return result.resize(out_size, Image.LANCZOS), seconds, name, mode, found


# ------------------------------------------------------------------ helpers


def _seed(seed) -> int:
    try:
        return min(max(int(float(seed)), 0), MAX_SEED)
    except (TypeError, ValueError):
        return 42


def _resolve_seed(seed, randomize) -> int:
    return random.randint(0, MAX_SEED) if randomize else _seed(seed)


def _requested_path(value) -> str:
    """Turn whatever the canvas/API client sends into a path."""
    if isinstance(value, str):
        return value
    if not isinstance(value, dict):
        raise gr.Error(f"Expected an image, got {type(value).__name__}")
    if path := value.get("path"):
        return path
    url = value.get("url") or ""
    if "/gradio_api/file=" in url:
        return unquote(url.split("/gradio_api/file=", 1)[1])
    if url.startswith(("http://", "https://")):
        raise gr.Error("Upload the image instead of passing a URL.")
    return unquote(urlparse(url).path)


def _local_path(value) -> str:
    """A path this process may open: only files inside gradio's upload folder,
    where /gradio_api/upload puts uploads and _publish puts results."""
    path = os.path.realpath(_requested_path(value))
    folder = os.path.realpath(get_upload_folder())
    if os.path.commonpath([path, folder]) != folder:
        raise gr.Error("Upload the image first.")
    return path


def _empty_media(value) -> bool:
    """True when the canvas/API sent no file: None, '', or a dict with no path/url."""
    if value is None or value == "":
        return True
    if isinstance(value, Image.Image):
        return False
    if isinstance(value, str):
        return not value.strip()
    if isinstance(value, dict):
        return not (value.get("path") or value.get("url"))
    return True


def _as_instruction(value) -> str:
    if value is None or isinstance(value, bool):
        return ""
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        for key in ("text", "value", "prompt", "out", "in"):
            item = value.get(key)
            if isinstance(item, str) and item.strip():
                return item
        return ""
    return str(value)


def _open_rgb(value) -> Image.Image:
    if isinstance(value, Image.Image):
        return value.convert("RGB")
    return ImageOps.exif_transpose(Image.open(_local_path(value))).convert("RGB")


def _open_mask(value):
    """A mask as grayscale, white = the object. Transparent pixels count as black. None when absent."""
    if _empty_media(value):
        return None
    try:
        image = value if isinstance(value, Image.Image) else ImageOps.exif_transpose(Image.open(_local_path(value)))
    except (gr.Error, OSError, ValueError, FileNotFoundError):
        return None
    if "A" in image.getbands() or image.mode == "P":
        image = Image.alpha_composite(Image.new("RGBA", image.size, (0, 0, 0, 255)), image.convert("RGBA"))
    return image.convert("L")


def _as_data_uri(image: Image.Image) -> str:
    """Inline a result image so the studio page never needs a file route."""
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode()


def _publish(image: Image.Image) -> str:
    """Save a result into gradio's upload folder, which it already serves, so the workflow canvas can render it."""
    folder = get_upload_folder()
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, f"pruna_{os.urandom(8).hex()}.png")
    image.save(path)
    return path


def _as_media(path: str) -> dict:
    """The dict shape the workflow canvas renders media from."""
    return {
        "path": path,
        "url": f"/gradio_api/file={client_utils.encode_file_path(path)}",
        "is_file": True,
    }


# ------------------------------------------------------------------ workflow fn


def pq_remove(image, prompt="", variant: str = DEFAULT_VARIANT, seed: str = "42", mask=None) -> dict:
    """Remove an object with a mask outline, a text instruction, or both (no CFG)."""
    if _empty_media(image):
        raise gr.Error("Connect a source image first.")
    result, _, _, _, _ = _remove(
        _open_rgb(image), _open_mask(mask), _as_instruction(prompt), variant, _seed(seed)
    )
    return _as_media(_publish(result))


# ------------------------------------------------------------------ studio server

server = gr.Server(title="P-Q-Image-2.1-Remove-Object")

EXAMPLE_FILES = {f"{eid}_{kind}" for eid, _, _ in EXAMPLES for kind in ("photo.jpg", "mask.png")}


@server.get("/", response_class=HTMLResponse)
def _studio_page():
    with open(os.path.join(APP_DIR, "studio.html"), encoding="utf-8") as handle:
        page = handle.read()
    config = {
        "examples": [{"id": eid, "label": label, "prompt": prompt} for eid, label, prompt in EXAMPLES],
        "variants": VARIANT_CHOICES,
        "defaultVariant": DEFAULT_VARIANT,
        "textMask": masker is not None,
    }
    return page.replace("__STUDIO_CONFIG__", json.dumps(config).replace("</", "<\\/"))


@server.get("/pruna-logo.png")
def _logo():
    return FileResponse(os.path.join(APP_DIR, "pruna-logo.png"), media_type="image/png")


@server.get("/examples/{name}")
def _example_file(name: str):
    if name not in EXAMPLE_FILES:
        raise HTTPException(status_code=404)
    return FileResponse(os.path.join(EXAMPLES_DIR, name))


@server.api(name="remove", concurrency_limit=1)
def api_remove(
    source: str,
    mask: str = "",
    instruction: str = "",
    variant: str = DEFAULT_VARIANT,
    seed: float = 42,
    randomize: bool = True,
    auto_mask: bool = False,
) -> dict:
    """Remove the masked object (white = remove) with P-Q-Image-2.1-Remove-Object. With no mask and
    `auto_mask`, SAM 3 finds the object the instruction names and its mask is returned as `mask_data`."""
    if not source:
        raise gr.Error("Upload a photo first.")
    seed_value = _resolve_seed(seed, randomize)
    image, seconds, name, mode, found = _remove(
        _open_rgb(source), _open_mask(mask), instruction, variant, seed_value, auto_mask
    )
    info = f"{image.width}×{image.height}, seed {seed_value} · {name} · {mode} · {seconds:.1f} s on the GPU"
    payload = {"image_data": _as_data_uri(image), "seed": seed_value, "info": info}
    if found is not None:
        payload["mask_data"] = _as_data_uri(found)
    return payload


@server.api(name="find_mask", concurrency_limit=2)
def api_find_mask(source: str, text: str = "") -> dict:
    """Mask (white = the object) of what `text` names in the photo, found with SAM 3."""
    if masker is None:
        raise gr.Error("Text-guided masks are not available on this Space.")
    if not source:
        raise gr.Error("Upload a photo first.")
    if not (text or "").strip():
        raise gr.Error("Write what to remove first.")
    mask = _find_gpu(_open_rgb(source), text.strip())
    if mask is None:
        raise gr.Error("SAM 3 found nothing matching the instruction. Try naming the object more simply.")
    return {"mask_data": _as_data_uri(mask)}


# The @server.api endpoints are deferred until Server.launch(), which also fires
# the ZeroGPU registration hook (pack weights + startup report) through the
# patched Blocks.launch path. Launch on a loopback port for registration, then
# serve the same app via uvicorn.
_port_env = os.environ.get("GRADIO_SERVER_PORT")
os.environ["GRADIO_SERVER_PORT"] = "7861"
try:
    server.launch(
        prevent_thread_lock=True,
        server_name="127.0.0.1",
        server_port=7861,
        quiet=True,
        show_error=False,
        ssr_mode=False,
    )
finally:
    if _port_env is None:
        del os.environ["GRADIO_SERVER_PORT"]
    else:
        os.environ["GRADIO_SERVER_PORT"] = _port_env

try:
    server.close()
except Exception:
    pass

# ------------------------------------------------------------------ workflow canvas
# Mounted only after the registration launch above. Mounting first starts the
# workflow's queue on the loopback server's event loop, and close() then strands
# it: jobs are accepted on port 7860 but never run.

workflow_demo = gr.Workflow(graph="workflow.json", bind={"pq_remove": pq_remove})
workflow_demo.title = "P-Q-Image-2.1-Remove-Object Workflow"

server = gr.mount_gradio_app(server, workflow_demo, path="/workflow", ssr_mode=False)

if __name__ == "__main__":
    uvicorn.run(server, host="0.0.0.0", port=7860)