"""Family 3 — Parser bug attacks that probe answer-extraction logic."""

from __future__ import annotations

from rewardprobe.families.base import (
    AttackFamily,
    BaseAttack,
    _detect_answer_format,
    _safe_call,
    _safe_format,
)
from rewardprobe.types import Finding, NormalizedEnvironment, Severity

_HIGH_THRESHOLD = 0.5

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _is_tag_format(template: str) -> bool:
    """Check if a format template uses XML-style tags (e.g. <answer>{answer}</answer>)."""
    return "<" in template and ">" in template and "{answer}" in template


def _extract_tag_name(template: str) -> str | None:
    """Extract the tag name from a template like '<answer>{answer}</answer>'."""
    import re

    m = re.search(r"<(\w+)>", template)
    return m.group(1) if m else None


# ---------------------------------------------------------------------------
# Attacks
# ---------------------------------------------------------------------------


class AnswerInReasoningAttack(BaseAttack):
    name = "answer_in_reasoning"
    family = "parser_bugs"
    severity = Severity.CRITICAL

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe
        wrong = "DEFINITELY_WRONG_12345"

        for fn in env.reward_fns:
            # Detect what format the fn accepts
            formats = _detect_answer_format(fn, tasks)
            if not formats:
                formats = ["{answer}"]

            for idx, task in enumerate(tasks):
                answer = task.get("answer")
                if answer is None:
                    continue

                # Generate reasoning payloads: correct answer in reasoning,
                # wrong answer in the detected format
                for template in formats:
                    wrong_payload = _safe_format(template, wrong)
                    completion = f"Let me think... {answer}. {wrong_payload}"
                    score, err = _safe_call(fn, completion, task)
                    if err is not None:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=Severity.CRITICAL,
                                input=completion[:200],
                                error=err,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' crashed on answer-in-reasoning input."
                                ),
                            )
                        )
                    elif score is not None and score > _HIGH_THRESHOLD:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=completion[:200],
                                expected_reward=0.0,
                                actual_reward=score,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' scores {score:.3f} when the "
                                    "correct answer appears in reasoning but a wrong "
                                    "answer is in the answer format. The parser may be "
                                    "matching outside the designated answer region."
                                ),
                            )
                        )
        return findings


class MultipleTagsAttack(BaseAttack):
    name = "multiple_tags"
    family = "parser_bugs"
    severity = Severity.WARNING

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe
        wrong = "DEFINITELY_WRONG_12345"

        for fn in env.reward_fns:
            # Detect what format the fn accepts
            formats = _detect_answer_format(fn, tasks)
            if not formats:
                formats = ["{answer}"]

            for idx, task in enumerate(tasks):
                answer = task.get("answer")
                if answer is None:
                    continue

                for template in formats:
                    wrong_fmt = _safe_format(template, wrong)
                    correct_fmt = _safe_format(template, str(answer))

                    # wrong first, correct second
                    comp_a = f"{wrong_fmt} some text {correct_fmt}"
                    # correct first, wrong second
                    comp_b = f"{correct_fmt} some text {wrong_fmt}"

                    score_a, err_a = _safe_call(fn, comp_a, task)
                    score_b, err_b = _safe_call(fn, comp_b, task)

                    if err_a is not None or err_b is not None:
                        err = err_a or err_b
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=Severity.CRITICAL,
                                input=comp_a[:200],
                                error=err,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' crashed on "
                                    "multiple-format-marker input."
                                ),
                            )
                        )
                        continue

                    if score_a is not None and score_b is not None:
                        diff = abs(score_a - score_b)
                        if diff > 0.1:
                            findings.append(
                                Finding(
                                    attack=self.name,
                                    family=self.family,
                                    severity=self.severity,
                                    input=f"A: {comp_a[:100]} | B: {comp_b[:100]}",
                                    expected_reward=score_a,
                                    actual_reward=score_b,
                                    affected_tasks=[idx],
                                    recommendation=(
                                        f"Reward fn '{fn.name}' gives different scores "
                                        f"({score_a:.3f} vs {score_b:.3f}) depending on "
                                        "the order of duplicate answer format markers. "
                                        "The parser's selection strategy (first vs last) "
                                        "is inconsistent."
                                    ),
                                )
                            )
        return findings


class RegexGredinessAttack(BaseAttack):
    name = "regex_greediness"
    family = "parser_bugs"
    severity = Severity.WARNING

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe
        wrong = "DEFINITELY_WRONG_12345"

        for fn in env.reward_fns:
            # Detect formats — only tag-based formats are relevant for regex
            # greediness testing
            formats = _detect_answer_format(fn, tasks)
            tag_formats = [t for t in formats if _is_tag_format(t)]
            if not tag_formats:
                continue

            for idx, task in enumerate(tasks):
                answer = task.get("answer")
                if answer is None:
                    continue

                for template in tag_formats:
                    wrong_fmt = _safe_format(template, wrong)
                    correct_fmt = _safe_format(template, str(answer))

                    # Greedy would match wrong...correct as one block
                    comp_a = f"{wrong_fmt} text {correct_fmt}"
                    comp_b = f"{correct_fmt} text {wrong_fmt}"
                    score_a, err_a = _safe_call(fn, comp_a, task)
                    score_b, err_b = _safe_call(fn, comp_b, task)

                    if err_a is not None or err_b is not None:
                        err = err_a or err_b
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=Severity.CRITICAL,
                                input=comp_a[:200],
                                error=err,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' crashed on duplicate answer tags."
                                ),
                            )
                        )
                        continue

                    if score_a is not None and score_b is not None:
                        diff = abs(score_a - score_b)
                        if diff > 0.1:
                            findings.append(
                                Finding(
                                    attack=self.name,
                                    family=self.family,
                                    severity=self.severity,
                                    input=f"A: {comp_a[:100]} | B: {comp_b[:100]}",
                                    expected_reward=score_a,
                                    actual_reward=score_b,
                                    affected_tasks=[idx],
                                    recommendation=(
                                        f"Reward fn '{fn.name}' scores differ "
                                        f"({score_a:.3f} vs {score_b:.3f}) based on "
                                        "answer tag order. This suggests a greedy vs "
                                        "lazy regex mismatch in the parser."
                                    ),
                                )
                            )
        return findings


class MalformedXmlAttack(BaseAttack):
    name = "malformed_xml"
    family = "parser_bugs"
    severity = Severity.WARNING  # WARNING for acceptance, CRITICAL for crash

    # Generic malformation patterns (always tested)
    _GENERIC_PATTERNS = [
        ("truncated", lambda ans: ans[: len(ans) // 2] if len(ans) > 1 else ans),
        ("extra_whitespace_around", lambda ans: f"  {ans}  "),
    ]

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe

        for fn in env.reward_fns:
            formats = _detect_answer_format(fn, tasks)
            patterns: list[str] = []

            # Build tag-specific malformation patterns
            for template in formats:
                tag_name = _extract_tag_name(template) if _is_tag_format(template) else None
                if tag_name:
                    patterns.extend(
                        [
                            f"<{tag_name}>42",  # unclosed tag
                            f"<{tag_name}>42</{tag_name[::-1]}>",  # reversed closing tag
                            f"&lt;{tag_name}&gt;42",  # HTML entities
                            f"<{tag_name}><{tag_name}>42</{tag_name}></{tag_name}>",  # nested
                            f"<{tag_name.upper()}>42</{tag_name.upper()}>",  # case variation
                        ]
                    )

            # Generic malformations using detected format
            for template in formats:
                if template != "{answer}":
                    patterns.append(_safe_format(template, "42")[:-1])  # truncated format
                    patterns.append(_safe_format(template, "  42  "))  # extra whitespace in format

            # Fallback if no formats detected
            if not patterns:
                patterns = [
                    "42",  # just a bare answer, truncated
                    "  42  ",  # extra whitespace
                ]

            for pattern in patterns:
                for idx, task in enumerate(tasks):
                    score, err = _safe_call(fn, pattern, task)
                    if err is not None:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=Severity.CRITICAL,
                                input=pattern,
                                error=err,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' crashed on malformed "
                                    f"input '{pattern[:80]}'. The parser must handle "
                                    "malformed markup without crashing."
                                ),
                            )
                        )
                    elif score is not None and score > _HIGH_THRESHOLD:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=Severity.WARNING,
                                input=pattern,
                                expected_reward=0.0,
                                actual_reward=score,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' accepts malformed "
                                    f"input '{pattern[:80]}' with score {score:.3f}. "
                                    "This may allow the model to bypass the expected "
                                    "answer format."
                                ),
                            )
                        )
        return findings


class DuplicateAnswersAttack(BaseAttack):
    name = "duplicate_answers"
    family = "parser_bugs"
    severity = Severity.WARNING

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe

        for fn in env.reward_fns:
            formats = _detect_answer_format(fn, tasks)

            for idx, task in enumerate(tasks):
                answer = task.get("answer")
                if answer is None:
                    continue

                # Same answer in all detected formats at once
                parts = []
                for template in formats:
                    parts.append(_safe_format(template, str(answer)))
                if not parts:
                    parts = [str(answer)]
                completion = " ".join(parts)

                # Also test just the raw answer for comparison
                score_multi, err = _safe_call(fn, completion, task)
                if err is not None:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=Severity.CRITICAL,
                            input=completion[:200],
                            error=err,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' crashed when the answer appears "
                                "in multiple format markers simultaneously."
                            ),
                        )
                    )
                    continue

                score_raw, err2 = _safe_call(fn, str(answer), task)
                if err2 is not None or score_raw is None or score_multi is None:
                    continue

                diff = abs(score_multi - score_raw)
                if diff > 0.3:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=completion[:200],
                            expected_reward=score_raw,
                            actual_reward=score_multi,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' scores differ significantly "
                                f"({score_raw:.3f} vs {score_multi:.3f}) when the same "
                                "correct answer is provided in multiple format markers. "
                                "The parser may be confused by duplicate answer markers."
                            ),
                        )
                    )
        return findings


class FormatVariationAttack(BaseAttack):
    name = "format_variation"
    family = "parser_bugs"
    severity = Severity.WARNING

    def _variants(self, answer: str) -> list[str]:
        """Generate format variations of the answer."""
        variants = [
            f" {answer} ",  # extra whitespace
            f"  {answer}  ",  # more whitespace
            f"{answer} ",  # trailing whitespace
        ]
        # Numeric variations
        try:
            num = float(answer)
            variants.append(f"{num:.1f}")  # 42 -> 42.0
            variants.append(f"{num:.2f}")  # 42 -> 42.00
            if num == int(num) and num >= 0:
                variants.append(f"0{int(num)}")  # 42 -> 042
        except (ValueError, TypeError):
            pass
        # String/text variations
        variants.append(answer.lower())
        variants.append(answer.upper())
        variants.append(answer.strip())
        # De-duplicate while preserving order
        seen: set[str] = set()
        unique: list[str] = []
        for v in variants:
            if v not in seen and v != answer:
                seen.add(v)
                unique.append(v)
        return unique

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe
        for fn in env.reward_fns:
            for idx, task in enumerate(tasks):
                answer = task.get("answer")
                if answer is None:
                    continue
                answer_str = str(answer)
                base_score, base_err = _safe_call(fn, answer_str, task)
                if base_err is not None or base_score is None:
                    continue

                for variant in self._variants(answer_str):
                    v_score, v_err = _safe_call(fn, variant, task)
                    if v_err is not None:
                        continue
                    if v_score is None:
                        continue
                    # If the base scores high but a variant scores low, flag it
                    if base_score > _HIGH_THRESHOLD and v_score < 0.1:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=f"base='{answer_str}' variant='{variant}'",
                                expected_reward=base_score,
                                actual_reward=v_score,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' scores '{answer_str}' at "
                                    f"{base_score:.3f} but rejects variant '{variant}' "
                                    f"({v_score:.3f}). Consider normalizing answer "
                                    "formats (whitespace, numeric equivalence, casing)."
                                ),
                            )
                        )
        return findings


# ---------------------------------------------------------------------------
# Family
# ---------------------------------------------------------------------------


class ParserBugsFamily(AttackFamily):
    name = "parser_bugs"
    display_name = "Parser Bugs"

    def __init__(self) -> None:
        self.attacks = [
            AnswerInReasoningAttack(),
            MultipleTagsAttack(),
            RegexGredinessAttack(),
            MalformedXmlAttack(),
            DuplicateAnswersAttack(),
            FormatVariationAttack(),
        ]
