from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from PIL import Image


def _detections_to_coco(
    detections: Any,  # supervision.Detections
    image_path: Path,
    image_id: int,
    category_id: int = 0,  # Assuming single class 'knot'
) -> dict[str, Any]:
    """Convert supervision.Detections to COCO annotation format for one image."""
    annotations = []
    for i in range(len(detections)):
        xyxy = detections.xyxy[i].tolist()
        conf = float(detections.confidence[i]) if detections.confidence is not None else 1.0
        # COCO bbox: [x_min, y_min, width, height]
        x_min, y_min, x_max, y_max = xyxy
        width = x_max - x_min
        height = y_max - y_min
        bbox = [x_min, y_min, width, height]

        ann = {
            "id": image_id * 1000 + i,  # Simple ID scheme
            "image_id": image_id,
            "category_id": category_id,
            "bbox": bbox,
            "area": width * height,
            "iscrowd": 0,
            "score": conf,  # Optional, for model confidence
        }
        annotations.append(ann)

    return {"annotations": annotations}


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Auto-label new images with trained RF-DETR, output COCO JSON."
    )
    parser.add_argument("--weights", type=Path, required=True, help="Path to trained checkpoint")
    parser.add_argument("--images-dir", type=Path, required=True, help="Directory with new images")
    parser.add_argument("--output-json", type=Path, required=True, help="Output COCO JSON file")
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--category-name", default="knot", help="Class name for annotations")
    args = parser.parse_args()

    if not args.weights.exists():
        raise SystemExit(f"Weights not found: {args.weights}")
    if not args.images_dir.exists():
        raise SystemExit(f"Images dir not found: {args.images_dir}")

    from rfdetr import RFDETRBase

    model = RFDETRBase(pretrain_weights=str(args.weights))

    # Collect all images
    image_paths = list(args.images_dir.glob("*.jpg")) + list(args.images_dir.glob("*.png"))
    if not image_paths:
        raise SystemExit(f"No .jpg or .png images found in {args.images_dir}")

    coco_data = {
        "images": [],
        "annotations": [],
        "categories": [{"id": 0, "name": args.category_name, "supercategory": "none"}],
    }

    ann_id_counter = 0
    for img_id, img_path in enumerate(sorted(image_paths)):
        image = Image.open(img_path).convert("RGB")
        detections = model.predict(image, threshold=args.threshold)

        # Add image entry
        width, height = image.size
        coco_data["images"].append({
            "id": img_id,
            "file_name": img_path.name,
            "width": width,
            "height": height,
        })

        # Convert detections to annotations
        ann_data = _detections_to_coco(detections, img_path, img_id, category_id=0)
        for ann in ann_data["annotations"]:
            ann["id"] = ann_id_counter
            ann_id_counter += 1
            coco_data["annotations"].append(ann)

    # Write COCO JSON
    with args.output_json.open("w", encoding="utf-8") as f:
        json.dump(coco_data, f, indent=2)

    print(f"Auto-labeled {len(image_paths)} images -> {args.output_json}")
    print(f"Total annotations: {len(coco_data['annotations'])}")
    return 0


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