"""Reward function simulator.

Generates diverse completions, scores them against reward functions,
and analyzes the reward landscape to predict training behavior.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from rewardprobe.families.base import _safe_call
from rewardprobe.types import NormalizedEnvironment

# ── Completion generation ────────────────────────────────────────────

DIVERSITY_SYSTEM = """\
You generate diverse completions to map how a reward function behaves \
across the full range of model outputs during RL training.

Your completions must look like realistic language model output. \
Study the task, the expected answer format, and the reward function \
description to produce domain-appropriate completions — do NOT assume \
any specific format (like \\boxed{{}} or <think> tags) unless the reward \
function description tells you it uses them."""

DIVERSITY_USER = """\
Task: {prompt}
Correct answer: {answer}
Reward function: {reward_description}

Generate exactly 10 diverse completions as a JSON array. Each must \
represent a different STRATEGY a model might learn during RL training. \
Adapt the strategies to this specific task domain:

1. perfect: Genuinely correct solution with clear reasoning
2. correct_lazy: Correct answer but minimal effort (just the answer, no work shown)
3. correct_verbose: Correct answer buried in excessive, rambling text
4. near_miss: Plausible attempt but answer is slightly wrong (off by one, \
close synonym, minor error)
5. wrong_confident: Confident and well-structured but fundamentally wrong
6. format_only: Uses the expected output format perfectly but content is \
wrong or nonsensical
7. echo: Copies or paraphrases the input/question instead of answering
8. hedge: Provides multiple possible answers hoping one matches
9. shortcut: Attempts to answer without actually doing the required work \
(guessing, pattern matching, or copying from the question)
10. garbage: Incoherent, off-topic, or random text

IMPORTANT: Each completion should look like what a REAL language model \
would produce. Match the expected format and style for this domain.

Respond with ONLY a JSON array of 10 objects:
[
  {{"strategy": "perfect", "completion": "...full completion text..."}},
  {{"strategy": "correct_lazy", "completion": "..."}},
  ...
]"""


@dataclass
class ScoredCompletion:
    """A completion scored against all reward functions."""

    strategy: str
    completion: str
    scores: dict[str, float] = field(default_factory=dict)
    rubric_score: float = 0.0
    is_correct: bool = False


@dataclass
class LandscapeInsight:
    """A single insight from reward landscape analysis."""

    severity: str  # critical, warning, info
    title: str
    description: str
    evidence: str = ""


@dataclass
class SimulationResult:
    """Full simulation result."""

    environment: str
    tasks_simulated: int
    completions_generated: int
    insights: list[LandscapeInsight] = field(default_factory=list)
    scored_completions: list[ScoredCompletion] = field(default_factory=list)


def generate_diverse_completions(
    task: dict[str, Any],
    client: Any,
    reward_description: str = "",
) -> list[dict[str, str]]:
    """Generate 10 diverse completions spanning the strategy space."""
    prompt_text = task.get("prompt", task.get("question", ""))
    if isinstance(prompt_text, list):
        prompt_text = next(
            (
                m.get("content", "")
                for m in reversed(prompt_text)
                if m.get("role") == "user"
            ),
            str(prompt_text),
        )

    answer = task.get("answer", "")
    desc = reward_description or "Unknown — generate completions in multiple formats."

    user_prompt = DIVERSITY_USER.format(
        prompt=prompt_text,
        answer=answer,
        reward_description=desc,
    )

    try:
        result = client.complete_json(DIVERSITY_SYSTEM, user_prompt)
        if isinstance(result, list):
            return result
        return result.get("completions", result.get("items", []))
    except Exception:
        return []


def score_completions(
    completions: list[dict[str, str]],
    env: NormalizedEnvironment,
    task: dict[str, Any],
) -> list[ScoredCompletion]:
    """Score each completion against all reward functions."""
    scored = []

    for comp_data in completions:
        strategy = comp_data.get("strategy", "unknown")
        text = comp_data.get("completion", "")
        if not text:
            continue

        sc = ScoredCompletion(strategy=strategy, completion=text)

        # Score against each reward function
        total_weight = sum(fn.weight for fn in env.reward_fns)
        weighted_sum = 0.0

        for fn in env.reward_fns:
            score, err = _safe_call(fn, text, task)
            val = score if score is not None and err is None else 0.0
            sc.scores[fn.name] = val
            weighted_sum += val * fn.weight

        sc.rubric_score = weighted_sum / total_weight if total_weight > 0 else 0.0

        # Heuristic: is this completion "correct"?
        # A completion is correct if it's in the "perfect" strategy category
        # or if the answer appears in the completion in a reasonable way
        sc.is_correct = strategy in ("perfect", "correct_lazy", "correct_verbose")

        scored.append(sc)

    return scored


# ── Landscape analysis ───────────────────────────────────────────────


def analyze_landscape(
    all_scored: list[ScoredCompletion],
    env: NormalizedEnvironment,
) -> list[LandscapeInsight]:
    """Analyze scored completions to find reward landscape issues."""
    insights = []

    if not all_scored:
        return insights

    # ── 1. Wrong completions that score high ──
    wrong_high = [
        s
        for s in all_scored
        if not s.is_correct and s.rubric_score > 0.5
    ]
    correct_scores = [s.rubric_score for s in all_scored if s.is_correct]
    avg_correct = sum(correct_scores) / len(correct_scores) if correct_scores else 0

    if wrong_high:
        worst = max(wrong_high, key=lambda s: s.rubric_score)
        pct = len(wrong_high) / max(len(all_scored), 1) * 100

        insights.append(
            LandscapeInsight(
                severity="critical",
                title="Wrong completions score high",
                description=(
                    f"{len(wrong_high)} of {len(all_scored)} completions "
                    f"({pct:.0f}%) score above 0.5 despite being incorrect. "
                    f"Correct completions average {avg_correct:.2f}. "
                    f"A model can achieve {worst.rubric_score:.2f} reward "
                    f"with the '{worst.strategy}' strategy."
                ),
                evidence=(
                    f"Strategy '{worst.strategy}': "
                    f"\"{worst.completion[:120]}...\" "
                    f"→ {worst.rubric_score:.2f}"
                ),
            )
        )

    # ── 2. Degenerate strategies ──
    strategy_scores: dict[str, list[float]] = {}
    for s in all_scored:
        strategy_scores.setdefault(s.strategy, []).append(s.rubric_score)

    degenerate_strategies = [
        "format_only",
        "echo",
        "garbage",
        "shortcut",
        "correct_lazy",
    ]

    for strat in degenerate_strategies:
        scores = strategy_scores.get(strat, [])
        if not scores:
            continue
        avg = sum(scores) / len(scores)
        if avg > 0.4:
            severity = "critical" if avg > 0.6 else "warning"
            label = {
                "format_only": "Format-only (no real content)",
                "echo": "Echo (copies the question)",
                "garbage": "Garbage (incoherent text)",
                "shortcut": "Shortcut (skips computation)",
                "correct_lazy": "Lazy correct (no reasoning shown)",
            }.get(strat, strat)

            insights.append(
                LandscapeInsight(
                    severity=severity,
                    title=f"'{label}' strategy scores {avg:.2f}",
                    description=(
                        f"A model using the '{strat}' strategy earns "
                        f"an average reward of {avg:.2f} "
                        f"({avg / max(avg_correct, 0.01) * 100:.0f}% of "
                        f"a correct answer's {avg_correct:.2f}). "
                        "This strategy requires less effort than genuine "
                        "problem-solving, so a model may converge toward it."
                    ),
                )
            )

    # ── 3. Weight dominance (multi-fn rubrics) ──
    if len(env.reward_fns) >= 2:
        # Check if any single fn can dominate
        for fn in env.reward_fns:
            # Find completions that score high on this fn but low on others
            fn_dominant = []
            for s in all_scored:
                fn_score = s.scores.get(fn.name, 0)
                other_avg = 0.0
                other_count = 0
                for other_fn in env.reward_fns:
                    if other_fn.name != fn.name:
                        other_avg += s.scores.get(other_fn.name, 0)
                        other_count += 1
                if other_count > 0:
                    other_avg /= other_count
                if fn_score > 0.8 and other_avg < 0.3:
                    fn_dominant.append(s)

            if fn_dominant:
                best = max(fn_dominant, key=lambda s: s.rubric_score)
                if best.rubric_score > 0.4:
                    insights.append(
                        LandscapeInsight(
                            severity="warning",
                            title=f"'{fn.name}' can dominate the rubric",
                            description=(
                                f"A model scoring 1.0 on '{fn.name}' "
                                f"(weight {fn.weight}) and 0.0 on other "
                                f"functions still achieves {best.rubric_score:.2f} "
                                f"overall reward. This means a model can "
                                f"optimize only for {fn.name} and ignore "
                                "other objectives."
                            ),
                        )
                    )

    # ── 4. Near-miss acceptance ──
    near_misses = [
        s for s in all_scored if s.strategy == "near_miss"
    ]
    accepted_near = [s for s in near_misses if s.rubric_score > 0.5]
    if accepted_near:
        insights.append(
            LandscapeInsight(
                severity="critical",
                title="Near-miss answers accepted",
                description=(
                    f"{len(accepted_near)} of {len(near_misses)} near-miss "
                    "completions (close but wrong answers) score above 0.5. "
                    "During training, the model receives positive reward for "
                    "almost-right answers, weakening the training signal."
                ),
                evidence=(
                    f"\"{accepted_near[0].completion[:120]}...\" "
                    f"→ {accepted_near[0].rubric_score:.2f}"
                ),
            )
        )

    # ── 5. Hedge strategy ──
    hedges = [s for s in all_scored if s.strategy == "hedge"]
    accepted_hedge = [s for s in hedges if s.rubric_score > 0.5]
    if accepted_hedge:
        insights.append(
            LandscapeInsight(
                severity="warning",
                title="Hedging strategy rewarded",
                description=(
                    "Completions that output multiple possible answers "
                    f"score {accepted_hedge[0].rubric_score:.2f}. A model "
                    "can learn to hedge its bets instead of committing to "
                    "one answer."
                ),
            )
        )

    return insights


# ── Main simulate function ───────────────────────────────────────────


def _build_reward_description(env: NormalizedEnvironment) -> str:
    """Build a description of the reward functions for the LLM."""
    parts = []
    for fn in env.reward_fns:
        # Try to get source code for a brief description
        try:
            import inspect

            target = fn.fn
            if hasattr(target, "__closure__") and target.__closure__:
                for cell in target.__closure__:
                    try:
                        val = cell.cell_contents
                        if callable(val) and hasattr(val, "__name__"):
                            target = val
                            break
                    except ValueError:
                        pass
            source = inspect.getsource(target)
            # Just the first 3 lines for context
            first_lines = "\n".join(source.strip().split("\n")[:3])
            parts.append(f"- {fn.name} (weight {fn.weight}): {first_lines}")
        except (OSError, TypeError):
            parts.append(f"- {fn.name} (weight {fn.weight})")
    return "Reward functions:\n" + "\n".join(parts)


def simulate(
    env: NormalizedEnvironment,
    client: Any,
    n_tasks: int = 5,
) -> SimulationResult:
    """Run a full reward landscape simulation."""
    tasks = env.tasks[:n_tasks]
    all_scored: list[ScoredCompletion] = []
    reward_desc = _build_reward_description(env)

    for task in tasks:
        completions = generate_diverse_completions(task, client, reward_desc)
        scored = score_completions(completions, env, task)
        all_scored.extend(scored)

    insights = analyze_landscape(all_scored, env)

    return SimulationResult(
        environment=env.name,
        tasks_simulated=len(tasks),
        completions_generated=len(all_scored),
        insights=insights,
        scored_completions=all_scored,
    )


# ── Terminal output ──────────────────────────────────────────────────


def print_simulation(result: SimulationResult) -> None:
    """Print simulation results to terminal."""
    from rich.console import Console
    from rich.text import Text

    con = Console(width=90)

    con.print()
    con.print(
        Text("rewardprobe simulate", style="bold cyan")
        + Text(f" — {result.environment}", style="bold")
    )
    con.print(
        Text(
            f"  {result.completions_generated} completions across "
            f"{result.tasks_simulated} tasks",
            style="dim",
        )
    )
    con.print()

    if not result.insights:
        con.print(Text("  No issues found.", style="bold green"))
        con.print()
        return

    # Count by severity
    n_crit = sum(1 for i in result.insights if i.severity == "critical")
    n_warn = sum(1 for i in result.insights if i.severity == "warning")

    parts = []
    if n_crit:
        parts.append(f"[bold red]{n_crit} critical[/]")
    if n_warn:
        parts.append(f"[bold yellow]{n_warn} warning[/]")
    con.print(f"  {', '.join(parts)} found\n")

    for i, insight in enumerate(result.insights, 1):
        if insight.severity == "critical":
            tag = Text(" critical ", style="bold white on red")
        else:
            tag = Text(" warning ", style="bold black on yellow")

        con.print(Text(f"  {i}. ", style="bold") + tag)
        con.print(Text(f"     {insight.title}", style="bold"))
        con.print(Text(f"     {insight.description}", style=""))
        if insight.evidence:
            con.print(Text(f"     {insight.evidence}", style="dim"))
        con.print()

    # Strategy scoreboard
    con.print(Text("  Strategy scoreboard:", style="bold"))
    strategy_avgs: dict[str, list[float]] = {}
    for s in result.scored_completions:
        strategy_avgs.setdefault(s.strategy, []).append(s.rubric_score)

    rows = []
    for strat, scores in sorted(
        strategy_avgs.items(), key=lambda x: -sum(x[1]) / len(x[1])
    ):
        avg = sum(scores) / len(scores)
        bar = "█" * int(avg * 20) + "░" * (20 - int(avg * 20))
        rows.append((strat, avg, bar))

    for strat, avg, bar in rows:
        is_correct_strat = strat.startswith("correct") or strat == "perfect"
        color = "green" if is_correct_strat else "red" if avg > 0.5 else "dim"
        con.print(
            Text(f"    {strat:20s} ", style="")
            + Text(f"{bar} ", style=color)
            + Text(f"{avg:.2f}", style="bold " + color)
        )

    con.print()
