"""Extractors used by composed evaluators.

Provides standalone extraction functions (``extract_mc_answer``,
``extract_yes_no_answer``) and thin ``IOutputExtractor`` wrappers
(``MCAnswerExtractor``, ``YesNoExtractor``) for use with
:class:`~ember.utils.eval.evaluators.ComposedEvaluator`.
"""

from __future__ import annotations

import abc
import re
from collections.abc import Sequence
from typing import Generic, TypeVar

T_out = TypeVar("T_out")
T_truth = TypeVar("T_truth")


class IOutputExtractor(Generic[T_out, T_truth], metaclass=abc.ABCMeta):
    """Convert raw system output into a comparable form."""

    @abc.abstractmethod
    def extract(self, system_output: T_out, **kwargs: object) -> T_truth:
        raise NotImplementedError


class RegexExtractor(IOutputExtractor[str, str]):
    """Extract the first capturing group from a regex match."""

    def __init__(self, pattern: str) -> None:
        self.compiled_pattern: re.Pattern[str] = re.compile(pattern)

    def extract(self, system_output: str, **kwargs: object) -> str:
        match = self.compiled_pattern.search(system_output)
        if match is None:
            return ""
        return str(match.group(1))


class FinalLetterExtractor(IOutputExtractor[str, str]):
    """Extract the final multiple-choice letter from free-form text."""

    _MARKDOWN_PATTERN: re.Pattern[str] = re.compile(r"\*\*|\*|__|_|~~|`")

    def __init__(self, valid_letters: str = "ABCD") -> None:
        collected = [letter.upper() for letter in valid_letters if letter.isalpha()]
        if not collected:
            collected = list("ABCD")
        self.valid_letters = tuple(collected)
        letter_class = "".join(self.valid_letters)
        self._patterns: Sequence[re.Pattern[str]] = (
            re.compile(rf"\\boxed\{{([{letter_class}])\}}", re.IGNORECASE),
            re.compile(
                rf"final\s*answer\s*(?:is\s+|[:=]\s*)\(?([{letter_class}])\)?",
                re.IGNORECASE,
            ),
            re.compile(
                rf"(?:the\s+)?(?:correct\s+)?(?:answer|statement|option)\s+is\s+"
                rf"\(?([{letter_class}])\)?(?:\s|$|\.)",
                re.IGNORECASE,
            ),
            re.compile(
                rf"(?:answer|choice)\s*[:=]\s*\(?([{letter_class}])\)?"
                rf"(?!\s*(?:and|/|,)\s*[{letter_class}])",
                re.IGNORECASE,
            ),
            re.compile(
                rf"option\s+([{letter_class}])\s+is\s+(?:the\s+)?correct",
                re.IGNORECASE,
            ),
            re.compile(
                rf"(?:i\s+)?(?:would\s+)?choose\s+\(?([{letter_class}])\)?",
                re.IGNORECASE,
            ),
            re.compile(
                rf"\b([{letter_class}])\s+is\s+(?:the\s+)?(?:best|correct|right)\s+"
                rf"(?:answer|choice|option)",
                re.IGNORECASE,
            ),
            re.compile(rf"^([{letter_class}])[\.\)]\s+\S", re.IGNORECASE),
            re.compile(rf"^[ \t]*([{letter_class}])[ \t]*$", re.MULTILINE),
            re.compile(rf"\(([{letter_class}])\)\s*\.?\s*$"),
        )

    def _strip_markdown(self, text: str) -> str:
        return self._MARKDOWN_PATTERN.sub("", text)

    def extract(self, system_output: str, **kwargs: object) -> str:
        text = self._strip_markdown(system_output.strip())
        if not text:
            return ""
        for pattern in self._patterns:
            # Use the LAST match per pattern so that model self-corrections
            # ("the answer is B... wait, the answer is A") resolve to A.
            last_match: re.Match[str] | None = None
            for m in pattern.finditer(text):
                last_match = m
            if last_match is not None:
                letter = last_match.group(1).upper()
                if letter in self.valid_letters:
                    return letter
        return ""


# ---------------------------------------------------------------------------
# Comprehensive MC answer extraction
# ---------------------------------------------------------------------------

# Answer keyword patterns for extract_mc_answer.  Searched via finditer
# with latest-position-wins semantics so that the model's *final*
# declared answer beats letters mentioned during reasoning.
_MC_ANSWER_PATTERNS: tuple[str, ...] = (
    # Bold answer: **Answer: A**, **Final Answer: B**
    r"\*\*(?:[Ff]inal\s+)?[Aa]nswer\s*[:=]\s*\(?([A-Ja-j])\)?",
    # "Answer: **A" or "Answer: *A" (markdown after colon)
    r"[Aa]nswer\s*[:=]\s*\*+\(?([A-Ja-j])\)?(?![A-Za-z])",
    # "Answer: A" or "Answer = A" (plain text)
    r"[Aa]nswer\s*[:=]\s*\(?([A-Ja-j])\)?(?:[.,\)\s]|$)",
    # "The answer is A" / "answer is A"
    r"[Tt]he\s+answer\s+is\s*[:=]?\s*\*?\*?\(?([A-Ja-j])\)?(?![A-Za-z])",
    r"[Aa]nswer\s+is\s*[:=]?\s*\*?\*?\(?([A-Ja-j])\)?(?![A-Za-z])",
    # "My answer is A" / "final answer is A"
    r"[Mm]y\s+(?:final\s+)?answer\s+is\s*[:=]?\s*\(?([A-Ja-j])\)?(?![A-Za-z])",
    r"[Ff]inal\s+[Aa]nswer\s*[:=]?\s*(?:is\s+)?\(?([A-Ja-j])\)?(?![A-Za-z])",
    # "Correct answer: A" / "The correct answer is A"
    r"[Cc]orrect\s+(?:answer|choice|option)\s*[:=]?\s*(?:is\s+)?"
    r"\(?([A-Ja-j])\)?(?![A-Za-z])",
    # "I choose A", "select B", "pick C"
    r"(?:[Ii]\s+)?(?:choose|select|pick)\s+(?:option\s+)?"
    r"\(?([A-Ja-j])\)?(?![A-Za-z])",
    # "Option A is correct", "Choice B"
    r"[Oo]ption\s+\*{0,2}\(?([A-Ja-j])\)?(?:\*{0,2})(?:\s+is|\s*[.,*]|$)",
    r"[Cc]hoice\s+\*{0,2}\(?([A-Ja-j])\)?(?:\*{0,2})(?:\s+is|\s*[.,*]|$)",
    # "Conclusion: E" (Gemini pattern)
    r"[Cc]onclusion\s*[:=]\s*\(?([A-Ja-j])\)?(?![A-Za-z])",
    # "Correct Choice:\nC. 12.5"
    r"[Cc]orrect\s+[Cc]hoice\s*:\s*\n\s*\(?([A-Ja-j])\)?[.\),:=\s]",
)

# CJK answer patterns for AGIEval
_MC_CJK_PATTERNS: tuple[str, ...] = (
    r"答案[是为]\s*[:：]?\s*\(?([A-Ja-j])\)?",
    r"选[择项]\s*[:：]?\s*\(?([A-Ja-j])\)?",
    r"选\s*\(?([A-Ja-j])\)?(?![A-Za-z])",
    r"应[该当]选\s*\(?([A-Ja-j])\)?",
)


def extract_mc_answer(text: str) -> str | None:
    """Extract a multiple-choice answer letter from LLM response text.

    Handles diverse formats: keyword patterns ("Answer: B"), bold
    markdown ("**C.**"), parenthesized letters ("(A)"), CJK patterns
    (AGIEval), LaTeX ``\\boxed{}``, and bare-letter-on-last-line.

    When multiple answer declarations appear (common in chain-of-thought),
    the latest positional match wins so that the model's final answer
    takes precedence over letters mentioned during reasoning.

    Args:
        text: Raw model response text.

    Returns:
        Single uppercase letter (A-J) if found, ``None`` otherwise.
    """
    if not text or not isinstance(text, str):
        return None

    text = text.strip()

    # Bare single letter (ground truth shorthand).
    if len(text) == 1 and text.upper() in "ABCDEFGHIJ":
        return text.upper()

    # Bold-wrapped letter at start: "**C. Full answer text.**"
    bold_start = re.match(
        r"^\*{1,2}\(?([A-HJ-Ja-hj-j])\)?[.\),:\s*]", text
    )
    if bold_start:
        return bold_start.group(1).upper()
    if re.match(r"^\*{1,2}I[.\),:]", text):
        return "I"

    # Letter at start followed by punctuation (not space+word).
    start_match = re.match(r"^([A-HJ-Ja-hj-j])(?:[.\),:]\s*|\s*\n)", text)
    if start_match:
        return start_match.group(1).upper()
    if len(text) > 1 and text[0].upper() == "I" and text[1] in ".):,\n":
        return "I"

    # Strip bold markdown for keyword matching.
    text_plain = text.replace("**", "")

    # Keyword patterns — latest positional match wins.
    best_letter: str | None = None
    best_pos = -1
    for pattern in _MC_ANSWER_PATTERNS:
        for m in re.finditer(pattern, text_plain, re.IGNORECASE):
            if m.start() > best_pos:
                best_pos = m.start()
                best_letter = m.group(1).upper()
    if best_letter is not None:
        return best_letter

    # Bold standalone letter: **A**, **(A)**, **A.**
    bold_match = re.search(
        r"\*\*\(?([A-Ja-j])\)?(?![A-Za-z])[.\):]?\*\*", text
    )
    if bold_match:
        return bold_match.group(1).upper()
    is_bold = re.search(r"is\s+\*\*\(?([A-Ja-j])\)?(?![A-Za-z])", text)
    if is_bold:
        return is_bold.group(1).upper()
    bold_sentence = re.search(r"\*\*([A-Ja-j])[.\)]\s+\w", text)
    if bold_sentence:
        return bold_sentence.group(1).upper()

    # Parenthesized letter with answer context.
    paren_with_context = re.search(
        r"(?:answer|choice|option|select|pick|is)\s+\(([A-Ja-j])\)",
        text, re.I,
    )
    if paren_with_context:
        return paren_with_context.group(1).upper()
    paren_candidates = re.findall(
        r"(?<![A-Za-z0-9])\(([A-Ja-j])\)", text
    )
    if paren_candidates:
        return paren_candidates[-1].upper()

    # LaTeX \boxed{}.
    boxed_mc = re.search(r"\\boxed\{?\(?([A-Ja-j])\)?\}?", text)
    if boxed_mc:
        return boxed_mc.group(1).upper()

    # CJK patterns (AGIEval).
    for pattern in _MC_CJK_PATTERNS:
        match = re.search(pattern, text)
        if match:
            return match.group(1).upper()

    # Letter at end with answer context.
    text_end = text[-150:] if len(text) > 150 else text
    end_with_context = re.search(
        r"(?:answer|therefore|thus|hence|so|consequently)"
        r"\s+(?:is\s+)?(?:\*\*)?([A-Ja-j])(?:\*\*)?[.!?]?\s*$",
        text_end, re.I,
    )
    if end_with_context:
        return end_with_context.group(1).upper()

    # Bare letter on the very last line (last resort).
    last_line = text.rstrip().rsplit("\n", 1)[-1].strip()
    if re.fullmatch(r"\(?([A-HJ-Ja-hj-j])\)?[.!?]?", last_line):
        return last_line.strip("().!? ").upper()

    return None


class MCAnswerExtractor(IOutputExtractor[str, str]):
    """Thin wrapper around :func:`extract_mc_answer` for composition."""

    def extract(self, system_output: str, **kwargs: object) -> str:
        return extract_mc_answer(system_output) or ""


# ---------------------------------------------------------------------------
# Yes/No answer extraction
# ---------------------------------------------------------------------------


def extract_yes_no_answer(text: str) -> str | None:
    """Extract Yes/No/Ambiguous answer from LLM response text.

    Handles five tiers of patterns:
    1. Starts with yes/no/ambiguous
    2. Markdown bold at start (``**Yes**``)
    3. Keyword patterns ("the answer is Yes")
    4. Standalone yes/no at end of response
    5. First word of any sentence is Yes/No

    Args:
        text: Raw model response text.

    Returns:
        ``"Yes"``, ``"No"``, or ``"Ambiguous"`` if found, ``None`` otherwise.
    """
    if not text or not isinstance(text, str):
        return None

    text_lower = text.strip().lower()

    # Tier 1: starts with yes/no/ambiguous.
    if re.match(r"^yes(?:[.,:\s]|$)", text_lower):
        return "Yes"
    if re.match(r"^no(?:[.,:\s]|$)", text_lower):
        return "No"
    if re.match(r"^ambiguous(?:[.,:\s]|$)", text_lower):
        return "Ambiguous"

    # Tier 2: bold at start.
    bold_match = re.match(r"^\*\*(yes|no|ambiguous)\*\*", text_lower)
    if bold_match:
        return bold_match.group(1).capitalize()

    # Tier 3: keyword patterns.
    _yn_patterns = (
        r"(?:the\s+)?answer\s+is\s*[:=]?\s*(?:\*\*)?"
        r"(yes|no|ambiguous)(?:\*\*)?",
        r"(?:my\s+)?(?:final\s+)?answer\s*[:=]\s*(?:\*\*)?"
        r"(yes|no|ambiguous)(?:\*\*)?",
        r"(?:correct|right)\s+answer\s*[:=]?\s*(?:is\s+)?(?:\*\*)?"
        r"(yes|no|ambiguous)(?:\*\*)?",
        r"(?:i\s+)?(?:would\s+)?(?:say|conclude|answer)\s*[:=]?\s*(?:\*\*)?"
        r"(yes|no|ambiguous)(?:\*\*)?",
    )
    for pattern in _yn_patterns:
        match = re.search(pattern, text_lower)
        if match:
            return match.group(1).capitalize()

    # Tier 4: standalone yes/no at end.
    text_end = text_lower[-100:] if len(text_lower) > 100 else text_lower
    end_match = re.search(r"\b(yes|no|ambiguous)[.!?]?\s*$", text_end)
    if end_match:
        return end_match.group(1).capitalize()

    # Tier 5: first word of any sentence.
    sentence_match = re.search(
        r"[.!?]\s*(yes|no|ambiguous)(?:[.,:\s]|$)", text_lower
    )
    if sentence_match:
        return sentence_match.group(1).capitalize()

    return None


class YesNoExtractor(IOutputExtractor[str, str]):
    """Thin wrapper around :func:`extract_yes_no_answer` for composition."""

    def extract(self, system_output: str, **kwargs: object) -> str:
        return extract_yes_no_answer(system_output) or ""
