"""Convert X-AnyLabeling rectangle JSON annotations to YOLO detection TXT."""

import argparse
import json
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(
        description="Convert X-AnyLabeling JSON files to YOLOv5 TXT labels."
    )
    parser.add_argument(
        "--input-dir",
        required=True,
        type=Path,
        help="Directory containing X-AnyLabeling JSON files.",
    )
    parser.add_argument(
        "--output-dir",
        required=True,
        type=Path,
        help="Directory in which YOLO TXT files will be written.",
    )
    parser.add_argument(
        "--classes",
        required=True,
        nargs="+",
        help="Class names in class-id order, for example: --classes gangqiu",
    )
    return parser.parse_args()


def convert_file(json_path, output_dir, class_to_id):
    with json_path.open("r", encoding="utf-8") as file:
        annotation = json.load(file)

    image_width = annotation.get("imageWidth")
    image_height = annotation.get("imageHeight")
    if not image_width or not image_height:
        raise ValueError("imageWidth or imageHeight is missing")

    lines = []
    for shape in annotation.get("shapes", []):
        if shape.get("shape_type") != "rectangle":
            raise ValueError(
                f"unsupported shape_type: {shape.get('shape_type')!r}; "
                "only rectangle annotations are allowed"
            )

        label = shape.get("label")
        if label not in class_to_id:
            raise ValueError(
                f"unknown class {label!r}; expected one of {list(class_to_id)}"
            )

        points = shape.get("points", [])
        if len(points) < 2:
            raise ValueError("a rectangle must contain at least two points")

        xs = [float(point[0]) for point in points]
        ys = [float(point[1]) for point in points]
        x_min = max(0.0, min(xs))
        y_min = max(0.0, min(ys))
        x_max = min(float(image_width), max(xs))
        y_max = min(float(image_height), max(ys))

        if x_max <= x_min or y_max <= y_min:
            raise ValueError("rectangle width and height must be positive")

        x_center = ((x_min + x_max) / 2.0) / image_width
        y_center = ((y_min + y_max) / 2.0) / image_height
        width = (x_max - x_min) / image_width
        height = (y_max - y_min) / image_height

        lines.append(
            f"{class_to_id[label]} "
            f"{x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
        )

    output_path = output_dir / f"{json_path.stem}.txt"
    output_path.write_text(
        "\n".join(lines) + ("\n" if lines else ""),
        encoding="utf-8",
    )


def main():
    args = parse_args()
    input_dir = args.input_dir.resolve()
    output_dir = args.output_dir.resolve()

    if not input_dir.is_dir():
        raise SystemExit(f"Input directory does not exist: {input_dir}")
    if len(set(args.classes)) != len(args.classes):
        raise SystemExit("Class names must not be duplicated.")

    output_dir.mkdir(parents=True, exist_ok=True)
    class_to_id = {name: index for index, name in enumerate(args.classes)}
    json_files = sorted(input_dir.glob("*.json"))
    if not json_files:
        raise SystemExit(f"No JSON files found in: {input_dir}")

    converted = 0
    for json_path in json_files:
        try:
            convert_file(json_path, output_dir, class_to_id)
            converted += 1
        except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError) as error:
            raise SystemExit(f"Failed to convert {json_path.name}: {error}") from error

    print(f"Converted {converted} JSON files.")
    print(f"YOLO labels: {output_dir}")
    print("Class mapping:")
    for name, class_id in class_to_id.items():
        print(f"  {class_id}: {name}")


if __name__ == "__main__":
    main()