from __future__ import annotations

from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw


BACKGROUND_COLORS: dict[str, tuple[int, int, int]] = {
    "black": (0, 0, 0),
    "white": (255, 255, 255),
    "gray": (128, 128, 128),
    "red": (220, 32, 32),
    "blue": (32, 96, 220),
}


def composite(rgb: np.ndarray, alpha: np.ndarray, color: tuple[int, int, int]) -> np.ndarray:
    fg = rgb.astype(np.float32)
    bg = np.zeros_like(fg)
    bg[..., 0] = color[0]
    bg[..., 1] = color[1]
    bg[..., 2] = color[2]
    a = alpha[..., None].astype(np.float32)
    return np.clip(fg * a + bg * (1.0 - a), 0, 255).astype(np.uint8)


def save_previews(debug_dir: str | Path, rgb: np.ndarray, alpha: np.ndarray) -> dict[str, Path]:
    debug = Path(debug_dir)
    debug.mkdir(parents=True, exist_ok=True)
    paths: dict[str, Path] = {}
    for name, color in BACKGROUND_COLORS.items():
        path = debug / f"preview_{name}.png"
        Image.fromarray(composite(rgb, alpha, color), mode="RGB").save(path)
        paths[name] = path
    return paths


def make_qa_grid(debug_dir: str | Path, tiles: dict[str, np.ndarray]) -> Path:
    debug = Path(debug_dir)
    debug.mkdir(parents=True, exist_ok=True)
    tile_items = list(tiles.items())
    if not tile_items:
        raise ValueError("No QA tiles provided")

    thumb_w = 320
    label_h = 28
    padding = 10
    cols = min(3, len(tile_items))
    rows = int(np.ceil(len(tile_items) / cols))
    first_h, first_w = tile_items[0][1].shape[:2]
    thumb_h = max(180, int(round(thumb_w * first_h / max(first_w, 1))))

    canvas = Image.new(
        "RGB",
        (
            cols * thumb_w + (cols + 1) * padding,
            rows * (thumb_h + label_h) + (rows + 1) * padding,
        ),
        (32, 32, 32),
    )
    draw = ImageDraw.Draw(canvas)

    for idx, (name, arr) in enumerate(tile_items):
        row = idx // cols
        col = idx % cols
        x = padding + col * (thumb_w + padding)
        y = padding + row * (thumb_h + label_h + padding)
        img = Image.fromarray(arr.astype(np.uint8), mode="RGB")
        img.thumbnail((thumb_w, thumb_h), Image.Resampling.LANCZOS)
        frame = Image.new("RGB", (thumb_w, thumb_h), (240, 240, 240))
        frame.paste(img, ((thumb_w - img.width) // 2, (thumb_h - img.height) // 2))
        canvas.paste(frame, (x, y + label_h))
        draw.text((x, y), name, fill=(245, 245, 245))

    path = debug / "qa_grid.png"
    canvas.save(path)
    return path