from typing import List
from typing import Tuple
from typing import Dict
from typing import Set


def _index_labels(
    original_labels: List[str],
    strip_prefixes: bool = False
) -> Tuple[Dict[str, int], Set[str], Set[str]]:
    """ Normalize and index a list of labels to:
    1. Generate a mapping from labels to indices
    2. Identify label prefixes (e.g., I, B, L, etc.)
    3. Identify list of labels without prefixes

    If `strip_prefixes` is True, we normalize IO/BIO/BILUO labels
    (e.g., B-PERSON and I-PERSON will be normalized PERSON). 
    
    If `strip_prefixes` is False, we assume the labels
    use IO/BIO/BILUO formats and mantain different labels/indices for labels
    like B-PERSON and I-PERSON.

    NB: We assume that the first label in labels is 'O' for the null token.
    """
    labels = []
    prefixes = set()
    labels_without_prefix = set()

    for original_label in original_labels:
        if "-" in original_label:
            # Normalize B-PER and I-PER to PER
            prefix, normalized_label = original_label.split("-", 1)
            prefixes.add(prefix)
        else:
            # No normalization required 
            normalized_label = original_label

        # Track labels without prefixes
        labels_without_prefix.add(normalized_label)

        if strip_prefixes:
            # Use normalized label for token labeling
            if normalized_label not in labels:
                labels.append(normalized_label)
        else:
            # Use original label for token labeling
            if original_label not in labels:
                labels.append(original_label)

    # Generate mapping of labels to label indices
    label2idx = {label: i for i, label in enumerate(labels)}

    return label2idx, prefixes, labels_without_prefix