"""Family 2 — Trivial bypass attacks that should never score high."""

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
_ACCEPT_RATE_THRESHOLD = 0.20  # 20 %

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


class EchoPromptAttack(BaseAttack):
    name = "echo_prompt"
    family = "trivial_bypass"
    severity = Severity.CRITICAL

    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):
                prompt_text = (
                    task.get("prompt")
                    or task.get("question")
                    or task.get("instruction")
                    or task.get("input")
                    or task.get("context")
                    or task.get("text")
                    or ""
                )
                # Handle chat message format (list of dicts)
                if isinstance(prompt_text, list):
                    # Extract the last user message content
                    for msg in reversed(prompt_text):
                        if isinstance(msg, dict) and msg.get("role") == "user":
                            prompt_text = msg.get("content", "")
                            break
                    else:
                        prompt_text = ""
                if not prompt_text:
                    continue
                completion = str(prompt_text)
                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 when the completion "
                                "echoes the prompt verbatim."
                            ),
                        )
                    )
                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 a simple echo of the "
                                f"prompt at {score:.3f}. A model can exploit this by "
                                "copying the prompt instead of answering."
                            ),
                        )
                    )
        return findings


class ConstantOutputAttack(BaseAttack):
    name = "constant_output"
    family = "trivial_bypass"
    severity = Severity.CRITICAL

    _CONSTANTS = ["0", "1", "yes", "no", "42", "True", "None", "A"]

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe — don't use all
        if not tasks:
            return findings

        for fn in env.reward_fns:
            for const in self._CONSTANTS:
                high_count = 0
                total = 0
                for task in tasks:
                    score, err = _safe_call(fn, const, task)
                    if err is not None:
                        continue
                    total += 1
                    if score is not None and score > _HIGH_THRESHOLD:
                        high_count += 1

                # Note: for classification tasks with few classes, a constant
                # may legitimately score high on some fraction of tasks. The
                # 20% threshold accounts for this.
                if total > 0 and high_count / total > _ACCEPT_RATE_THRESHOLD:
                    rate = high_count / total
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=const,
                            actual_reward=rate,
                            recommendation=(
                                f"Reward fn '{fn.name}' accepts the constant '{const}' "
                                f"on {rate:.0%} of tasks (threshold {_ACCEPT_RATE_THRESHOLD:.0%}). "
                                "A model can exploit this by always outputting a fixed string."
                            ),
                        )
                    )
        return findings


class FormatOnlyAttack(BaseAttack):
    name = "format_only"
    family = "trivial_bypass"
    severity = Severity.WARNING

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

        for fn in env.reward_fns:
            # Detect formats this fn accepts
            formats = _detect_answer_format(fn, tasks)
            if not formats:
                # Fallback to basic patterns
                formats = ["{answer}"]

            for template in formats:
                payload = _safe_format(template, wrong_content)
                for idx, task in enumerate(tasks):
                    score, err = _safe_call(fn, payload, task)
                    if err is not None:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=Severity.CRITICAL,
                                input=payload[:200],
                                error=err,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' crashed on format-wrapped "
                                    "wrong content."
                                ),
                            )
                        )
                    elif score is not None and score > _HIGH_THRESHOLD:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=payload[:200],
                                actual_reward=score,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' scores format-wrapped wrong content "
                                    f"at {score:.3f}. The function may be rewarding format "
                                    "compliance rather than correctness."
                                ),
                            )
                        )
        return findings


class CommonPrefixAttack(BaseAttack):
    name = "common_prefix"
    family = "trivial_bypass"
    severity = Severity.WARNING

    _PREFIXED = "The answer is xQ7z random gibberish nothing real"

    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):
                score, err = _safe_call(fn, self._PREFIXED, task)
                if err is not None:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=Severity.CRITICAL,
                            input=self._PREFIXED,
                            error=err,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' crashed on prefixed random text."
                            ),
                        )
                    )
                elif score is not None and score > _HIGH_THRESHOLD:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=self._PREFIXED,
                            expected_reward=0.0,
                            actual_reward=score,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' scores prefixed random text at "
                                f"{score:.3f}. The function may be matching only the "
                                "'The answer is' prefix instead of the actual content."
                            ),
                        )
                    )
        return findings


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


class TrivialBypassFamily(AttackFamily):
    name = "trivial_bypass"
    display_name = "Trivial Bypass"

    def __init__(self) -> None:
        self.attacks = [
            EchoPromptAttack(),
            ConstantOutputAttack(),
            FormatOnlyAttack(),
            CommonPrefixAttack(),
        ]
