#!/usr/bin/env python
"""把 labelme / VOC 标注批量转换成 YOLO 风格的 .txt。

    python scripts/convert_annotations.py --format labelme \
        --input data/raw/film/labelme --images data/raw/film/images \
        --output data/raw/film/labels --data-yaml data/raw/film/data.yaml

每张图输出一个同名 .txt，每行 `class_id cx cy w h`（相对宽高归一化）；
类别顺序默认按扫描到的名字排序，也可以用 `--classes scratch bubble ...` 固定。
"""

from __future__ import annotations

import argparse
import json
import sys
import xml.etree.ElementTree as ET
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from datasets.yolo import IMAGE_SUFFIXES, format_label_lines, write_data_yaml  # noqa: E402


def from_labelme(path: Path) -> dict:
    data = json.loads(path.read_text(encoding="utf-8"))
    boxes, labels = [], []
    for shape in data.get("shapes", []):
        points = shape["points"]
        xs = [p[0] for p in points]
        ys = [p[1] for p in points]
        boxes.append([min(xs), min(ys), max(xs), max(ys)])
        labels.append(shape["label"])
    return {
        "image": Path(data.get("imagePath", path.stem + ".png")).name,
        "width": data.get("imageWidth"),
        "height": data.get("imageHeight"),
        "boxes": boxes,
        "labels": labels,
    }


def from_voc(path: Path) -> dict:
    root = ET.parse(path).getroot()
    size = root.find("size")
    boxes, labels = [], []
    for obj in root.findall("object"):
        bnd = obj.find("bndbox")
        if bnd is None:
            continue
        boxes.append([float(bnd.findtext(key, "0")) for key in ("xmin", "ymin", "xmax", "ymax")])
        labels.append(obj.findtext("name", "unknown"))
    return {
        "image": Path(root.findtext("filename") or (path.stem + ".jpg")).name,
        "width": int(size.findtext("width", "0")) if size is not None else None,
        "height": int(size.findtext("height", "0")) if size is not None else None,
        "boxes": boxes,
        "labels": labels,
    }


PARSERS = {"labelme": (from_labelme, "*.json"), "voc": (from_voc, "*.xml")}


def image_size(record: dict, images_dir: Path | None) -> tuple[int, int]:
    """标注里没写尺寸就去读图片本身——归一化坐标必须要有宽高。"""
    if record.get("width") and record.get("height"):
        return int(record["width"]), int(record["height"])
    if images_dir is None:
        raise ValueError(f"{record['image']} 的标注里没有图像尺寸，请用 --images 指定图片目录")

    from PIL import Image

    stem = Path(record["image"]).stem
    for suffix in IMAGE_SUFFIXES:
        candidate = images_dir / f"{stem}{suffix}"
        if candidate.exists():
            with Image.open(candidate) as image:
                return image.size
    raise FileNotFoundError(f"在 {images_dir} 下找不到 {record['image']}")


def main() -> int:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--format", choices=sorted(PARSERS), required=True)
    parser.add_argument("--input", required=True, help="标注文件所在目录")
    parser.add_argument("--output", required=True, help="输出的 labels 目录")
    parser.add_argument("--images", default=None, help="图片目录；标注里没有尺寸时用来读宽高")
    parser.add_argument("--classes", nargs="*", default=None, help="固定类别顺序（缺省按名字排序）")
    parser.add_argument("--data-yaml", default=None, help="顺带写一份 YOLO data.yaml")
    args = parser.parse_args()

    convert, pattern = PARSERS[args.format]
    files = sorted(Path(args.input).rglob(pattern))
    if not files:
        print(f"在 {args.input} 下没有找到 {pattern}")
        return 1

    records = [convert(file) for file in files]
    classes = args.classes or sorted({label for r in records for label in r["labels"]})
    class_to_idx = {name: i for i, name in enumerate(classes)}

    unknown = {label for r in records for label in r["labels"] if label not in class_to_idx}
    if unknown:
        print(f"标注中出现了 --classes 未列出的类别: {sorted(unknown)}")
        return 1

    images_dir = Path(args.images) if args.images else None
    out_dir = Path(args.output)
    out_dir.mkdir(parents=True, exist_ok=True)

    total_boxes = 0
    for record in records:
        width, height = image_size(record, images_dir)
        lines = format_label_lines(
            record["boxes"], [class_to_idx[label] for label in record["labels"]], width, height
        )
        total_boxes += len(lines)
        # 无缺陷图也写空文件，方便核对"标了但没有目标" vs "漏标"
        (out_dir / f"{Path(record['image']).stem}.txt").write_text(
            "\n".join(lines) + ("\n" if lines else ""), encoding="utf-8"
        )

    print(f"已转换 {len(records)} 个文件、{total_boxes} 个框 -> {out_dir}")
    print(f"类别顺序（class_id 即下标）: {classes}")
    if args.data_yaml:
        path = write_data_yaml(args.data_yaml, classes, root=out_dir.parent)
        print(f"data.yaml 已写入 {path}")
    return 0


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