"""Exact Match metric implementation for Scorebook."""

import string
from typing import Any, Dict, List, Tuple

from scorebook.metrics import MetricBase, scorebook_metric


@scorebook_metric
class ExactMatch(MetricBase):
    """Exact Match metric for evaluating string predictions.

    Compares strings for exact equality with optional preprocessing.

    Args:
        case_insensitive: If True, convert strings to lowercase before comparison.
            Defaults to True.
        strip: If True, strip leading and trailing whitespace before comparison.
            Defaults to True.
        strip_punctuation: If True, strip leading and trailing punctuation before
            comparison. Defaults to False.
    """

    @property
    def name(self) -> str:
        """Return the metric name."""
        return "exact_match"

    def __init__(
        self,
        case_insensitive: bool = True,
        strip: bool = True,
        strip_punctuation: bool = False,
    ) -> None:
        """Initialize ExactMatch metric with preprocessing options.

        Args:
            case_insensitive: If True, convert strings to lowercase before comparison.
                Defaults to True.
            strip: If True, strip leading and trailing whitespace before comparison.
                Defaults to True.
            strip_punctuation: If True, strip leading and trailing punctuation before
                comparison. Defaults to False.
        """
        self.case_insensitive = case_insensitive
        self.strip = strip
        self.strip_punctuation = strip_punctuation

    def _preprocess(self, value: Any) -> Any:
        """Apply preprocessing to a value if it's a string.

        Args:
            value: The value to preprocess.

        Returns:
            The preprocessed value (string) or original value (non-string).
        """
        if not isinstance(value, str):
            return value

        result = value
        if self.strip:
            result = result.strip()
        if self.strip_punctuation:
            result = result.strip(string.punctuation)
        if self.case_insensitive:
            result = result.lower()
        return result

    def score(
        self, outputs: List[Any], labels: List[Any]
    ) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
        """Calculate the exact match score between predictions and references.

        Args:
            outputs: A list of inference outputs.
            labels: A list of ground truth labels.

        Returns:
            The aggregate exact match score for all items (matches / total).
            Per-item scores as dicts: [{"exact_match": True/False}, ...].
        """
        if not outputs:
            return {"exact_match": 0.0}, []

        # Calculate item scores with preprocessing
        matches = [
            self._preprocess(output) == self._preprocess(label)
            for output, label in zip(outputs, labels)
        ]
        item_scores = [{"exact_match": match} for match in matches]

        # Calculate aggregate score
        num_matches = sum(matches)
        total = len(outputs)
        aggregate_scores = {"exact_match": num_matches / total}

        return aggregate_scores, item_scores
