"""Family 6 — Stateful environment bugs for multi-turn RL environments."""

from __future__ import annotations

from rewardprobe.adapters.auto import is_answer_aware
from rewardprobe.families.base import (
    AttackFamily,
    BaseAttack,
    _best_answer_completion,
    _detect_answer_format,
    _safe_call,
)
from rewardprobe.types import Finding, NormalizedEnvironment, Severity

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


class StateLeakAttack(BaseAttack):
    """Answer visible in environment observations/state."""

    name = "state_leak"
    family = "stateful"
    severity = Severity.CRITICAL

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        for idx, task in enumerate(env.tasks):
            answer = task.get("answer")
            if not answer:
                continue
            answer_str = str(answer)
            # Only flag non-trivial answers to avoid false positives
            if len(answer_str) <= 2:
                continue
            for key, value in task.items():
                if key == "answer":
                    continue
                value_str = str(value)
                if answer_str in value_str:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=f"answer '{answer_str}' found in task['{key}']",
                            affected_tasks=[idx],
                            recommendation=(
                                f"The ground-truth answer '{answer_str}' appears "
                                f"in the task's '{key}' field. A model can extract "
                                "the answer from the input without solving the task."
                            ),
                        )
                    )
                    break  # One finding per task
        return findings


class ToolSkipAttack(BaseAttack):
    """Model gets reward without calling required tools."""

    name = "tool_skip"
    family = "stateful"
    severity = Severity.WARNING

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        # Only relevant for environments that have tools
        has_tools = env.metadata.get("has_tools", False)
        if not has_tools:
            return findings

        for fn in env.reward_fns:
            if not is_answer_aware(fn, env.tasks):
                continue
            formats = _detect_answer_format(fn, env.tasks)
            for idx, task in enumerate(env.tasks):
                answer = task.get("answer")
                if not answer:
                    continue
                # Generate a completion with the right answer but NO tool calls
                completion, score = _best_answer_completion(fn, task, formats)
                if score > 0.5:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=completion[:200],
                            actual_reward=score,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' scores {score:.3f} for a "
                                "plain-text answer with no tool calls in a "
                                "tool-use environment. The model can skip tool "
                                "use entirely and still get rewarded."
                            ),
                        )
                    )
                    break  # One example per reward fn suffices
        return findings


class StaleStateAttack(BaseAttack):
    """Reward function depends on call order / previous results."""

    name = "stale_state"
    family = "stateful"
    severity = Severity.CRITICAL

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        if len(env.tasks) < 2:
            return findings

        for fn in env.reward_fns:
            formats = _detect_answer_format(fn, env.tasks)
            task_a = env.tasks[0]
            task_b = env.tasks[1]

            comp_a, _ = _best_answer_completion(fn, task_a, formats)

            # Call on task A
            score_a1, err1 = _safe_call(fn, comp_a, task_a)
            # Call on task B (potential state contamination)
            _safe_call(fn, "different completion", task_b)
            # Call on task A again
            score_a2, err2 = _safe_call(fn, comp_a, task_a)

            if err1 is None and err2 is None and score_a1 != score_a2:
                findings.append(
                    Finding(
                        attack=self.name,
                        family=self.family,
                        severity=self.severity,
                        input=comp_a[:200],
                        expected_reward=score_a1,
                        actual_reward=score_a2,
                        recommendation=(
                            f"Reward fn '{fn.name}' returns different scores "
                            f"({score_a1:.3f} vs {score_a2:.3f}) for the same "
                            "input after processing a different task. The function "
                            "may have hidden state from previous evaluations."
                        ),
                    )
                )
        return findings


class EarlyTerminationAttack(BaseAttack):
    """Reward is positive even when episode is trivially short."""

    name = "early_termination"
    family = "stateful"
    severity = Severity.WARNING

    _MULTITURN_TYPES = {
        "multiturn",
        "multi_turn",
        "tool",
        "stateful_tool",
        "toolenv",
        "multiturnenv",
        "statefultoolenv",
    }

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        # Only meaningful for multi-turn environments
        env_type = env.metadata.get("env_type", "").lower()
        if env_type not in self._MULTITURN_TYPES:
            return findings

        for fn in env.reward_fns:
            if not is_answer_aware(fn, env.tasks):
                continue
            formats = _detect_answer_format(fn, env.tasks)
            for idx, task in enumerate(env.tasks):
                answer = task.get("answer")
                if not answer:
                    continue
                # Minimal completion — just the answer, no multi-turn interaction
                completion, score = _best_answer_completion(fn, task, formats)
                if score > 0.5:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=completion[:200],
                            actual_reward=score,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' scores {score:.3f} for a "
                                "minimal one-turn answer in a multi-turn environment. "
                                "The model can skip the multi-turn interaction entirely "
                                "and just output the answer directly."
                            ),
                        )
                    )
                    break  # One example suffices
        return findings


class ToolSpamAttack(BaseAttack):
    """Repeated identical inputs cause different scores."""

    name = "tool_spam"
    family = "stateful"
    severity = Severity.WARNING

    _TEST_INPUT = "I'll try again: the answer is 42"

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        for fn in env.reward_fns:
            task = env.tasks[0] if env.tasks else {}
            scores: list[float] = []
            for _ in range(5):
                score, err = _safe_call(fn, self._TEST_INPUT, task)
                if err is not None:
                    break
                scores.append(score)
            else:
                if len(set(scores)) > 1:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=self._TEST_INPUT,
                            recommendation=(
                                f"Reward fn '{fn.name}' returns varying scores "
                                f"({scores}) for repeated identical calls. "
                                "The function may accumulate state between calls."
                            ),
                        )
                    )
        return findings


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


class StatefulFamily(AttackFamily):
    name = "stateful"
    display_name = "Stateful"

    def __init__(self) -> None:
        self.attacks = [
            StateLeakAttack(),
            StaleStateAttack(),
            ToolSpamAttack(),
            EarlyTerminationAttack(),
            ToolSkipAttack(),
        ]
