import os
from .wrapper import DatasetWrapper, DataSplit
from PIL import Image
from pycocotools.coco import COCO
import contextlib

class SignWrapper(DatasetWrapper):
    def __init__(self, root, data_split='train'):
        self.root = root
        self.data_split = DataSplit(data_split)
        if self.data_split not in DataSplit:
            raise ValueError(f"Invalid data split: {data_split}. Must be one of {list(DataSplit)}")

        # Assume all splits share the same annotation file
        ann_file = os.path.join(self.root, data_split, '_annotations.coco.json')
        self.signs = COCO(ann_file)
        # with open(os.devnull, 'w') as devnull:
        #     with contextlib.redirect_stdout(devnull):
                
        self.ids = list(sorted(self.signs.imgs.keys()))
        self.categories = self.signs.loadCats(self.signs.getCatIds())

        self.__format_annotations__()

    def __format_annotations__(self):
        # Roboflow typically stores all images in a single directory, e.g., "images"
        image_base_path = os.path.join(self.root, 'images')

        print(f'[SignWrapper] Formatting annotations for {self.data_split.value} (N={len(self.ids)}) from {image_base_path}')

        annotations = {}
        for img_id in self.ids:
            img_info = self.signs.imgs[img_id]
            ann_ids = self.signs.getAnnIds(imgIds=img_id)
            anns = self.signs.loadAnns(ann_ids)

            bboxes = []
            category_ids = []
            poison_masks = []
            target_ids = []

            for ann in anns:
                bboxes.append(ann['bbox'])
                category_ids.append(ann['category_id'])
                poison_masks.append(False)
                target_ids.append(-1)

            annotations[img_id] = [{
                'sub_id': 0,
                'bbox': bboxes,
                'category_id': category_ids,
                'poison_mask': poison_masks,
                'target_id': target_ids,
                'clean_img_path': os.path.join(image_base_path, img_info['file_name']),
                'bd_img_path': None
            }]
            
        self.annotations = annotations
