import argparse
import json
import sys
from pathlib import Path

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


def main() -> int:
    parser = argparse.ArgumentParser(description="Check COCO segmentation dataset structure.")
    parser.add_argument("--dataset-dir", required=True)
    args = parser.parse_args()
    dataset_dir = Path(args.dataset_dir)
    for split in ["train", "valid", "test"]:
        ann = dataset_dir / split / "_annotations.coco.json"
        if not ann.exists():
            raise FileNotFoundError(f"missing: {ann}")
        data = json.loads(ann.read_text(encoding="utf-8"))
        if "images" not in data or "annotations" not in data or "categories" not in data:
            raise ValueError(f"invalid COCO json: {ann}")
    print("COCO segmentation dataset structure looks valid.")
    return 0


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