"""
Build an Ultralytics-style YOLO dataset YAML from a COCO annotations file.

Reads ``categories`` from the COCO JSON in list order (not sorted by ``id``) and
assigns YOLO class indices 0, 1, 2, ... in that same order.

Usage:
    poetry run coco-yolo-dataset-yaml \\
        --coco-json ./annotations/coco.json \\
        --path /data/my-dataset \\
        --kpt-shape 2 3 \\
        -o ./dataset.yaml

Detection-only (omit pose fields):

    poetry run coco-yolo-dataset-yaml \\
        -i ./annotations/coco.json \\
        -p /data/my-dataset \\
        --no-keypoints \\
        -o ./dataset.yaml
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import List


def _load_coco_categories(coco_json_path: Path) -> List[dict]:
    with open(coco_json_path, "r", encoding="utf-8") as f:
        coco = json.load(f)
    categories = coco.get("categories")
    if not categories:
        print("Error: COCO JSON has no 'categories' or it is empty.", file=sys.stderr)
        sys.exit(1)
    return list(categories)


def _yaml_scalar(s: str) -> str:
    """Emit a YAML scalar; quote when needed."""
    if s == "" or any(c in s for c in (":", "#", "\n", '"', "'")):
        return json.dumps(s)
    return s


def _format_int_list(values: List[int]) -> str:
    inner = ", ".join(str(v) for v in values)
    return f"[{inner}]"


def build_yaml_lines(
    dataset_path: str,
    train_subdir: str,
    val_subdir: str,
    names_by_index: List[str],
    keypoints: bool,
    kpt_shape: tuple[int, int] | None,
    flip_idx: List[int] | None,
) -> List[str]:
    lines: List[str] = [
        f"path: {_yaml_scalar(dataset_path)}",
        f"train: {_yaml_scalar(train_subdir)}",
        f"val: {_yaml_scalar(val_subdir)}",
    ]
    if keypoints:
        assert kpt_shape is not None
        lines.append(f"kpt_shape: {_format_int_list([kpt_shape[0], kpt_shape[1]])}")
        assert flip_idx is not None
        lines.append(f"flip_idx: {_format_int_list(flip_idx)}")
    lines.append("names:")
    for i, name in enumerate(names_by_index):
        lines.append(f"  {i}: {_yaml_scalar(name)}")
    return lines


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate a YOLO dataset YAML from COCO categories (COCO list order -> indices 0..N-1)."
    )
    parser.add_argument(
        "--coco-json",
        "-i",
        type=Path,
        required=True,
        help="Path to COCO JSON (uses categories[] order, not category id order).",
    )
    parser.add_argument(
        "--path",
        "-p",
        type=str,
        required=True,
        help="Dataset root directory (written as `path:` in the YAML).",
    )
    parser.add_argument(
        "--train-subdir",
        type=str,
        default="images/train",
        help="Train images path relative to dataset root (default: images/train).",
    )
    parser.add_argument(
        "--val-subdir",
        type=str,
        default="images/val",
        help="Val images path relative to dataset root (default: images/val).",
    )
    keygroup = parser.add_mutually_exclusive_group()
    keygroup.add_argument(
        "--keypoints",
        dest="keypoints",
        action="store_true",
        help="Include kpt_shape and flip_idx (default).",
    )
    keygroup.add_argument(
        "--no-keypoints",
        dest="keypoints",
        action="store_false",
        help="Omit pose fields; detection-style YAML only.",
    )
    parser.set_defaults(keypoints=True)
    parser.add_argument(
        "--kpt-shape",
        type=int,
        nargs=2,
        metavar=("NUM_KEYPOINTS", "DIMS"),
        help="Keypoint layout [num_keypoints, dims], e.g. `--kpt-shape 2 3`. Required with --keypoints.",
    )
    parser.add_argument(
        "--flip-idx",
        type=str,
        default=None,
        help=(
            "Comma-separated keypoint indices for horizontal flip, same length as num keypoints "
            "(e.g. `0,1` for identity). Default: 0,1,...,N-1 from --kpt-shape."
        ),
    )
    parser.add_argument(
        "--output",
        "-o",
        type=Path,
        required=True,
        help="Output YAML file path.",
    )

    args = parser.parse_args()

    if not args.keypoints and args.kpt_shape is not None:
        print("Error: --kpt-shape is only valid with keypoints enabled (omit --no-keypoints).", file=sys.stderr)
        sys.exit(1)
    if not args.keypoints and args.flip_idx is not None:
        print("Error: --flip-idx is only valid with keypoints enabled (omit --no-keypoints).", file=sys.stderr)
        sys.exit(1)

    if not args.coco_json.is_file():
        print(f"Error: COCO JSON not found: {args.coco_json}", file=sys.stderr)
        sys.exit(1)

    if args.keypoints and args.kpt_shape is None:
        print("Error: --kpt-shape NUM DIMS is required when keypoints are enabled.", file=sys.stderr)
        sys.exit(1)

    categories = _load_coco_categories(args.coco_json)
    names = [str(c["name"]) for c in categories]

    kpt_shape: tuple[int, int] | None = None
    flip_idx: List[int] | None = None
    if args.keypoints:
        n_kpt, dims = args.kpt_shape[0], args.kpt_shape[1]
        if n_kpt < 1 or dims < 1:
            print("Error: kpt_shape values must be positive integers.", file=sys.stderr)
            sys.exit(1)
        kpt_shape = (n_kpt, dims)
        if args.flip_idx is not None:
            parts = [p.strip() for p in args.flip_idx.split(",") if p.strip() != ""]
            try:
                flip_idx = [int(p) for p in parts]
            except ValueError:
                print("Error: --flip-idx must be comma-separated integers.", file=sys.stderr)
                sys.exit(1)
        else:
            flip_idx = list(range(n_kpt))
        if len(flip_idx) != n_kpt:
            print(
                f"Error: flip_idx has {len(flip_idx)} entries but num keypoints is {n_kpt}.",
                file=sys.stderr,
            )
            sys.exit(1)

    print("YOLO class index (COCO categories list order):")
    for i, name in enumerate(names):
        print(f"  {i}: {name}")

    lines = build_yaml_lines(
        dataset_path=args.path,
        train_subdir=args.train_subdir,
        val_subdir=args.val_subdir,
        names_by_index=names,
        keypoints=args.keypoints,
        kpt_shape=kpt_shape,
        flip_idx=flip_idx,
    )
    text = "\n".join(lines) + "\n"

    out_path = args.output
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(text, encoding="utf-8")
    print(f"Wrote {out_path}")


if __name__ == "__main__":
    main()
