"""Stage 2: LLM-generated adversarial completions."""

from dataclasses import dataclass
from typing import Any

from rewardprobe.families.base import _safe_call
from rewardprobe.tier2.client import LLMClient
from rewardprobe.tier2.code_analyzer import CodeAnalysis
from rewardprobe.tier2.prompts import ADVERSARIAL_SYSTEM, ADVERSARIAL_USER
from rewardprobe.types import Finding, NormalizedRewardFn, Severity


@dataclass
class AdversarialResult:
    attack_type: str  # near_miss, wrong_work_right_answer, confident_wrong, format_exploit
    completion: str
    why_wrong: str
    score: float | None = None
    error: str | None = None
    is_exploit: bool = False  # True if score > threshold for a wrong completion


def generate_adversarial(
    fn: NormalizedRewardFn,
    task: dict[str, Any],
    analysis: CodeAnalysis,
    client: LLMClient,
    threshold: float = 0.5,
) -> list[AdversarialResult]:
    """Generate adversarial completions for a single task and test them."""
    answer = task.get("answer", "")
    prompt_text = task.get("prompt", task.get("question", ""))
    if isinstance(prompt_text, list):
        # Chat format — extract last user message
        prompt_text = next(
            (m.get("content", "") for m in reversed(prompt_text) if m.get("role") == "user"),
            str(prompt_text),
        )

    user_prompt = ADVERSARIAL_USER.format(
        prompt=prompt_text,
        answer=answer,
        function_description=f"{analysis.what_it_checks} (type: {analysis.function_type})",
    )

    try:
        result = client.complete_json(ADVERSARIAL_SYSTEM, user_prompt)
    except Exception:
        return []

    results = []
    for attack_type in [
        "near_miss",
        "wrong_work_right_answer",
        "confident_wrong",
        "format_exploit",
    ]:
        entry = result.get(attack_type, {})
        if not entry:
            continue
        completion = entry.get("completion", "")
        why = entry.get(
            "why_wrong", entry.get("why_problematic", entry.get("what_it_exploits", ""))
        )

        if not completion:
            continue

        # Run the completion against the actual reward function
        score, err = _safe_call(fn, completion, task)

        is_exploit = False
        if err is None and score is not None and score > threshold:
            # This wrong completion scored high — it's an exploit!
            # Exception: wrong_work_right_answer has the correct answer,
            # so scoring high is expected for correctness fns.
            # But it's still a finding because the reasoning is wrong.
            is_exploit = True

        results.append(
            AdversarialResult(
                attack_type=attack_type,
                completion=completion,
                why_wrong=why,
                score=score,
                error=err,
                is_exploit=is_exploit,
            )
        )

    return results


def adversarial_to_findings(
    results: list[AdversarialResult],
    fn_name: str,
    task_idx: int,
) -> list[Finding]:
    """Convert adversarial results to Finding objects."""
    findings = []
    for r in results:
        if not r.is_exploit:
            continue

        severity = (
            Severity.CRITICAL
            if r.attack_type in ("near_miss", "confident_wrong")
            else Severity.WARNING
        )

        findings.append(
            Finding(
                attack=f"adversarial_{r.attack_type}",
                family="adversarial",
                severity=severity,
                input=r.completion[:300],
                expected_reward=0.0,
                actual_reward=r.score,
                affected_tasks=[task_idx],
                recommendation=(
                    f"[AI-generated] Reward fn '{fn_name}' scores a "
                    f"{r.attack_type.replace('_', ' ')} completion "
                    f"at {r.score:.3f}. {r.why_wrong}"
                ),
            )
        )
    return findings
