import spaces  # noqa: F401 - must be imported before torch / CUDA-touching imports

import time

import gradio as gr
import torch

# The InternVL3.5 remote code predates transformers 5.x and never calls `post_init()`,
# so `all_tied_weights_keys` is missing when `from_pretrained` inspects tied weights.
from transformers.modeling_utils import PreTrainedModel

if not hasattr(PreTrainedModel, "all_tied_weights_keys"):
    PreTrainedModel.all_tied_weights_keys = {}

from PIL import Image
from torchvision import transforms as T
from torchvision.transforms.functional import InterpolationMode
from transformers import AutoModel, AutoTokenizer

MODEL_ID = "Trendyol/Trendyol-Vision-Flash"
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
INPUT_SIZE = 448

model = (
    AutoModel.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
        use_flash_attn=False,
    )
    .eval()
    .to("cuda")
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=False)

# Exactly the preprocessing from the official model card: a single 448x448 tile per image.
_transform = T.Compose(
    [
        T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
        T.Resize((INPUT_SIZE, INPUT_SIZE), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    ]
)


def _to_pixel_values(images):
    tensors = [_transform(img).unsqueeze(0) for img in images]
    return torch.cat(tensors, dim=0).to(dtype=torch.bfloat16, device="cuda")


def _generate(images, prompt, max_new_tokens, num_patches_list=None):
    pixel_values = _to_pixel_values(images)
    generation_config = {"max_new_tokens": int(max_new_tokens), "do_sample": False}
    started = time.perf_counter()
    response = model.chat(
        tokenizer,
        pixel_values,
        prompt,
        generation_config,
        num_patches_list=num_patches_list,
    )
    elapsed = time.perf_counter() - started
    print(f"[trendyol-vision-flash] {len(images)} image(s), {elapsed:.2f}s", flush=True)
    return response.strip(), f"⚡ generated in **{elapsed:.2f}s** on ZeroGPU"


# ----------------------------------------------------------------------------- inference


@spaces.GPU(duration=10)
def analyze_product(
    image: Image.Image,
    prompt: str,
    max_new_tokens: int = 512,
) -> tuple:
    """Run a catalog task on a single product image with Trendyol-Vision-Flash.

    Args:
        image: the product photo to analyse.
        prompt: the task instruction (Turkish or English). A `<image>` placeholder is
            prepended automatically when the prompt does not contain one.
        max_new_tokens: maximum number of tokens to generate.

    Returns:
        The model's answer and a short latency note.
    """
    if image is None:
        raise gr.Error("Please upload a product image first.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt (or pick a task template above).")

    prompt = prompt.strip()
    n_slots = prompt.count("<image>")
    if n_slots == 0:
        prompt = "<image>\n" + prompt
    elif n_slots > 1:
        raise gr.Error(
            f"This tab handles one image, but the prompt has {n_slots} `<image>` "
            "placeholders. Use the 'Compare two products' tab instead."
        )
    return _generate([image], prompt, max_new_tokens)


@spaces.GPU(duration=6)
def compare_products(
    image_before: Image.Image,
    image_after: Image.Image,
    prompt: str,
    max_new_tokens: int = 32,
) -> tuple:
    """Compare two product listings with Trendyol-Vision-Flash.

    Args:
        image_before: photo of the "before" listing.
        image_after: photo of the "after" listing.
        prompt: the task instruction, containing two `<image>` placeholders.
        max_new_tokens: maximum number of tokens to generate.

    Returns:
        The model's answer and a short latency note.
    """
    if image_before is None or image_after is None:
        raise gr.Error("Please upload both product images.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    prompt = prompt.strip()
    n_slots = prompt.count("<image>")
    if n_slots != 2:
        raise gr.Error(
            f"The prompt needs exactly two `<image>` placeholders (found {n_slots}) — "
            "one for the 'before' listing and one for the 'after' listing."
        )
    return _generate([image_before, image_after], prompt, max_new_tokens, num_patches_list=[1, 1])


# ----------------------------------------------------------------------------- prompts

SINGLE_TASKS = {
    "Brand detection": (
        "<image>\n"
        "Görsellerden, {category} kategorisinde yer alan ürünün markasını çıkar.\n"
        "Sadece verilen görsellerde doğrulanabilen bilgilere dayan.\n"
        'Emin olmadığında "Unknown" şeklinde cevap ver.\n'
        "Sadece marka adını döndür."
    ),
    "Attribute extraction": (
        "<image>\n"
        "Bu görseldeki, başlığı '{title}' ve açıklaması '{description}' olan ürünün "
        "{attribute_list} bilgilerini json formatında çıkarır mısın?"
    ),
    "Title generation": (
        "<image>\n"
        "Ürün fotoğrafı ile '{reference_title}' bilgisini karşılaştırıp, Trendyol katalog "
        "moderasyon kurallarına göre yanıltıcı ifadelerden kaçınarak net bir başlık üret; "
        "çıktıyı düz metin ver."
    ),
    "Product caption (EN)": (
        "<image>\n"
        "Without speculating about details you cannot see, describe this product based on "
        "the image and the information provided.\n"
        "Product title: {title}\n"
        "Brand: {brand}\n"
        "First decide which object is the product, review OCR for brand/model/title clues, "
        "then analyze colors, shape, material, pattern, and other grounded details."
    ),
    "Content safety classification": (
        "<image>\n"
        "Ürün Başlığı: {title}\n\n"
        "Bu ürün görselini ve başlığını inceleyerek moderasyon sınıflandırması yap. "
        "Sonucu 0 (Forbidden), 1 (Fantasy) veya 2 (Safe) olarak ver."
    ),
    "Per-pack quantity extraction": (
        "<image>\n"
        "Extract total product quantity from the title and image. Title is the primary "
        "source; use the image only when the title is missing or ambiguous.\n"
        'Return only JSON: {"amount": <number|null>, "unit": "KG"|"L"|"PIECE"|null}\n'
        "Convert g→KG, ml/cc→L, counts→PIECE. Ignore model numbers, storage, wattage, and "
        'dimensions. If unclear, return {"amount": null, "unit": null}.\n\n'
        "Product Title: {title}"
    ),
}

SIMILARITY_TEMPLATE = (
    "E-ticaret kataloğundaki ürün benzerliği konusunda uzmansınız.\n"
    "İki ürünün önceki ve sonraki başlık/görsellerini karşılaştırarak aynı ürün olup "
    "olmadığını belirle.\n"
    "Ambalaj, arka plan veya model farkları tek başına fark sayılmaz; renk, boyut, miktar "
    "veya varyant farkları fark sayılır.\n"
    'Sadece "1" veya "0" döndür.\n\n'
    "Önceki Ürün Başlığı:\n{before_title}\n"
    "Önceki Ürün Resimleri:\n<image>\n"
    "Sonraki Ürün Başlığı:\n{after_title}\n"
    "Sonraki Ürün Resimleri:\n<image>"
)

BRAND_PROMPT = SINGLE_TASKS["Brand detection"].replace("{category}", "Eldiven")

CHAIR_PROMPT = (
    "<image>\n"
    "Bu görseldeki, başlığı 'Helen Bar Sandalyesi-mavi-9519q0119' ve açıklaması "
    "'Maksimum kargolanma süresi: Sipariş tarihinden sonraki 7. gün Ortalama montaj süresi: "
    "1 dakika Genişlik: 62cm Derinlik: 52cm Yükseklik: 96cm Oturak Yüksekliği: 44cm Ürün "
    "Ağırlığı: 10kg Kullanılan malzeme: Metal Kullanılan sünger: Yüksek yoğunluklu gri "
    "sünger Kullanılan kol: Metal Kollu Kullanılan ayak: Metal Ayaklı + Kromajlı Garanti "
    "süresi - Menşei: 24 Ay - Yerli' olan ürünün Garanti Süresi, Materyal, Model, Sandalye "
    "Kumaşı, Sandalye Sayısı, Tema / Stil bilgilerini json formatında çıkarır mısın?"
)

SUNGLASSES_PROMPT = SINGLE_TASKS["Title generation"].replace(
    "{reference_title}",
    "%100 Polarize UV400 Korumalı Füme Renk Orijinal Oval Erkek Güneş Gözlüğü DKG6273C3",
)

DRESS_PROMPT = (
    SINGLE_TASKS["Product caption (EN)"]
    .replace("{title}", "çok renkli çiçek desenli anne elbisesi")
    .replace("{brand}", "mihvera")
)

SAFETY_PROMPT = SINGLE_TASKS["Content safety classification"].replace(
    "{title}",
    "8. SINIF LGS - MEB BÖYLE SORAR SARMAL BRANŞ DENEME SETİ İKİNCİ DOZ (2025-LGS)",
)

QUANTITY_PROMPT = SINGLE_TASKS["Per-pack quantity extraction"].replace(
    "{title}", "Seda 6'lı Çay Bardağı"
)

OIL_PROMPT = SIMILARITY_TEMPLATE.replace(
    "{before_title}", "Bambu Kapaklı Vakumlu Borosilikat Yağdanlık | 3 Parça - 2lt"
).replace("{after_title}", "Bambu Kapaklı Vakumlu Borosilikat Yağdanlık Seti")

GLOVE_PROMPT = SIMILARITY_TEMPLATE.replace(
    "{before_title}", "Adidas Erkek Kışlık Eldiven"
).replace("{after_title}", "Adidas Unisex Kışlık Eldiven Siyah")

# ----------------------------------------------------------------------------- UI

CSS = """
#col-container { max-width: 1180px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

HEAD = """
# ⚡ Trendyol-Vision-Flash

The lightweight (1.06B) **e-commerce catalog** vision-language model from Trendyol — brand
detection, attribute extraction, title generation, product captioning, content-safety
moderation, per-pack quantity and product similarity, in Turkish and English.

[model](https://huggingface.co/Trendyol/Trendyol-Vision-Flash) ·
[base: InternVL3.5-1B](https://huggingface.co/OpenGVLab/InternVL3_5-1B-Instruct) ·
[larger sibling: Trendyol-Vision-Master](https://huggingface.co/Trendyol/Trendyol-Vision-Master)
"""

with gr.Blocks(title="Trendyol-Vision-Flash") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(HEAD)

        with gr.Tabs():
            # ------------------------------------------------ single-image catalog tasks
            with gr.Tab("Catalog task"):
                with gr.Row():
                    with gr.Column(scale=1):
                        image = gr.Image(label="Product image", type="pil", height=340)
                    with gr.Column(scale=1):
                        task = gr.Dropdown(
                            label="Task template",
                            choices=list(SINGLE_TASKS),
                            value="Brand detection",
                            info="Loads the official prompt template — replace the {placeholders}.",
                        )
                        prompt = gr.Textbox(
                            label="Prompt",
                            lines=8,
                            placeholder="Ask something about the product image…",
                            value=BRAND_PROMPT,
                        )
                        run = gr.Button("Run", variant="primary")

                answer = gr.Textbox(label="Model output", lines=6)
                latency = gr.Markdown()

                with gr.Accordion("Advanced", open=False):
                    max_tokens = gr.Slider(
                        label="Max new tokens",
                        minimum=8,
                        maximum=1024,
                        step=8,
                        value=512,
                    )

                gr.Examples(
                    examples=[
                        ["examples/glove_1.jpg", BRAND_PROMPT],
                        ["examples/bar_chair.jpg", CHAIR_PROMPT],
                        ["examples/sunglasses.jpg", SUNGLASSES_PROMPT],
                        ["examples/dress.jpg", DRESS_PROMPT],
                        ["examples/test_books.jpg", SAFETY_PROMPT],
                        ["examples/tea_glasses.jpg", QUANTITY_PROMPT],
                    ],
                    inputs=[image, prompt],
                    outputs=[answer, latency],
                    fn=analyze_product,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Examples from the official model card",
                )

                task.change(
                    lambda name: SINGLE_TASKS.get(name, ""),
                    inputs=task,
                    outputs=prompt,
                )
                run.click(
                    analyze_product,
                    inputs=[image, prompt, max_tokens],
                    outputs=[answer, latency],
                    api_name="analyze_product",
                )

            # ------------------------------------------------------ product similarity
            with gr.Tab("Compare two products"):
                gr.Markdown(
                    "Decide whether two catalog listings are the **same SKU** — "
                    "`1` = same product, `0` = different."
                )
                with gr.Row():
                    image_before = gr.Image(label="Before listing", type="pil", height=300)
                    image_after = gr.Image(label="After listing", type="pil", height=300)
                sim_prompt = gr.Textbox(
                    label="Prompt (needs two <image> placeholders)",
                    lines=10,
                    value=SIMILARITY_TEMPLATE,
                )
                sim_run = gr.Button("Compare", variant="primary")
                sim_answer = gr.Textbox(label="Model output", lines=2)
                sim_latency = gr.Markdown()

                with gr.Accordion("Advanced", open=False):
                    sim_max_tokens = gr.Slider(
                        label="Max new tokens", minimum=4, maximum=256, step=4, value=32
                    )

                gr.Examples(
                    examples=[
                        [
                            "examples/oil_bottle_before.jpg",
                            "examples/oil_bottle_after.jpg",
                            OIL_PROMPT,
                        ],
                        ["examples/glove_1.jpg", "examples/glove_2.jpg", GLOVE_PROMPT],
                    ],
                    inputs=[image_before, image_after, sim_prompt],
                    outputs=[sim_answer, sim_latency],
                    fn=compare_products,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Examples from the official model card",
                )

                sim_run.click(
                    compare_products,
                    inputs=[image_before, image_after, sim_prompt, sim_max_tokens],
                    outputs=[sim_answer, sim_latency],
                    api_name="compare_products",
                )

        gr.Markdown(
            "Prompts and example product photos come from the official "
            "[Trendyol-Vision-Flash](https://huggingface.co/Trendyol/Trendyol-Vision-Flash) "
            "and [Trendyol-Vision-Master](https://huggingface.co/Trendyol/Trendyol-Vision-Master) "
            "model cards (CC BY 4.0, © Trendyol). Greedy decoding, single 448×448 tile per "
            "image — exactly the reference inference path from the model card."
        )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)