"""画像の読み込みとレターボックス化（SPEC_02「前処理」）.

アスペクト比を保って長辺を `image_size` に合わせ，余白を黒 (0, 0, 0) で埋める．マスク外の
領域も黒であるため，パディングも黒で統一する．学習を速くするため，全画像をレターボックス化した
uint8 配列としてメモリに載せる（1085 枚 × 224² × 3 ≒ 160 MB）．
"""

from __future__ import annotations

import logging
from collections.abc import Sequence
from pathlib import Path

import numpy as np
from PIL import Image

from tias_judge_model.data.manifest import ManifestRecord, resolve_image_path

logger = logging.getLogger(__name__)


def letterbox(image: Image.Image, size: int) -> np.ndarray:
    """画像をアスペクト比を保って `size × size` に収め，余白を黒で埋める.

    長辺が `size` になるよう縮小（または拡大）し，中央に配置する．

    Args:
        image: RGB の PIL 画像．
        size: 出力の一辺（ピクセル）．

    Returns:
        `(size, size, 3)` の uint8 配列（RGB）．
    """
    width, height = image.size
    scale = size / max(width, height)
    new_w = max(1, round(width * scale))
    new_h = max(1, round(height * scale))
    resized = image.resize((new_w, new_h), Image.Resampling.BILINEAR)
    canvas = Image.new("RGB", (size, size), (0, 0, 0))
    canvas.paste(resized, ((size - new_w) // 2, (size - new_h) // 2))
    return np.asarray(canvas, dtype=np.uint8)


def load_letterboxed_images(
    records: Sequence[ManifestRecord], images_dir: Path, size: int
) -> np.ndarray:
    """manifest の全画像を読み，レターボックス化して 1 つの配列にまとめる.

    読めない画像は学習を歪めるため止める（fail fast）．例外メッセージには画像のパスを出さず，
    manifest のデータ行の番号（1 始まり）で示す．

    Args:
        records: manifest レコード．
        images_dir: 画像ディレクトリ（`<data_root>/images`）．
        size: 出力の一辺（ピクセル）．

    Returns:
        `(N, size, size, 3)` の uint8 配列（RGB）．並びは `records` と同じ．

    Raises:
        ValueError: 画像が無い，または読めない場合．
    """
    out = np.zeros((len(records), size, size, 3), dtype=np.uint8)
    for i, record in enumerate(records):
        path = resolve_image_path(images_dir, record)
        if not path.is_file():
            raise ValueError(f"manifest のデータ {i + 1} 件目: 画像ファイルが無い")
        try:
            with Image.open(path) as img:
                img.load()
                rgb = img.convert("RGB")
        except (OSError, ValueError, Image.DecompressionBombError) as exc:
            raise ValueError(
                f"manifest のデータ {i + 1} 件目: 画像を読めない（{type(exc).__name__}）"
            ) from None
        out[i] = letterbox(rgb, size)
    logger.info("画像 %d 件をレターボックス化した（%d×%d）", len(records), size, size)
    return out