#!/usr/bin/env python3
"""Build labeled contact sheets from serial_scenario webcam snapshots."""

from __future__ import annotations

import argparse
import math
from pathlib import Path

import cv2
import numpy as np


CROP_PRESETS = {
    "full": (0.00, 0.00, 1.00, 1.00),
    "panel": (0.03, 0.10, 0.98, 0.94),
    "right-stack": (0.70, 0.23, 0.92, 0.86),
    "lcd": (0.00, 0.22, 0.38, 0.78),
}


def main() -> int:
    args = build_arg_parser().parse_args()
    images = sorted(args.snapshot_dir.glob(args.glob))
    if args.only_candidates:
        images = [path for path in images if "candidate_" in path.name]
    if not images:
        raise SystemExit(f"no images found in {args.snapshot_dir}")

    output_dir = args.output_dir or args.snapshot_dir.with_name(args.snapshot_dir.name + "-sheets")
    output_dir.mkdir(parents=True, exist_ok=True)

    per_sheet = args.cols * args.rows
    total_sheets = math.ceil(len(images) / per_sheet)
    written: list[Path] = []
    for sheet_index in range(total_sheets):
        group = images[sheet_index * per_sheet : (sheet_index + 1) * per_sheet]
        sheet = build_sheet(group, args)
        out = output_dir / f"sheet-{sheet_index + 1:03d}.jpg"
        cv2.imwrite(str(out), sheet)
        written.append(out)

    print(f"images={len(images)} sheets={len(written)} output_dir={output_dir}")
    if written:
        print(f"first={written[0]}")
        print(f"last={written[-1]}")
    return 0


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Create labeled contact sheets from webcam snapshots.")
    parser.add_argument("snapshot_dir", type=Path, help="directory containing serial_scenario snapshots")
    parser.add_argument("--output-dir", type=Path, help="directory for generated contact sheets")
    parser.add_argument("--glob", default="*.jpg", help="image glob within snapshot_dir")
    parser.add_argument("--cols", type=int, default=4, help="thumbnail columns per sheet")
    parser.add_argument("--rows", type=int, default=5, help="thumbnail rows per sheet")
    parser.add_argument("--thumb-width", type=int, default=360, help="thumbnail image width")
    parser.add_argument("--label-height", type=int, default=48, help="label area height per thumbnail")
    parser.add_argument(
        "--crop",
        choices=sorted(CROP_PRESETS),
        default="panel",
        help="crop preset before thumbnailing",
    )
    parser.add_argument(
        "--only-candidates",
        action="store_true",
        help="include only filenames containing candidate_",
    )
    return parser


def build_sheet(paths: list[Path], args: argparse.Namespace) -> np.ndarray:
    cells = [build_cell(path, args) for path in paths]
    blank = np.full_like(cells[0], 245)
    while len(cells) < args.cols * args.rows:
        cells.append(blank.copy())

    rows = []
    for row_index in range(args.rows):
        row_cells = cells[row_index * args.cols : (row_index + 1) * args.cols]
        rows.append(np.hstack(row_cells))
    return np.vstack(rows)


def build_cell(path: Path, args: argparse.Namespace) -> np.ndarray:
    image = cv2.imread(str(path))
    if image is None:
        image = np.full((240, 320, 3), 220, dtype=np.uint8)
    image = crop_image(image, CROP_PRESETS[args.crop])
    thumb_height = max(1, int(image.shape[0] * args.thumb_width / max(1, image.shape[1])))
    thumb = cv2.resize(image, (args.thumb_width, thumb_height), interpolation=cv2.INTER_AREA)

    label = snapshot_label(path)
    cell = np.full((args.label_height + thumb_height, args.thumb_width, 3), 245, dtype=np.uint8)
    cell[args.label_height : args.label_height + thumb_height, 0 : args.thumb_width] = thumb
    draw_label(cell, label, args.label_height)
    return cell


def crop_image(image: np.ndarray, crop: tuple[float, float, float, float]) -> np.ndarray:
    h, w = image.shape[:2]
    x1, y1, x2, y2 = crop
    left = min(w - 1, max(0, int(w * x1)))
    top = min(h - 1, max(0, int(h * y1)))
    right = min(w, max(left + 1, int(w * x2)))
    bottom = min(h, max(top + 1, int(h * y2)))
    return image[top:bottom, left:right]


def draw_label(cell: np.ndarray, text: str, label_height: int) -> None:
    font = cv2.FONT_HERSHEY_SIMPLEX
    scale = 0.48
    thickness = 1
    max_chars = 45
    line1 = text[:max_chars]
    line2 = text[max_chars : max_chars * 2]
    cv2.putText(cell, line1, (6, 18), font, scale, (0, 0, 0), thickness, cv2.LINE_AA)
    if line2:
        cv2.putText(cell, line2, (6, min(label_height - 8, 38)), font, scale, (0, 0, 0), thickness, cv2.LINE_AA)


def snapshot_label(path: Path) -> str:
    name = path.name
    if "_tx_0500ms_" in name:
        label = name.split("_tx_0500ms_", 1)[1]
    elif "_tx_" in name:
        label = name.split("_tx_", 1)[1]
    else:
        label = name
    return label.rsplit("_", 1)[0]


if __name__ == "__main__":
    raise SystemExit(main())