"""Compare two uploaded images with the gated SushiUI style encoders."""

from __future__ import annotations

import json
import threading
from pathlib import Path

import gradio as gr
import spaces
import torch
from huggingface_hub import get_hf_file_metadata, hf_hub_download, hf_hub_url
from huggingface_hub.errors import GatedRepoError, HfHubHTTPError
from PIL import Image

from model import load_encoder


MODEL_REPO = "celstk/style_embedder"
MODEL_REVISION = "dbf4e0b4530876cdf150d4006625e8638bc338db"
CALIBRATION = json.loads(Path(__file__).with_name("thresholds.json").read_text(encoding="utf-8"))
MODELS: dict[str, tuple] = {}
MODEL_LOCK = threading.RLock()


def _check_image(image: Image.Image | None, name: str) -> Image.Image:
    if image is None:
        raise gr.Error(f"{name} の画像を入力してください。")
    if image.width * image.height > 40_000_000:
        raise gr.Error(f"{name} は 4,000 万画素以下の画像にしてください。")
    return image.convert("RGB")


def _require_model_access(variant: str, token: str | None) -> None:
    if not token:
        raise gr.Error("Hugging Face にログインし、モデルへのアクセス承認を受けてください。")
    url = hf_hub_url(MODEL_REPO, f"{variant}/best_style_encoder.safetensors",
                     revision=MODEL_REVISION)
    try:
        # Check every request: a shared model cache must not bypass the model's gate.
        get_hf_file_metadata(url, token=token)
    except (GatedRepoError, HfHubHTTPError) as exc:
        raise gr.Error("モデルへのアクセス権を確認できません。celstk/style_embedder でアクセスを申請し、"
                       "承認後に再度ログインしてください。") from exc


def _model_for(variant: str, token: str) -> tuple:
    cached = MODELS.get(variant)
    if cached is not None:
        return cached
    filename = f"{variant}/best_style_encoder.safetensors"
    path = hf_hub_download(MODEL_REPO, filename, revision=MODEL_REVISION, token=token)
    hf_hub_download(MODEL_REPO, f"{variant}/best_metadata.json",
                    revision=MODEL_REVISION, token=token)
    model, processor, max_patches = load_encoder(path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
    model = model.to(device, dtype=dtype).eval()
    cached = (model, processor, max_patches, device, dtype)
    MODELS[variant] = cached
    return cached


def _embedding(image: Image.Image, model, processor, max_patches, device, dtype):
    inputs = processor(images=[image], max_num_patches=max_patches, return_tensors="pt")
    with torch.inference_mode():
        return model(inputs["pixel_values"].to(device, dtype=dtype),
                     inputs["pixel_attention_mask"].to(device),
                     inputs["spatial_shapes"].to(device))[0].float()


@spaces.GPU(duration=240)
def compare(image_a: Image.Image | None, image_b: Image.Image | None,
            variant: str, target_recall: str,
            oauth_token: gr.OAuthToken | None = None) -> tuple[float, str]:
    if variant not in CALIBRATION or target_recall not in CALIBRATION[variant]["targets"]:
        raise gr.Error("無効なモデルまたは判定設定です。")
    token = oauth_token.token if oauth_token is not None else None
    _require_model_access(variant, token)
    first = _check_image(image_a, "A")
    second = _check_image(image_b, "B")
    with MODEL_LOCK:
        model, processor, max_patches, device, dtype = _model_for(variant, token)
        a = _embedding(first, model, processor, max_patches, device, dtype)
        b = _embedding(second, model, processor, max_patches, device, dtype)
        cosine = float((a * b).sum().clamp(-1, 1))
    operating_point = CALIBRATION[variant]["targets"][target_recall]
    threshold = operating_point["threshold"]
    verdict = "同一 style とみなす" if cosine >= threshold else "同一 style と判定しない"
    result = (f"### {verdict}\n\n"
              f"**Cosine 類似度:** `{cosine:.4f}`　|　**判定閾値:** `{threshold:.4f}`\n\n"
              f"この閾値は未学習作者の同作者画像ペア **{operating_point['same_recall']:.1%}** を拾い、"
              f"異作者ペアも **{operating_point['different_acceptance']:.1%}** 受け入れた設定です。"
              "作者ラベルで校正した参考判定であり、画風の同一性を証明するものではありません。")
    return round(cosine, 4), result


CSS = """
.gradio-container {max-width: 1060px !important; margin: auto !important;}
.intro {line-height: 1.55;}
.result-box {min-height: 170px;}
"""


def create_ui() -> gr.Blocks:
    with gr.Blocks(title="Style Embedder · A/B 比較") as demo:
        gr.Markdown("# Style Embedder · A/B 比較")
        gr.Markdown(
            "画像 A と B の画風ベクトルの cosine 類似度を測ります。まず "
            "[モデルのアクセス申請](https://huggingface.co/celstk/style_embedder)を行い、承認後にログインしてください。"
            "入力画像はこの Space の推論に使い、モデルリポジトリには送信しません。",
            elem_classes="intro",
        )
        gr.LoginButton()
        with gr.Row():
            image_a = gr.Image(type="pil", sources=["upload"], label="画像 A", height=320)
            image_b = gr.Image(type="pil", sources=["upload"], label="画像 B", height=320)
        with gr.Row():
            variant = gr.Dropdown(choices=[("Baseline（推奨）", "baseline"), ("HSIC", "hsic")],
                                  value="baseline", label="モデル")
            target_recall = gr.Dropdown(
                choices=[("再現率 90%", "0.9"), ("再現率 95%", "0.95"),
                         ("再現率 99%・見逃し重視", "0.99"), ("検証ペアの見逃し 0", "1.0")],
                value="0.99", label="判定設定",
                info="同作者ペアの見逃しを減らすほど、異作者も同一と判定しやすくなります。",
            )
        run = gr.Button("比較する", variant="primary")
        cosine = gr.Number(label="Style cosine 類似度", precision=4)
        result = gr.Markdown(elem_classes="result-box")
        run.click(compare, inputs=[image_a, image_b, variant, target_recall],
                  outputs=[cosine, result], concurrency_limit=1)
        gr.Markdown(
            "判定は dataset 39 の未学習作者376人・490枚（同作者ペア298組）で校正しました。"
            "best checkpoint の選択にもこの検証集合を使用したため、独立した外部テスト結果ではありません。"
            "作者名や画像は Space のコード・モデルには含まれません。"
        )
    demo.queue()
    return demo


if __name__ == "__main__":
    create_ui().launch(css=CSS)