from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any

import torch
from huggingface_hub import snapshot_download
from peft import PeftModel
from PIL import Image
from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    BitsAndBytesConfig,
)

BASE_MODEL_ID = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
MODEL_REPO_ID = "TuWaveGod/Puker_Judge"
MAX_LENGTH = 2048
MAX_IMAGE_LONGEST_EDGE = 1280

BINARY_PROMPT = (
    "Judge whether this geometrically assembled playing card has coherent rank, "
    "suit, border, portrait, symbols, and continuous artwork. A whole-card "
    "180-degree rotation is valid. Answer VALID or INVALID only."
)


def rank_prompt(candidate_count: int) -> str:
    if not 2 <= candidate_count <= 4:
        raise ValueError("Rank inference requires 2 to 4 candidates.")
    labels = ", ".join(str(index) for index in range(1, candidate_count + 1))
    return (
        "All displayed candidates are geometrically valid reconstructions made "
        "from the same playing-card pieces. Select the candidate whose rank, suit, "
        "outer border, portrait, symbols, and line artwork form one coherent "
        "original playing card. A whole-card 180-degree rotation is equivalent. "
        f"The available labels are {labels}. Answer with one label only."
    )


def resize_for_model(image: Image.Image) -> Image.Image:
    image = image.convert("RGB")
    longest = max(image.size)
    if longest <= MAX_IMAGE_LONGEST_EDGE:
        return image
    scale = MAX_IMAGE_LONGEST_EDGE / longest
    return image.resize(
        (
            max(1, int(round(image.width * scale))),
            max(1, int(round(image.height * scale))),
        ),
        Image.Resampling.LANCZOS,
    )


def load_adapter(
    adapter_name: str,
    *,
    repo_id: str = MODEL_REPO_ID,
    base_model_id: str = BASE_MODEL_ID,
    int4: bool = False,
) -> tuple[Any, Any, torch.device, dict[str, float]]:
    if adapter_name not in {"binary_adapter", "rank_adapter"}:
        raise ValueError(f"Unknown adapter: {adapter_name}")
    if not torch.cuda.is_available():
        raise RuntimeError("A CUDA GPU is required by these example scripts.")

    local_repo = Path(repo_id).expanduser()
    if local_repo.is_dir():
        snapshot_path = local_repo.resolve()
        download_seconds = 0.0
    else:
        download_started = time.perf_counter()
        snapshot_path = Path(
            snapshot_download(
                repo_id=repo_id,
                allow_patterns=[
                    f"{adapter_name}/*",
                    "processor/*",
                ],
            )
        )
        download_seconds = time.perf_counter() - download_started

    processor = AutoProcessor.from_pretrained(snapshot_path / "processor")
    load_kwargs: dict[str, Any] = {
        "torch_dtype": torch.bfloat16,
        "attn_implementation": "sdpa",
    }
    if int4:
        load_kwargs.update(
            {
                "quantization_config": BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_quant_type="nf4",
                    bnb_4bit_compute_dtype=torch.bfloat16,
                    bnb_4bit_use_double_quant=True,
                ),
                "device_map": {"": 0},
            }
        )

    load_started = time.perf_counter()
    base_model = AutoModelForImageTextToText.from_pretrained(
        base_model_id,
        **load_kwargs,
    )
    if not int4:
        base_model = base_model.to("cuda:0")
    model = PeftModel.from_pretrained(
        base_model,
        snapshot_path / adapter_name,
    ).eval()
    torch.cuda.synchronize()
    load_seconds = time.perf_counter() - load_started
    return (
        model,
        processor,
        torch.device("cuda:0"),
        {
            "snapshot_download_seconds": download_seconds,
            "model_load_seconds": load_seconds,
        },
    )


def encode_image_prompt(
    processor: Any,
    image: Image.Image,
    prompt: str,
    device: torch.device,
) -> dict[str, Any]:
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": prompt},
            ],
        }
    ]
    text = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=False,
    )
    inputs = processor(
        text=text,
        images=resize_for_model(image),
        return_tensors="pt",
        truncation=True,
        max_length=MAX_LENGTH,
    )
    moved: dict[str, Any] = {}
    for key, value in inputs.items():
        if not isinstance(value, torch.Tensor):
            moved[key] = value
        elif key == "pixel_values":
            moved[key] = value.to(device=device, dtype=torch.bfloat16)
        else:
            moved[key] = value.to(device=device)
    return moved


def generate_answer(
    model: Any,
    processor: Any,
    inputs: dict[str, Any],
    *,
    max_new_tokens: int = 4,
) -> tuple[str, float]:
    input_length = int(inputs["input_ids"].shape[1])
    torch.cuda.synchronize()
    started = time.perf_counter()
    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            do_sample=False,
            max_new_tokens=max_new_tokens,
        )
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - started
    answer = processor.decode(
        output_ids[0, input_length:],
        skip_special_tokens=True,
    ).strip()
    return answer, elapsed


def print_json(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, ensure_ascii=False, indent=2))