"""Base classes for attack families."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Optional

from rewardprobe.types import (
    FamilyResult,
    Finding,
    NormalizedEnvironment,
    NormalizedRewardFn,
    Severity,
)


def _safe_call(
    fn: NormalizedRewardFn,
    completion: str,
    task: dict[str, Any],
) -> tuple[Optional[float], Optional[str]]:
    """Call a reward function safely; return (score, error_or_None)."""
    try:
        score = fn.fn(completion, task)
        return score, None
    except Exception as exc:
        return None, f"{type(exc).__name__}: {exc}"


# ---------------------------------------------------------------------------
# Format-aware answer wrapping
# ---------------------------------------------------------------------------

# Common answer formats used in RL training environments
_ANSWER_FORMATS = [
    ("raw", "{answer}"),
    ("boxed", "\\boxed{{{answer}}}"),
    ("hash", "#### {answer}"),
    ("xml_answer", "<answer>{answer}</answer>"),
    ("think_boxed", "<think>\nLet me work through this.\n</think>\n\\boxed{{{answer}}}"),
    ("think_xml", "<think>\nLet me work through this.\n</think>\n<answer>{answer}</answer>"),
    ("sentence", "The answer is {answer}."),
    ("final_answer", "\n#### The final answer is {answer}"),
]


def _detect_answer_format(
    fn: NormalizedRewardFn,
    tasks: list[dict[str, Any]],
    n_probe: int = 3,
) -> list[str]:
    """Probe a reward function to discover which answer formats it accepts.

    Returns a list of format template strings that scored > 0.5 on at least
    one task, ordered by highest average score.
    """
    scores_by_fmt: dict[str, list[float]] = {name: [] for name, _ in _ANSWER_FORMATS}

    for task in tasks[:n_probe]:
        answer = task.get("answer")
        if answer is None:
            continue
        for fmt_name, template in _ANSWER_FORMATS:
            completion = _safe_format(template, str(answer))
            score, err = _safe_call(fn, completion, task)
            if err is None and score is not None:
                scores_by_fmt[fmt_name].append(score)

    # Return formats that scored > 0.5 on average, best first
    accepted = []
    for fmt_name, template in _ANSWER_FORMATS:
        scores = scores_by_fmt[fmt_name]
        if scores and sum(scores) / len(scores) > 0.5:
            accepted.append(template)

    return accepted


def _safe_format(template: str, answer: str) -> str:
    """Format template with answer, safely handling braces in the answer."""
    safe_answer = str(answer).replace("{", "{{").replace("}", "}}")
    return template.format(answer=safe_answer)


def _wrap_answer(answer: str, formats: list[str]) -> list[str]:
    """Wrap an answer in all detected formats. Always includes raw."""
    results = []
    for template in formats:
        results.append(_safe_format(template, str(answer)))
    # Always include raw as fallback
    if not any("{answer}" == t for t in formats):
        results.append(str(answer))
    return results


def _best_answer_completion(
    fn: NormalizedRewardFn,
    task: dict[str, Any],
    formats: list[str],
) -> tuple[str, float]:
    """Return the completion + score using the best-scoring format."""
    answer = task.get("answer", "")
    best_completion = str(answer)
    best_score = -1.0

    for template in formats:
        completion = _safe_format(template, str(answer))
        score, err = _safe_call(fn, completion, task)
        if err is None and score is not None and score > best_score:
            best_score = score
            best_completion = completion

    # Fallback to raw
    if best_score < 0:
        score, _ = _safe_call(fn, str(answer), task)
        return str(answer), score if score is not None else 0.0

    return best_completion, best_score


class BaseAttack(ABC):
    """A single attack probe."""

    name: str
    family: str
    severity: Severity

    @abstractmethod
    def run(self, env: NormalizedEnvironment) -> list[Finding]: ...


class AttackFamily(ABC):
    """A group of related attack probes."""

    name: str
    display_name: str
    attacks: list[BaseAttack]

    def run(self, env: NormalizedEnvironment) -> FamilyResult:
        """Run every attack in the family and aggregate results."""
        all_findings: list[Finding] = []
        passed = 0
        for attack in self.attacks:
            findings = attack.run(env)
            if not findings:
                passed += 1
            else:
                all_findings.extend(findings)

        # Deduplicate: merge findings with same (attack, recommendation)
        # into one finding with combined affected_tasks.
        aggregated = _aggregate_findings(all_findings)

        return FamilyResult(
            name=self.name,
            display_name=self.display_name,
            total=len(self.attacks),
            passed=passed,
            findings=aggregated,
        )


def _aggregate_findings(findings: list[Finding]) -> list[Finding]:
    """Merge per-task findings into per-attack findings."""
    groups: dict[tuple, Finding] = {}
    for f in findings:
        # Group by (attack, severity, full recommendation)
        key = (f.attack, f.severity, f.recommendation)
        if key not in groups:
            groups[key] = f.model_copy()
        else:
            existing = groups[key]
            # Merge affected_tasks
            existing.affected_tasks = list(set(existing.affected_tasks) | set(f.affected_tasks))
    return list(groups.values())


# ---------------------------------------------------------------------------
# Answer key normalization
# ---------------------------------------------------------------------------

_ANSWER_KEY_ALIASES = [
    "answer",
    "solution",
    "ground_truth",
    "target",
    "label",
    "expected",
    "expected_output",
    "gold",
    "reference",
    "correct_answer",
    "final_answer",
]


def _detect_answer_key(tasks: list[dict]) -> str | None:
    """Find the answer key name used in the task dicts."""
    if not tasks:
        return None
    sample = tasks[0]
    for key in _ANSWER_KEY_ALIASES:
        if key in sample:
            return key
    return None


def _normalize_answer_key(tasks: list[dict]) -> list[dict]:
    """Ensure all tasks have an 'answer' key by detecting aliases."""
    key = _detect_answer_key(tasks)
    if key is None or key == "answer":
        return tasks
    return [{**t, "answer": t[key]} if key in t and "answer" not in t else t for t in tasks]
