"""Text similarity scorers for comparing output against reference."""

import math
from collections import Counter
from difflib import SequenceMatcher
from typing import Optional

from openevalkit.scorers.base import Scorer
from openevalkit.score import Score
from openevalkit.run import Run
from openevalkit.errors import ScoringError


class LevenshteinDistance(Scorer):
    """
    Normalized Levenshtein similarity between output and reference.

    Returns a float in [0.0, 1.0] where 1.0 means identical strings.
    Computed as ``1 - (edit_distance / max(len(output), len(reference)))``.

    Examples:
        >>> scorer = LevenshteinDistance()
        >>> run = Run(id="1", input="q", output="kitten", reference="sitting")
        >>> score = scorer.score(run)
        >>> 0.0 <= score.value <= 1.0
        True
    """

    name = "levenshtein"
    requires_reference = True
    cacheable = False

    def score(self, run: Run) -> Score:
        if run.reference is None:
            raise ScoringError(
                f"{self.name} requires reference field. "
                f"Run {run.id} has no reference."
            )

        output_str = str(run.output)
        reference_str = str(run.reference)

        dist = self._edit_distance(output_str, reference_str)
        max_len = max(len(output_str), len(reference_str))
        similarity = 1.0 - (dist / max_len) if max_len > 0 else 1.0

        return Score(
            value=similarity,
            reason=f"Edit distance {dist}, similarity {similarity:.4f}",
            metadata={
                "edit_distance": dist,
                "output_length": len(output_str),
                "reference_length": len(reference_str),
            },
        )

    @staticmethod
    def _edit_distance(a: str, b: str) -> int:
        """Wagner-Fischer algorithm with O(min(m,n)) space."""
        if len(a) < len(b):
            a, b = b, a

        prev = list(range(len(b) + 1))
        for i in range(1, len(a) + 1):
            curr = [i] + [0] * len(b)
            for j in range(1, len(b) + 1):
                cost = 0 if a[i - 1] == b[j - 1] else 1
                curr[j] = min(
                    curr[j - 1] + 1,       # insertion
                    prev[j] + 1,            # deletion
                    prev[j - 1] + cost,     # substitution
                )
            prev = curr
        return prev[len(b)]


class FuzzyMatch(Scorer):
    """
    Fuzzy string similarity using Python's ``difflib.SequenceMatcher``.

    Returns a float in [0.0, 1.0] where 1.0 means identical strings.
    Handles reordering and partial matches better than edit distance.

    Examples:
        >>> scorer = FuzzyMatch()
        >>> run = Run(id="1", input="q", output="hello world", reference="hello world!")
        >>> score = scorer.score(run)
        >>> score.value > 0.9
        True
    """

    name = "fuzzy_match"
    requires_reference = True
    cacheable = False

    def score(self, run: Run) -> Score:
        if run.reference is None:
            raise ScoringError(
                f"{self.name} requires reference field. "
                f"Run {run.id} has no reference."
            )

        output_str = str(run.output)
        reference_str = str(run.reference)
        ratio = SequenceMatcher(None, output_str, reference_str).ratio()

        return Score(
            value=ratio,
            reason=f"Fuzzy similarity {ratio:.4f}",
            metadata={
                "output_length": len(output_str),
                "reference_length": len(reference_str),
            },
        )


class BLEU(Scorer):
    """
    BLEU score measuring n-gram precision with brevity penalty.

    Standard machine translation / NLG metric. Returns a float in [0.0, 1.0].

    Examples:
        >>> scorer = BLEU()
        >>> run = Run(
        ...     id="1", input="translate",
        ...     output="the cat sat on the mat",
        ...     reference="the cat is on the mat"
        ... )
        >>> score = scorer.score(run)
        >>> 0.0 <= score.value <= 1.0
        True
    """

    name = "bleu"
    requires_reference = True
    cacheable = False

    def __init__(self, max_n: int = 4):
        """
        Args:
            max_n: Maximum n-gram order. Default: 4.
        """
        self.max_n = max_n

    def score(self, run: Run) -> Score:
        if run.reference is None:
            raise ScoringError(
                f"{self.name} requires reference field. "
                f"Run {run.id} has no reference."
            )

        output_tokens = str(run.output).split()
        reference_tokens = str(run.reference).split()

        if not output_tokens or not reference_tokens:
            return Score(
                value=0.0,
                reason="Empty output or reference",
                metadata={"brevity_penalty": 0.0},
            )

        # Modified n-gram precisions
        precisions = []
        for n in range(1, self.max_n + 1):
            p = self._modified_precision(output_tokens, reference_tokens, n)
            precisions.append(p)

        # If any precision is 0, BLEU is 0
        if any(p == 0.0 for p in precisions):
            bleu = 0.0
        else:
            # Geometric mean of precisions (uniform weights)
            log_avg = sum(math.log(p) for p in precisions) / len(precisions)
            bleu = math.exp(log_avg)

        # Brevity penalty
        bp = self._brevity_penalty(len(output_tokens), len(reference_tokens))
        bleu *= bp

        return Score(
            value=bleu,
            reason=f"BLEU-{self.max_n}: {bleu:.4f}",
            metadata={
                "precisions": precisions,
                "brevity_penalty": bp,
                "output_length": len(output_tokens),
                "reference_length": len(reference_tokens),
            },
        )

    @staticmethod
    def _get_ngrams(tokens, n):
        """Extract n-grams from token list."""
        return [tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)]

    @staticmethod
    def _modified_precision(output_tokens, reference_tokens, n):
        """Clipped n-gram precision."""
        output_ngrams = BLEU._get_ngrams(output_tokens, n)
        reference_ngrams = BLEU._get_ngrams(reference_tokens, n)

        if not output_ngrams:
            return 0.0

        ref_counts = Counter(reference_ngrams)
        out_counts = Counter(output_ngrams)

        clipped = sum(
            min(count, ref_counts[ngram]) for ngram, count in out_counts.items()
        )
        total = sum(out_counts.values())

        return clipped / total if total > 0 else 0.0

    @staticmethod
    def _brevity_penalty(output_len, reference_len):
        """Brevity penalty for short outputs."""
        if output_len >= reference_len:
            return 1.0
        return math.exp(1 - reference_len / output_len)


class ROUGE(Scorer):
    """
    ROUGE-L score based on longest common subsequence (LCS).

    Returns the F1 measure of LCS-based precision and recall.
    Standard summarization metric. Returns a float in [0.0, 1.0].

    Examples:
        >>> scorer = ROUGE()
        >>> run = Run(
        ...     id="1", input="summarize",
        ...     output="the cat sat on the mat",
        ...     reference="the cat is on the mat"
        ... )
        >>> score = scorer.score(run)
        >>> 0.0 <= score.value <= 1.0
        True
    """

    name = "rouge_l"
    requires_reference = True
    cacheable = False

    def score(self, run: Run) -> Score:
        if run.reference is None:
            raise ScoringError(
                f"{self.name} requires reference field. "
                f"Run {run.id} has no reference."
            )

        output_tokens = str(run.output).split()
        reference_tokens = str(run.reference).split()

        if not output_tokens or not reference_tokens:
            return Score(
                value=0.0,
                reason="Empty output or reference",
                metadata={"precision": 0.0, "recall": 0.0, "lcs_length": 0},
            )

        lcs_len = self._lcs_length(output_tokens, reference_tokens)
        precision = lcs_len / len(output_tokens)
        recall = lcs_len / len(reference_tokens)

        if precision + recall == 0:
            f1 = 0.0
        else:
            f1 = 2 * precision * recall / (precision + recall)

        return Score(
            value=f1,
            reason=f"ROUGE-L F1: {f1:.4f} (P={precision:.4f}, R={recall:.4f})",
            metadata={
                "precision": precision,
                "recall": recall,
                "lcs_length": lcs_len,
                "output_length": len(output_tokens),
                "reference_length": len(reference_tokens),
            },
        )

    @staticmethod
    def _lcs_length(a, b):
        """Compute length of longest common subsequence using DP."""
        m, n = len(a), len(b)
        prev = [0] * (n + 1)
        for i in range(1, m + 1):
            curr = [0] * (n + 1)
            for j in range(1, n + 1):
                if a[i - 1] == b[j - 1]:
                    curr[j] = prev[j - 1] + 1
                else:
                    curr[j] = max(curr[j - 1], prev[j])
            prev = curr
        return prev[n]
