import json
import os
import pathlib
import re
from dataclasses import dataclass


@dataclass(frozen=True, eq=True)
class BlenderLineImage:
    """Reference to generated image within instance folder."""

    id: str
    path: pathlib.Path


@dataclass(frozen=True, eq=True)
class BlenderLineMask:
    """Reference to generated mask within instance folder."""

    id: str
    path: pathlib.Path
    label: str  # Label index (string, as label indices in label_mapping.json are strings)


def get_blenderline_image(instance_path: pathlib.Path) -> BlenderLineImage:
    """Get reference to generated image in a BlenderLine instance folder.

    Args:
        instance_path (pathlib.Path): path to BlenderLine instance folder.

    Returns:
        BlenderLineImage: reference to image containing path and image ID.
    """
    # BlenderLine names images with format "image__<image ID>__0001.png", where the image
    # ID is a 6 byte random hex string.
    image_pattern = re.compile(r"image__([0-9a-f]{12})__0001.png")

    matches: list[BlenderLineImage] = []

    for path in instance_path.iterdir():
        if match := image_pattern.search(path.name):
            matches.append(BlenderLineImage(id=match.group(1), path=path))

    # Check that the instance contains exactly one generated image.
    if len(matches) != 1:
        raise Exception(
            f"Instance at {instance_path} does not contain a valid BlenderLine image."
        )

    return matches[0]


def get_blenderline_masks(instance_path: pathlib.Path) -> list[BlenderLineMask]:
    """Get references to generated masks in a BlenderLine instance folder.

    Args:
        instance_path (pathlib.Path): path to BlenderLine instance folder.

    Returns:
        list[BlenderLineMask]: references to masks containing path, mask ID, and label ID.
    """
    # BlenderLine names masks with format "mask__<label ID>__<image ID>__0001.png", where
    # the mask ID is a 6 byte random hex string, and the label ID must be a digit.
    # TODO: there is currently no check for label IDs that are not digits, which are
    # incompatible with the YOLO metadata format.
    mask_pattern = re.compile(r"mask__([0-9]{1,10})__([0-9a-f]{12})__0001.png")

    matches: list[BlenderLineMask] = []

    for path in instance_path.iterdir():
        if match := mask_pattern.search(path.name):
            matches.append(
                BlenderLineMask(id=match.group(2), path=path, label=match.group(1))
            )

    return matches


def get_label_mapping(source_path: pathlib.Path) -> dict[str, str]:
    """Get label mapping generated by BlenderLine.

    Args:
        source_path (pathlib.Path): path to BlenderLine generated dataset root folder.

    Returns:
        dict[str, str]: mappign between label IDs and human-readable label names.
    """
    # Check for label mapping file existence.
    if not "label_mapping.json" in os.listdir(source_path):
        raise Exception(
            f"Source folder at {source_path} does not contain a label mapping."
        )

    with open(source_path / "label_mapping.json") as file:
        label_mapping = json.load(file)

    return label_mapping
