from typing import Any, Callable, Dict, Literal, Optional

from typing_extensions import TypedDict

METRICS = Literal[
    "damerau_levenshtein",
    "levenshtein",
    "jaro",
    "jaro_winkler",
    "hamming",
    "indel",
]


class EditDistanceConfig(TypedDict, total=False):
    metric: METRICS
    normalize_score: bool


class EditDistance:
    def __init__(
        self,
        config: Optional[EditDistanceConfig] = None,
    ):
        config = config or {}
        metric = config.get("metric") or "damerau_levenshtein"
        self.metric = self._get_metric(
            metric, normalize_score=config.get("normalize_score", True)
        )

    def evaluate(
        self,
        prediction: str,
        reference: Optional[str] = None,
    ) -> float:
        return self.metric(prediction, reference)

    @staticmethod
    def _get_metric(distance: str, normalize_score: bool = True) -> Callable:
        try:
            from rapidfuzz import (  # type: ignore[import-not-found]
                distance as rf_distance,
            )
        except ImportError:
            raise ImportError(
                "This operation requires the rapidfuzz library to use."
                "Please install it with `pip install -U rapidfuzz`."
            )

        module_map: Dict[str, Any] = {
            "damerau_levenshtein": rf_distance.DamerauLevenshtein,
            "levenshtein": rf_distance.Levenshtein,
            "jaro": rf_distance.Jaro,
            "jaro_winkler": rf_distance.JaroWinkler,
            "hamming": rf_distance.Hamming,
            "indel": rf_distance.Indel,
        }
        if distance not in module_map:
            raise ValueError(
                f"Invalid distance metric: {distance}"
                f"\nMust be one of: {list(module_map)}"
            )
        module = module_map[distance]
        if normalize_score:
            return module.normalized_distance
        else:
            return module.distance