# scripts/convert_roboflow_to_triplets.py
import json
import shutil
from pathlib import Path
from PIL import Image
import numpy as np

def convert_roboflow_to_triplets(roboflow_dir: str, output_dir: str):
    """
    Convert Roboflow COCO format to patch triplet format.
    
    Roboflow COCO structure:
        train/
            _annotations.coco.json
            image1.jpg
            image2.jpg
    """
    roboflow_path = Path(roboflow_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Load COCO annotations
    anno_file = roboflow_path / "train" / "_annotations.coco.json"
    with open(anno_file) as f:
        coco = json.load(f)
    
    # Build image ID → filename map
    images = {img['id']: img['file_name'] for img in coco['images']}
    
    # Group annotations by image
    from collections import defaultdict
    image_annotations = defaultdict(list)
    for ann in coco['annotations']:
        image_annotations[ann['image_id']].append(ann)
    
    triplet_id = 0
    for img_id, filename in images.items():
        img_path = roboflow_path / "train" / filename
        if not img_path.exists():
            continue
        
        img = Image.open(img_path).convert('RGB')
        img_array = np.array(img)
        H, W = img_array.shape[:2]
        
        anns = image_annotations.get(img_id, [])
        if not anns:
            continue
        
        for ann in anns:
            # 1. Skip non-defect category annotations (category_id == 2)
            if ann.get('category_id') == 2:
                continue

            # 2. Skip full-image bounding boxes
            x, y, w, h = ann['bbox']
            if w >= W and h >= H:
                continue
            
            triplet_id += 1
            triplet_dir = output_path / f"defect_{triplet_id:04d}"
            triplet_dir.mkdir(exist_ok=True)
            
            # Get bbox
            x, y, w, h = ann['bbox']
            x1, y1, x2, y2 = int(x), int(y), int(x+w), int(y+h)
            
            # Create mask from segmentation if available, else bbox
            if 'segmentation' in ann and ann['segmentation']:
                # COCO polygon segmentation
                from pycocotools import mask as maskUtils
                rles = maskUtils.frPyObjects(ann['segmentation'], H, W)
                mask = maskUtils.decode(rles)
                if len(mask.shape) == 3:
                    mask = np.any(mask, axis=2).astype(np.uint8) * 255
            else:
                # Fallback: bbox mask
                mask = np.zeros((H, W), dtype=np.uint8)
                mask[y1:y2, x1:x2] = 255
            
            # Save clean image (original_target)
            img.save(triplet_dir / "original_target.png")
            
            # Save mask (original_masked)
            Image.fromarray(mask).save(triplet_dir / "original_masked.png")
            
            # For artifact_target, we need the defective version.
            # Since this is a real defect dataset, the original image IS the defect.
            # For synthetic training, you may want to inpaint the defect out to create "clean",
            # but for RAG retrieval, we can use the same image as artifact_target.
            img.save(triplet_dir / "artifact_target.png")
            
            print(f"Created triplet {triplet_id}: {filename} → {triplet_dir}")
    
    print(f"\nTotal triplets created: {triplet_id}")
    print(f"Output: {output_path}")

if __name__ == "__main__":
    convert_roboflow_to_triplets(
        roboflow_dir="./data/external/Manufacturing_Defect_Detection",
        output_dir="data/external/roboflow_manufacturing"
    )