#!/usr/bin/env python3
"""从文件名解析检测框并画到图片上。

文件名格式:
  有框: {id}_{label}_c{conf}_l{left}_t{top}_r{right}_b{bottom}.ext
  无框: {id}_no-result_c0_nobox.ext
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}

# {id}_{label}_c{conf}_l{L}_t{T}_r{R}_b{B}
BOX_RE = re.compile(
    r"^(?P<id>[^_]+)_(?P<label>.+?)_c(?P<conf>\d+)"
    r"_l(?P<left>\d+)_t(?P<top>\d+)_r(?P<right>\d+)_b(?P<bottom>\d+)$"
)
# {id}_no-result_c0_nobox 或 任意含 nobox 的名字
NOBOX_RE = re.compile(r"^(?P<id>[^_]+)_(?P<label>.+?)_c(?P<conf>\d+)_nobox$")


def parse_filename(stem: str) -> dict | None:
    """解析文件名 stem，返回框信息或 None（无法解析）。"""
    m = BOX_RE.match(stem)
    if m:
        d = m.groupdict()
        return {
            "id": d["id"],
            "label": d["label"],
            "conf": int(d["conf"]),
            "box": (
                int(d["left"]),
                int(d["top"]),
                int(d["right"]),
                int(d["bottom"]),
            ),
        }

    m = NOBOX_RE.match(stem)
    if m:
        d = m.groupdict()
        return {
            "id": d["id"],
            "label": d["label"],
            "conf": int(d["conf"]),
            "box": None,
        }

    return None


def draw_box(
    img: Image.Image,
    info: dict,
    color: tuple[int, int, int] = (0, 255, 0),
    width: int = 2,
) -> Image.Image:
    """在图片上画框和标签文字。"""
    out = img.convert("RGB")
    draw = ImageDraw.Draw(out)

    try:
        font = ImageFont.truetype(
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14
        )
    except OSError:
        font = ImageFont.load_default()

    label = f"{info['label']} c{info['conf']}"
    box = info["box"]

    if box is not None:
        l, t, r, b = box
        # 裁剪到图片范围，避免异常坐标
        w, h = out.size
        l, t = max(0, l), max(0, t)
        r, b = min(w - 1, r), min(h - 1, b)
        if r > l and b > t:
            draw.rectangle([l, t, r, b], outline=color, width=width)

            # 标签背景
            text_bbox = draw.textbbox((l, t), label, font=font)
            pad = 2
            bg = [
                text_bbox[0] - pad,
                max(0, text_bbox[1] - pad - 16),
                text_bbox[2] + pad,
                max(0, text_bbox[1] - pad - 16) + (text_bbox[3] - text_bbox[1]) + pad * 2,
            ]
            # 把文字放到框上方；如果贴顶则放到框内
            ty = t - (text_bbox[3] - text_bbox[1]) - 4
            if ty < 0:
                ty = t + 2
            bg = [
                l,
                ty - pad,
                l + (text_bbox[2] - text_bbox[0]) + pad * 2,
                ty + (text_bbox[3] - text_bbox[1]) + pad,
            ]
            draw.rectangle(bg, fill=color)
            draw.text((l + pad, ty), label, fill=(0, 0, 0), font=font)
    else:
        # 无框：左上角写标签
        draw.rectangle([0, 0, 160, 22], fill=(255, 0, 0))
        draw.text((4, 3), label, fill=(255, 255, 255), font=font)

    return out


def collect_images(src: Path) -> list[Path]:
    files = []
    for p in sorted(src.rglob("*")):
        if p.is_file() and p.suffix.lower() in IMAGE_EXTS:
            files.append(p)
    return files


def process_dir(src_dir: Path, out_dir: Path | None = None) -> None:
    src_dir = src_dir.resolve()
    if not src_dir.is_dir():
        print(f"错误: 目录不存在: {src_dir}", file=sys.stderr)
        sys.exit(1)

    if out_dir is None:
        out_dir = src_dir.parent / f"{src_dir.name}_boxed"
    out_dir = out_dir.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    images = collect_images(src_dir)
    if not images:
        print(f"未找到图片: {src_dir}")
        return

    ok = skip = fail = 0
    for img_path in images:
        rel = img_path.relative_to(src_dir)
        info = parse_filename(img_path.stem)
        if info is None:
            print(f"[跳过] 无法解析: {rel}")
            skip += 1
            continue

        try:
            with Image.open(img_path) as im:
                drawn = draw_box(im, info)
            save_path = out_dir / rel
            save_path.parent.mkdir(parents=True, exist_ok=True)
            drawn.save(save_path, quality=95)
            box_str = str(info["box"]) if info["box"] else "nobox"
            print(f"[OK] {rel}  ->  {box_str}")
            ok += 1
        except Exception as e:
            print(f"[失败] {rel}: {e}", file=sys.stderr)
            fail += 1

    print(f"\n完成: 成功 {ok}, 跳过 {skip}, 失败 {fail}")
    print(f"输出目录: {out_dir}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="根据文件名中的框信息，把检测框画到图片上"
    )
    parser.add_argument("dir", type=Path, help="输入图片目录（支持子目录）")
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=None,
        help="输出目录，默认 <输入目录名>_boxed",
    )
    parser.add_argument(
        "--color",
        default="0,255,0",
        help="框颜色 R,G,B，默认 0,255,0",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=2,
        help="框线宽，默认 2",
    )
    args = parser.parse_args()

    # 把颜色/线宽通过闭包注入（简单起见直接改默认参数调用）
    global draw_box  # noqa: PLW0603 — 仅 CLI 层覆盖

    color = tuple(int(x) for x in args.color.split(","))
    width = args.width
    _orig = draw_box

    def draw_box_wrapped(img, info, color=color, width=width):  # type: ignore[no-redef]
        return _orig(img, info, color=color, width=width)

    draw_box = draw_box_wrapped  # type: ignore[misc]

    process_dir(args.dir, args.output)


if __name__ == "__main__":
    main()