from __future__ import annotations

import argparse
import json
from pathlib import Path

from PIL import Image


def coco_to_label_studio(
    image_paths: list[Path],
    coco_predictions_path: Path,
    category_name: str = "knot",
) -> list[dict]:
    """Convert COCO predictions to Label Studio import format with pre-annotations."""
    
    # Load COCO predictions
    with coco_predictions_path.open("r", encoding="utf-8") as f:
        coco_data = json.load(f)
    
    # Build image_id -> annotations mapping
    annotations_by_image = {}
    for ann in coco_data.get("annotations", []):
        img_id = ann["image_id"]
        if img_id not in annotations_by_image:
            annotations_by_image[img_id] = []
        annotations_by_image[img_id].append(ann)
    
    # Build image_id -> image info mapping
    image_info = {img["id"]: img for img in coco_data.get("images", [])}
    
    # Convert to Label Studio format
    ls_tasks = []
    for img_id, img_data in image_info.items():
        file_name = img_data["file_name"]
        width = img_data["width"]
        height = img_data["height"]
        
        # Find the actual file path
        matching_paths = [p for p in image_paths if p.name == file_name]
        if not matching_paths:
            continue
        
        img_path = matching_paths[0]
        
        # Build predictions (pre-annotations)
        predictions = []
        for ann in annotations_by_image.get(img_id, []):
            # COCO bbox: [x_min, y_min, width, height]
            x, y, w, h = ann["bbox"]
            
            # Label Studio uses percentages
            x_percent = (x / width) * 100
            y_percent = (y / height) * 100
            w_percent = (w / width) * 100
            h_percent = (h / height) * 100
            
            predictions.append({
                "value": {
                    "x": x_percent,
                    "y": y_percent,
                    "width": w_percent,
                    "height": h_percent,
                    "rectanglelabels": [category_name]
                },
                "from_name": "label",
                "to_name": "image",
                "type": "rectanglelabels",
                "score": ann.get("score", 1.0)
            })
        
        # Create Label Studio task with predictions
        task = {
            "data": {
                "image": f"/data/local-files/?d={img_path.absolute()}"
            },
            "predictions": [{
                "result": predictions,
                "model_version": "rfdetr-auto-label"
            }]
        }
        
        ls_tasks.append(task)
    
    return ls_tasks


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Convert COCO predictions to Label Studio format for review."
    )
    parser.add_argument("--coco-json", type=Path, required=True, help="COCO predictions JSON from auto_label_images.py")
    parser.add_argument("--images-dir", type=Path, required=True, help="Directory with the images")
    parser.add_argument("--output-json", type=Path, required=True, help="Output Label Studio tasks JSON")
    parser.add_argument("--category-name", default="knot", help="Label name in Label Studio")
    args = parser.parse_args()
    
    if not args.coco_json.exists():
        raise SystemExit(f"COCO JSON not found: {args.coco_json}")
    if not args.images_dir.exists():
        raise SystemExit(f"Images dir not found: {args.images_dir}")
    
    # 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 images found in {args.images_dir}")
    
    # Convert
    ls_tasks = coco_to_label_studio(image_paths, args.coco_json, args.category_name)
    
    # Write Label Studio import format
    with args.output_json.open("w", encoding="utf-8") as f:
        json.dump(ls_tasks, f, indent=2)
    
    print(f"✓ Converted {len(ls_tasks)} tasks with predictions")
    print(f"✓ Output: {args.output_json}")
    print(f"\nImport into Label Studio:")
    print(f"  1. Open your project")
    print(f"  2. Settings → Storage → Add Local Storage")
    print(f"     Absolute path: {args.images_dir.absolute()}")
    print(f"  3. Import → Upload Files → select {args.output_json.name}")
    print(f"  4. Start reviewing predictions!")
    
    return 0


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