#!/usr/bin/env python3
from __future__ import annotations

import argparse
import re
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

from puker_judge_utils import (
    BASE_MODEL_ID,
    MODEL_REPO_ID,
    encode_image_prompt,
    generate_answer,
    load_adapter,
    print_json,
    rank_prompt,
)

BOARD_SIZE = (1280, 820)
CANDIDATE_SIZE = (600, 360)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Choose the best of two to four playing-card candidates."
    )
    parser.add_argument(
        "candidates",
        type=Path,
        nargs="*",
        help="Two to four rectified candidate images.",
    )
    parser.add_argument(
        "--board-image",
        type=Path,
        help="Use an already constructed board instead of candidate files.",
    )
    parser.add_argument(
        "--candidate-count",
        type=int,
        help="Required with --board-image; must be between 2 and 4.",
    )
    parser.add_argument(
        "--board-output",
        type=Path,
        default=Path("candidate_board.jpg"),
    )
    parser.add_argument("--repo-id", default=MODEL_REPO_ID)
    parser.add_argument("--base-model-id", default=BASE_MODEL_ID)
    parser.add_argument(
        "--int4",
        action="store_true",
        help="Use bitsandbytes NF4 weights with BF16 compute.",
    )
    return parser.parse_args()


def load_font(size: int) -> ImageFont.ImageFont:
    for path in (
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "DejaVuSans-Bold.ttf",
    ):
        try:
            return ImageFont.truetype(path, size=size)
        except OSError:
            continue
    return ImageFont.load_default()


def build_board(candidate_paths: list[Path]) -> Image.Image:
    count = len(candidate_paths)
    if not 2 <= count <= 4:
        raise ValueError("Provide exactly 2, 3, or 4 candidate images.")

    board = Image.new("RGB", BOARD_SIZE, (226, 231, 237))
    draw = ImageDraw.Draw(board)
    font = load_font(34)
    outer_margin = 28
    column_gap = 24
    row_gap = 24
    cell_width = (BOARD_SIZE[0] - 2 * outer_margin - column_gap) // 2
    cell_height = (BOARD_SIZE[1] - 2 * outer_margin - row_gap) // 2
    label_height = 46

    for index, path in enumerate(candidate_paths):
        column = index % 2
        row = index // 2
        x0 = outer_margin + column * (cell_width + column_gap)
        y0 = outer_margin + row * (cell_height + row_gap)
        x1 = x0 + cell_width
        y1 = y0 + cell_height
        draw.rounded_rectangle(
            (x0, y0, x1, y1),
            radius=14,
            fill=(244, 246, 248),
            outline=(178, 184, 192),
            width=2,
        )
        draw.rounded_rectangle(
            (x0 + 12, y0 + 8, x0 + 76, y0 + label_height),
            radius=10,
            fill=(255, 218, 72),
            outline=(45, 48, 52),
            width=2,
        )
        draw.text(
            (x0 + 44, y0 + 8 + label_height // 2),
            str(index + 1),
            font=font,
            fill=(25, 31, 42),
            anchor="mm",
        )
        with Image.open(path) as source:
            candidate = source.convert("RGB").resize(
                CANDIDATE_SIZE,
                Image.Resampling.LANCZOS,
            )
        available_width = cell_width - 36
        available_height = cell_height - label_height - 28
        candidate.thumbnail(
            (available_width, available_height),
            Image.Resampling.LANCZOS,
        )
        paste_x = x0 + (cell_width - candidate.width) // 2
        paste_y = y0 + label_height + (
            cell_height - label_height - candidate.height
        ) // 2
        draw.rectangle(
            (
                paste_x + 4,
                paste_y + 5,
                paste_x + 4 + candidate.width,
                paste_y + 5 + candidate.height,
            ),
            fill=(185, 190, 196),
        )
        board.paste(candidate, (paste_x, paste_y))
    return board


def main() -> None:
    args = parse_args()
    if args.board_image is not None:
        if args.candidates:
            raise SystemExit(
                "Use either candidate files or --board-image, not both."
            )
        if args.candidate_count is None or not 2 <= args.candidate_count <= 4:
            raise SystemExit(
                "--candidate-count 2..4 is required with --board-image."
            )
        with Image.open(args.board_image) as source:
            board = source.convert("RGB").copy()
        candidate_count = args.candidate_count
        candidate_paths: list[Path] = []
    else:
        if args.candidate_count is not None:
            raise SystemExit(
                "--candidate-count is inferred when candidate files are used."
            )
        candidate_paths = args.candidates
        candidate_count = len(candidate_paths)
        board = build_board(candidate_paths)
        args.board_output.parent.mkdir(parents=True, exist_ok=True)
        board.save(args.board_output, "JPEG", quality=96, subsampling=0)

    model, processor, device, timings = load_adapter(
        "rank_adapter",
        repo_id=args.repo_id,
        base_model_id=args.base_model_id,
        int4=args.int4,
    )
    inputs = encode_image_prompt(
        processor,
        board,
        rank_prompt(candidate_count),
        device,
    )
    raw_output, generation_seconds = generate_answer(
        model,
        processor,
        inputs,
    )
    match = re.search(r"[1-4]", raw_output)
    if match is None or int(match.group(0)) > candidate_count:
        raise SystemExit(f"Model returned an invalid answer: {raw_output!r}")
    label = int(match.group(0))
    payload = {
        "selected_label": label,
        "raw_output": raw_output,
        "candidate_count": candidate_count,
        "quantization": "int4-nf4" if args.int4 else "bf16",
        "generation_seconds": round(generation_seconds, 4),
        **{key: round(value, 4) for key, value in timings.items()},
    }
    if candidate_paths:
        payload["selected_file"] = str(candidate_paths[label - 1].resolve())
        payload["board_image"] = str(args.board_output.resolve())
    else:
        payload["board_image"] = str(args.board_image.resolve())
    print_json(payload)


if __name__ == "__main__":
    main()