import json
import os
from typing import Dict, List, Optional

from tqdm import tqdm


def discretize_categories(categories: List[Dict[str, int]]) -> Dict[int, int]:
    """
    Maps each unique 'id' in the list of category dictionaries to a sequential integer index.
    Indices are assigned based on the sorted 'id' values.
    """
    sorted_categories = sorted(categories, key=lambda category: category["id"])
    return {category["id"]: index for index, category in enumerate(sorted_categories)}


def process_annotations(
    image_annotations: Dict[int, List[Dict]],
    image_info_dict: Dict[int, tuple],
    output_dir: str,
    id_to_idx: Optional[Dict[int, int]] = None,
) -> None:
    """
    Process and save annotations to files, with option to remap category IDs.
    """
    for image_id, annotations in tqdm(image_annotations.items(), desc="Processing annotations"):
        file_path = os.path.join(output_dir, f"{image_id:0>12}.txt")
        if not annotations:
            continue
        with open(file_path, "w") as file:
            for annotation in annotations:
                process_annotation(annotation, image_info_dict[image_id], id_to_idx, file)


def process_annotation(annotation: Dict, image_dims: tuple, id_to_idx: Optional[Dict[int, int]], file) -> None:
    """
    Convert a single annotation's segmentation and write it to the open file handle.
    """
    category_id = annotation["category_id"]
    segmentation = (
        annotation["segmentation"][0]
        if annotation["segmentation"] and isinstance(annotation["segmentation"][0], list)
        else None
    )

    if segmentation is None:
        return

    img_width, img_height = image_dims
    normalized_segmentation = normalize_segmentation(segmentation, img_width, img_height)

    if id_to_idx:
        category_id = id_to_idx.get(category_id, category_id)

    file.write(f"{category_id} {' '.join(normalized_segmentation)}\n")


def normalize_segmentation(segmentation: List[float], img_width: int, img_height: int) -> List[str]:
    """
    Normalize and format segmentation coordinates.
    """
    normalized = [
        f"{coord / img_width:.6f}" if index % 2 == 0 else f"{coord / img_height:.6f}"
        for index, coord in enumerate(segmentation)
    ]
    return normalized


def convert_annotations(json_file: str, output_dir: str) -> None:
    """
    Load annotation data from a JSON file and process all annotations.
    """
    with open(json_file) as file:
        data = json.load(file)

    os.makedirs(output_dir, exist_ok=True)

    image_info_dict = {img["id"]: (img["width"], img["height"]) for img in data.get("images", [])}
    id_to_idx = discretize_categories(data.get("categories", [])) if "categories" in data else None
    image_annotations = {img_id: [] for img_id in image_info_dict}

    for annotation in data.get("annotations", []):
        if not annotation.get("iscrowd", False):
            image_annotations[annotation["image_id"]].append(annotation)

    process_annotations(image_annotations, image_info_dict, output_dir, id_to_idx)


if __name__ == "__main__":
    convert_annotations("./data/coco/annotations/instances_train2017.json", "./data/coco/labels/train2017/")
    convert_annotations("./data/coco/annotations/instances_val2017.json", "./data/coco/labels/val2017/")