"""Token-level scorers for comparing output against reference."""

import re
from collections import Counter

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


class TokenF1(Scorer):
    """
    Token-level F1 score (harmonic mean of precision and recall).

    Computes token overlap between output and reference using whitespace
    tokenization. Standard QA metric (used by SQuAD).

    Returns a float in [0.0, 1.0].

    Examples:
        >>> scorer = TokenF1()
        >>> run = Run(
        ...     id="1", input="q",
        ...     output="the quick brown fox",
        ...     reference="the quick fox jumps"
        ... )
        >>> score = scorer.score(run)
        >>> score.metadata["precision"]
        0.75
    """

    name = "token_f1"
    requires_reference = True
    cacheable = False

    def __init__(self, ignore_case: bool = False, ignore_punctuation: bool = False):
        """
        Args:
            ignore_case: Lowercase both strings before tokenizing.
            ignore_punctuation: Remove punctuation before tokenizing.
        """
        self.ignore_case = ignore_case
        self.ignore_punctuation = ignore_punctuation

    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)

        if self.ignore_case:
            output_str = output_str.lower()
            reference_str = reference_str.lower()

        if self.ignore_punctuation:
            output_str = re.sub(r"[^\w\s]", "", output_str)
            reference_str = re.sub(r"[^\w\s]", "", reference_str)

        output_tokens = Counter(output_str.split())
        reference_tokens = Counter(reference_str.split())

        # Token overlap (intersection with counts)
        common = sum((output_tokens & reference_tokens).values())
        num_output = sum(output_tokens.values())
        num_reference = sum(reference_tokens.values())

        if num_output == 0 and num_reference == 0:
            return Score(value=1.0, reason="Both empty", metadata={
                "precision": 1.0, "recall": 1.0, "common_tokens": 0,
            })

        precision = common / num_output if num_output > 0 else 0.0
        recall = common / num_reference if num_reference > 0 else 0.0

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

        return Score(
            value=f1,
            reason=f"F1: {f1:.4f} (P={precision:.4f}, R={recall:.4f})",
            metadata={
                "precision": precision,
                "recall": recall,
                "common_tokens": common,
                "output_tokens": num_output,
                "reference_tokens": num_reference,
            },
        )


class LengthRatio(Scorer):
    """
    Ratio of output length to reference length (character-based).

    A ratio of 1.0 means the output and reference are the same length.
    Values > 1.0 indicate the output is longer; < 1.0 indicates shorter.

    Examples:
        >>> scorer = LengthRatio()
        >>> run = Run(id="1", input="q", output="hello world", reference="hello")
        >>> score = scorer.score(run)
        >>> score.value
        2.2
    """

    name = "length_ratio"
    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_len = len(str(run.output))
        reference_len = len(str(run.reference))

        if reference_len == 0:
            ratio = 0.0 if output_len == 0 else float("inf")
            reason = "Reference is empty"
        else:
            ratio = output_len / reference_len
            if ratio > 1.0:
                reason = f"Output is {ratio:.2f}x longer than reference"
            elif ratio < 1.0:
                reason = f"Output is {ratio:.2f}x shorter than reference"
            else:
                reason = "Output and reference are the same length"

        return Score(
            value=ratio,
            reason=reason,
            metadata={
                "output_length": output_len,
                "reference_length": reference_len,
            },
        )
