"""Family 5 — Schema attacks that test input/output contract robustness."""

from __future__ import annotations

from typing import Any

from rewardprobe.families.base import AttackFamily, BaseAttack, _safe_call
from rewardprobe.types import Finding, NormalizedEnvironment, Severity

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


class MissingColumnAttack(BaseAttack):
    name = "missing_column"
    family = "schema"
    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):
                keys = list(task.keys())
                for key in keys:
                    partial = {k: v for k, v in task.items() if k != key}
                    score, err = _safe_call(fn, "test", partial)
                    if err is not None:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=f"(task dict missing key '{key}')",
                                error=err,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' crashes when task dict is "
                                    f"missing the '{key}' key. Use .get() with defaults "
                                    "or validate the task schema upfront."
                                ),
                            )
                        )
                # Only test one task per fn to avoid excessive noise
                break
        return findings


class WrongArgumentAttack(BaseAttack):
    name = "wrong_argument"
    family = "schema"
    severity = Severity.WARNING

    _EXTRA_KEYS = {
        "_debug": True,
        "unexpected_field": "surprise",
        "__metadata__": {"nested": "dict"},
        "score_override": 1.0,
    }

    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):
                augmented = {**task, **self._EXTRA_KEYS}
                score, err = _safe_call(fn, "test", augmented)
                if err is not None:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=f"(task dict with extra keys: {list(self._EXTRA_KEYS)})",
                            error=err,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' crashes when task dict has "
                                "unexpected extra keys. The function should ignore "
                                "unknown keys."
                            ),
                        )
                    )
                # One task per fn is sufficient
                break
        return findings


class MessageVsStringAttack(BaseAttack):
    name = "message_vs_string"
    family = "schema"
    severity = Severity.INFO

    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):
                # Pass a list[dict] as completion to check cross-framework compat.
                # This is INFO-level because most fns are designed for one format.
                message_completion: Any = [{"role": "assistant", "content": "test answer"}]
                score, err = _safe_call(fn, message_completion, task)
                if err is not None:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input="(completion passed as list[dict] instead of str)",
                            error=err,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' crashes on chat message format "
                                "(list[dict]) input. This is expected if the function "
                                "is designed for string completions only. Only relevant "
                                "if the fn will be used across frameworks."
                            ),
                        )
                    )
                break
        return findings


class CompletionStructureAttack(BaseAttack):
    name = "completion_structure"
    family = "schema"
    severity = Severity.WARNING

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe
        edge_cases = {
            "single_char": "x",
            "very_long": "A" * 10_000,
            "unicode": "\u2603 \U0001f4a9 \u2603 snowman and pile of poo",
            "null_bytes": "answer\x00with\x00nulls",
            "extra_whitespace": "   \n\n  answer  \n\t  ",
        }
        for fn in env.reward_fns:
            for idx, task in enumerate(tasks):
                for label, completion in edge_cases.items():
                    score, err = _safe_call(fn, completion, task)
                    if err is not None:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=f"({label}: {completion[:60]!r})",
                                error=err,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' crashes on {label} input. "
                                    "Ensure the function handles edge-case string inputs "
                                    "(unicode, long text, null bytes, whitespace)."
                                ),
                            )
                        )
                # One task per fn
                break
        return findings


class BatchOrderingAttack(BaseAttack):
    name = "batch_ordering"
    family = "schema"
    severity = Severity.WARNING

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe
        test_input = "This is a batch ordering test."
        for fn in env.reward_fns:
            for idx, task in enumerate(tasks):
                # Call several times in a row with interleaved other calls
                scores = []
                for _ in range(3):
                    # Interleave with a different call to simulate batch context
                    _safe_call(fn, "unrelated dummy call", task)
                    score, err = _safe_call(fn, 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=test_input,
                                expected_reward=scores[0],
                                actual_reward=scores[-1],
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' returns different scores "
                                    f"({scores}) for the same input when called in "
                                    "different positions. The function may have hidden "
                                    "state or depend on call order."
                                ),
                            )
                        )
        return findings


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


class SchemaFamily(AttackFamily):
    name = "schema"
    display_name = "Schema"

    def __init__(self) -> None:
        self.attacks = [
            MissingColumnAttack(),
            WrongArgumentAttack(),
            MessageVsStringAttack(),
            CompletionStructureAttack(),
            BatchOrderingAttack(),
        ]
