"""Family 4 — Composition attacks that test multi-function rubric behaviour."""

from __future__ import annotations

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

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

_HIGH_THRESHOLD = 0.5


def _has_multiple_fns(env: NormalizedEnvironment) -> bool:
    return len(env.reward_fns) >= 2


def _skip_result(family: "CompositionFamily") -> FamilyResult:
    """Return an all-passed result when the family is not applicable."""
    return FamilyResult(
        name=family.name,
        display_name=family.display_name,
        total=len(family.attacks),
        passed=len(family.attacks),
        findings=[],
    )


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


class ZeroWeightLeakAttack(BaseAttack):
    name = "zero_weight_leak"
    family = "composition"
    severity = Severity.WARNING

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        if not _has_multiple_fns(env):
            return []
        findings: list[Finding] = []
        tasks = env.tasks  # Already sampled by Probe
        zero_fns = [fn for fn in env.reward_fns if fn.weight == 0.0]
        if not zero_fns:
            return []

        for task in tasks:
            # Two completions that differ on zero-weight fn
            comp_a = ""
            comp_b = "DEFINITELY_WRONG_12345 this should differ"

            try:
                score_a = env.compute_rubric_score(comp_a, task)
                score_b = env.compute_rubric_score(comp_b, task)
            except Exception:
                continue

            # Check if the zero-weight fn returns different values
            for zfn in zero_fns:
                za, _ = _safe_call(zfn, comp_a, task)
                zb, _ = _safe_call(zfn, comp_b, task)
                if za is None or zb is None:
                    continue
                if za == zb:
                    continue  # zero-weight fn gives same score, no leak to detect

                # The zero-weight fn differs — but rubric should NOT differ
                # because of it. However, the difference could come from
                # non-zero-weight fns too. We check if the rubric difference
                # exceeds what non-zero fns alone explain.
                if abs(score_a - score_b) > 0.001:
                    # Recompute without zero-weight fns
                    non_zero = [f for f in env.reward_fns if f.weight != 0.0]
                    tw = sum(f.weight for f in non_zero)
                    if tw == 0:
                        continue
                    nz_a = (
                        sum(
                            (s or 0.0) * f.weight
                            for f in non_zero
                            for s in [_safe_call(f, comp_a, task)[0]]
                        )
                        / tw
                    )
                    nz_b = (
                        sum(
                            (s or 0.0) * f.weight
                            for f in non_zero
                            for s in [_safe_call(f, comp_b, task)[0]]
                        )
                        / tw
                    )

                    expected_diff = abs(nz_a - nz_b)
                    actual_diff = abs(score_a - score_b)
                    if actual_diff > expected_diff + 0.001:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=f"comp_a='' vs comp_b='{comp_b[:50]}'",
                                expected_reward=expected_diff,
                                actual_reward=actual_diff,
                                recommendation=(
                                    f"Zero-weight fn '{zfn.name}' leaks into the "
                                    "rubric score. A fn with weight=0 should have "
                                    "no effect on the final score."
                                ),
                            )
                        )
        return findings


class WrongSignAttack(BaseAttack):
    name = "wrong_sign"
    family = "composition"
    severity = Severity.CRITICAL

    _INPUTS = [
        "",
        "hello",
        "42",
        "The answer is yes.",
        "DEFINITELY_WRONG_12345",
    ]

    def run(self, env: NormalizedEnvironment) -> list[Finding]:
        if not _has_multiple_fns(env):
            return []
        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
                    if score is not None and score < 0.0:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=inp,
                                actual_reward=score,
                                affected_tasks=[idx],
                                recommendation=(
                                    f"Reward fn '{fn.name}' returned negative value "
                                    f"{score:.4f}. Negative rewards can flip weighted "
                                    "composition and produce degenerate training signals."
                                ),
                            )
                        )
                        break  # one finding per fn is enough
                else:
                    continue
                break
        return findings


class NoneHandlingAttack(BaseAttack):
    name = "none_handling"
    family = "composition"
    severity = Severity.CRITICAL

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

        for fn_idx, fn in enumerate(env.reward_fns):
            # Create a wrapper that returns None
            original_fn = fn.fn

            def none_fn(completion: str, task: dict, _orig=original_fn) -> float:
                return None  # type: ignore[return-value]

            # Temporarily swap
            fn.fn = none_fn  # type: ignore[assignment]
            try:
                for task in tasks:
                    try:
                        env.compute_rubric_score("test", task)
                    except Exception as exc:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input="(reward fn returns None)",
                                error=f"{type(exc).__name__}: {exc}",
                                recommendation=(
                                    f"Rubric crashes when '{fn.name}' returns None. "
                                    "Add a None-check or type coercion in the rubric's "
                                    "score aggregation."
                                ),
                            )
                        )
                        break
            finally:
                fn.fn = original_fn  # type: ignore[assignment]

        return findings


class WeightDominanceAttack(BaseAttack):
    name = "weight_dominance"
    family = "composition"
    severity = Severity.WARNING

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

        sorted_fns = sorted(env.reward_fns, key=lambda f: f.weight)
        lowest = sorted_fns[0]
        highest = sorted_fns[-1]
        if lowest.name == highest.name:
            return []

        # Try to find a completion that scores 1.0 on lowest and 0.0 on highest
        # We test the ground-truth answer (likely scores high on correctness fn)
        # and also empty/random inputs.
        candidates = [
            "",
            "DEFINITELY_WRONG_12345",
            "42",
            "yes",
            "no",
        ]
        for task in tasks:
            answer = task.get("answer")
            if answer is not None:
                candidates.append(str(answer))

        for task in tasks:
            for comp in candidates:
                low_score, low_err = _safe_call(lowest, comp, task)
                high_score, high_err = _safe_call(highest, comp, task)
                if low_err or high_err or low_score is None or high_score is None:
                    continue
                if low_score > 0.9 and high_score < 0.1:
                    try:
                        rubric = env.compute_rubric_score(comp, task)
                    except Exception:
                        continue
                    if rubric > _HIGH_THRESHOLD:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input=comp[:100],
                                actual_reward=rubric,
                                recommendation=(
                                    f"Lowest-weight fn '{lowest.name}' (w={lowest.weight}) "
                                    f"dominates the rubric: scores 1.0 on it, 0.0 on "
                                    f"highest-weight '{highest.name}' (w={highest.weight}), "
                                    f"yet rubric = {rubric:.3f}. Weight distribution "
                                    "may need rebalancing."
                                ),
                            )
                        )
                        return findings  # one example suffices
        return findings


class SingleFnFailureAttack(BaseAttack):
    name = "single_fn_failure"
    family = "composition"
    severity = Severity.CRITICAL

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

        for fn_idx, fn in enumerate(env.reward_fns):
            original_fn = fn.fn

            def exploding_fn(completion: str, task: dict) -> float:
                raise RuntimeError("Simulated reward fn failure")

            fn.fn = exploding_fn  # type: ignore[assignment]
            try:
                for task in tasks:
                    try:
                        env.compute_rubric_score("test completion", task)
                        # If it doesn't crash, that's fine — rubric is resilient
                    except Exception as exc:
                        findings.append(
                            Finding(
                                attack=self.name,
                                family=self.family,
                                severity=self.severity,
                                input="(one fn raises RuntimeError)",
                                error=f"{type(exc).__name__}: {exc}",
                                recommendation=(
                                    f"The entire rubric crashes when '{fn.name}' "
                                    "raises an exception. Consider wrapping individual "
                                    "reward fn calls in try/except so other fns can "
                                    "still contribute."
                                ),
                            )
                        )
                        break
            finally:
                fn.fn = original_fn  # type: ignore[assignment]

        return findings


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


class CompositionFamily(AttackFamily):
    name = "composition"
    display_name = "Composition"

    def __init__(self) -> None:
        self.attacks = [
            ZeroWeightLeakAttack(),
            WrongSignAttack(),
            WeightDominanceAttack(),
        ]

    def run(self, env: NormalizedEnvironment) -> FamilyResult:
        if len(env.reward_fns) < 2:
            return _skip_result(self)
        return super().run(env)
