"""Family 1 — Basic sanity checks for reward functions."""

from __future__ import annotations

import random as _random

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

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

_HIGH_THRESHOLD = 0.5


def _random_tokens(seed: int, n_tokens: int = 6) -> str:
    """Deterministic gibberish string."""
    rng = _random.Random(seed)
    chars = "abcdefghijklmnopqrstuvwxyz0123456789!? "
    return " ".join(
        "".join(rng.choice(chars) for _ in range(rng.randint(2, 5))) for _ in range(n_tokens)
    )


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


class EmptyAttack(BaseAttack):
    name = "empty"
    family = "basic_sanity"
    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):
                score, err = _safe_call(fn, "", task)
                if err is not None:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=Severity.CRITICAL,
                            input="(empty string)",
                            error=err,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' crashed on empty input. "
                                "Ensure the function handles empty completions gracefully."
                            ),
                        )
                    )
                elif score is not None and score > _HIGH_THRESHOLD:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input="(empty string)",
                            expected_reward=0.0,
                            actual_reward=score,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' scores empty string at {score:.3f}. "
                                "An empty completion should receive a near-zero reward."
                            ),
                        )
                    )
        return findings


class WhitespaceAttack(BaseAttack):
    name = "whitespace"
    family = "basic_sanity"
    severity = Severity.CRITICAL

    _INPUT = " \n\t\n "

    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._INPUT, task)
                if err is not None:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=Severity.CRITICAL,
                            input=repr(self._INPUT),
                            error=err,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' crashed on whitespace-only input. "
                                "Strip/normalize whitespace before scoring."
                            ),
                        )
                    )
                elif score is not None and score > _HIGH_THRESHOLD:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=repr(self._INPUT),
                            expected_reward=0.0,
                            actual_reward=score,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' scores whitespace at {score:.3f}. "
                                "Pure whitespace should receive a near-zero reward."
                            ),
                        )
                    )
        return findings


class RandomTokensAttack(BaseAttack):
    name = "random_tokens"
    family = "basic_sanity"
    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:
            for idx, task in enumerate(tasks):
                gibberish = _random_tokens(seed=idx * 997)
                score, err = _safe_call(fn, gibberish, task)
                if err is not None:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=Severity.CRITICAL,
                            input=gibberish,
                            error=err,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' crashed on random tokens. "
                                "The function must tolerate arbitrary input."
                            ),
                        )
                    )
                elif score is not None and score > _HIGH_THRESHOLD:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=gibberish,
                            expected_reward=0.0,
                            actual_reward=score,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' scores random gibberish at {score:.3f}. "
                                "Random tokens should not receive a high reward."
                            ),
                        )
                    )
        return findings


class RewardRangeAttack(BaseAttack):
    name = "reward_range"
    family = "basic_sanity"
    severity = Severity.WARNING  # Not CRITICAL — many legit fns use other ranges

    _INPUTS = ["", "hello world", _random_tokens(0), "42", "The answer is yes."]

    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):
                for inp in self._INPUTS:
                    score, err = _safe_call(fn, inp, task)
                    if err is not None:
                        continue  # crashes are caught by other attacks
                    if score is not None and (score < 0.0 or score > 1.0):
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=inp,
                                expected_reward=None,
                                actual_reward=score,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' returned {score:.4f} which is "
                                    "outside [0.0, 1.0]. Most RL frameworks expect normalized "
                                    "rewards. If your framework handles unnormalized rewards, "
                                    "this may be intentional."
                                ),
                            )
                        )
        return findings


class DeterminismAttack(BaseAttack):
    name = "determinism"
    family = "basic_sanity"
    severity = Severity.CRITICAL

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe
        test_input = "This is a determinism test input."
        for fn in env.reward_fns:
            for idx, task in enumerate(tasks):
                score1, err1 = _safe_call(fn, test_input, task)
                score2, err2 = _safe_call(fn, test_input, task)
                if err1 is not None or err2 is not None:
                    continue  # crashes caught elsewhere
                if score1 != score2:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=test_input,
                            expected_reward=score1,
                            actual_reward=score2,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' is non-deterministic: "
                                f"same input returned {score1:.4f} then {score2:.4f}. "
                                "Reward functions must be deterministic for stable RL training."
                            ),
                        )
                    )
        return findings


class GroundTruthAttack(BaseAttack):
    name = "ground_truth"
    family = "basic_sanity"
    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:
            # Skip format-only fns — they don't check correctness
            if not is_answer_aware(fn, tasks):
                continue

            # Detect what answer format this fn expects
            formats = _detect_answer_format(fn, tasks)

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

                # Try the best-scoring detected format
                completion, score = _best_answer_completion(fn, task, formats)

                if score < _HIGH_THRESHOLD:
                    findings.append(
                        Finding(
                            attack=self.name,
                            family=self.family,
                            severity=self.severity,
                            input=completion,
                            expected_reward=1.0,
                            actual_reward=score,
                            affected_tasks=[idx],
                            recommendation=(
                                f"Reward fn '{fn.name}' scores the correct answer "
                                f"at only {score:.3f} even after trying "
                                f"{len(formats) + 1} common answer formats. "
                                "Ground-truth answers should score > 0.5 in at "
                                "least one format."
                            ),
                        )
                    )
        return findings


class SignalVarianceAttack(BaseAttack):
    name = "signal_variance"
    family = "basic_sanity"
    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:
            # Detect format so we test correct answer in the right format
            formats = _detect_answer_format(fn, tasks)
            answer_aware = is_answer_aware(fn, tasks)

            flat_count = 0
            tested = 0
            for idx, task in enumerate(tasks):
                answer = task.get("answer")
                if answer is None:
                    continue
                tested += 1
                # Use format-aware correct answer
                correct_completion, _ = _best_answer_completion(fn, task, formats)
                inputs = ["", _random_tokens(idx), correct_completion]
                scores: list[float] = []
                for inp in inputs:
                    s, err = _safe_call(fn, inp, task)
                    if err is not None or s is None:
                        break
                    scores.append(s)
                else:
                    if len(set(scores)) == 1:
                        flat_count += 1

            if tested > 0 and flat_count == tested:
                # If fn is format-only (not answer-aware), flat signal
                # across empty/random/raw-answer is expected — skip.
                # But if all scores are HIGH (>0.5), that's still a problem
                # even for format fns (accepts garbage).
                if not answer_aware:
                    # Check if scores are all low — that's expected for format fns
                    sample_score, _ = _safe_call(fn, "", tasks[0])
                    if sample_score is not None and sample_score < _HIGH_THRESHOLD:
                        continue  # Format fn returning 0 for everything — expected

                findings.append(
                    Finding(
                        attack=self.name,
                        family=self.family,
                        severity=self.severity,
                        input="(empty, random, ground_truth)",
                        recommendation=(
                            f"Reward fn '{fn.name}' returns the same score for empty, "
                            "random, and correct completions across all tested tasks. "
                            "The function provides no training signal."
                        ),
                    )
                )
        return findings


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


class BasicSanityFamily(AttackFamily):
    name = "basic_sanity"
    display_name = "Basic Sanity"

    def __init__(self) -> None:
        self.attacks = [
            EmptyAttack(),
            WhitespaceAttack(),
            RandomTokensAttack(),
            RewardRangeAttack(),
            DeterminismAttack(),
            GroundTruthAttack(),
            SignalVarianceAttack(),
        ]
