"""
Logic for running comparative experiments (AI baselines, House Rules, Racer Impact).
"""

from __future__ import annotations

import logging
import random
import time
import uuid
from collections import defaultdict
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, get_args

from tqdm import tqdm

from magsim.ai.baseline_agent import BaselineAgent
from magsim.ai.smart_agent import SmartAgent
from magsim.core.state import GameRules
from magsim.core.types import BoardName, RacerName
from magsim.engine.board import BOARD_DEFINITIONS
from magsim.engine.scenario import GameScenario, RacerConfig
from magsim.simulation.combinations import generate_combinations

if TYPE_CHECKING:
    from magsim.core.agent import Agent
    from magsim.simulation.config import GameConfig

logger = logging.getLogger(__name__)


@dataclass
class RacerStats:
    """Accumulator for racer performance statistics during experiments."""

    games_played_control: int = 0
    games_played_treatment: int = 0
    wins_control: int = 0
    wins_treatment: int = 0
    total_vp_control: int = 0
    total_vp_treatment: int = 0
    max_vp_control: int = 0
    max_vp_treatment: int = 0

    def update_control(self, *, is_winner: bool, vp: int) -> None:
        self.games_played_control += 1
        if is_winner:
            self.wins_control += 1
        self.total_vp_control += vp
        self.max_vp_control = max(self.max_vp_control, vp)

    def update_treatment(self, *, is_winner: bool, vp: int) -> None:
        self.games_played_treatment += 1
        if is_winner:
            self.wins_treatment += 1
        self.total_vp_treatment += vp
        self.max_vp_treatment = max(self.max_vp_treatment, vp)


@dataclass(frozen=True)
class SimulationResult:
    """Structured result of a single simulation run."""

    winner_idx: int
    vp_scores: dict[int, int]


@dataclass
class ExperimentResult:
    run_id: str
    timestamp: datetime
    racer: RacerName
    games_played: int

    # Win Rates
    winrate_control: float
    winrate_treatment: float

    # VP
    vp_control: float
    vp_treatment: float
    max_vp_control: int
    max_vp_treatment: int

    # Timing
    elapsed_control: float
    elapsed_treatment: float

    @property
    def winrate_delta(self) -> float:
        return self.winrate_treatment - self.winrate_control

    @property
    def winrate_pct_change(self) -> float:
        if self.winrate_control == 0:
            return 0.0
        return (self.winrate_treatment - self.winrate_control) / self.winrate_control

    @property
    def vp_delta(self) -> float:
        return self.vp_treatment - self.vp_control

    @property
    def vp_pct_change(self) -> float:
        if self.vp_control == 0:
            return 0.0
        return (self.vp_treatment - self.vp_control) / self.vp_control

    @property
    def speed_control(self) -> float:
        """Games per second (Control)."""
        return (
            self.games_played / self.elapsed_control
            if self.elapsed_control > 0
            else 0.0
        )

    @property
    def speed_treatment(self) -> float:
        """Games per second (Treatment)."""
        return (
            self.games_played / self.elapsed_treatment
            if self.elapsed_treatment > 0
            else 0.0
        )

    @property
    def speed_delta(self) -> float:
        return self.speed_treatment - self.speed_control

    @property
    def speed_pct_change(self) -> float:
        if self.speed_control == 0:
            return 0.0
        return (self.speed_treatment - self.speed_control) / self.speed_control


def run_ai_comparison(
    target_racer: RacerName,
    n_games: int = 500,
    seed_offset: int = 0,
    max_turns: int = 200,
) -> ExperimentResult:
    """Compare BaselineAgent vs SmartAgent for a specific racer."""

    all_racers = list(get_args(RacerName))
    all_boards = list(get_args(BoardName))
    estimated_needed = n_games * 10

    gen = generate_combinations(
        eligible_racers=all_racers,
        racer_counts=[4, 5, 6],
        boards=all_boards,
        runs_per_combination=1,
        max_total_runs=estimated_needed,
        seed_offset=seed_offset,
    )

    stats = RacerStats()
    time_control = 0.0
    time_treatment = 0.0

    games_processed = 0
    run_id = str(uuid.uuid4())[:8]

    with tqdm(total=n_games, desc=f"Comparing {target_racer}") as pbar:
        for config in gen:
            if games_processed >= n_games:
                break
            if target_racer not in config.racers:
                continue

            tqdm.write(f"[Sim {games_processed + 1}] {config.repr}")

            target_idx = config.racers.index(target_racer)

            # Control Run
            t0 = time.perf_counter()
            res_c = _run_config_with_setup(
                config,
                {target_idx: BaselineAgent()},
                GameRules(),
                max_turns=max_turns,
            )
            t1 = time.perf_counter()

            if res_c is None:
                tqdm.write(f"⚠️ Skipped Sim {games_processed}: Control timeout")
                continue

            time_control += t1 - t0
            stats.update_control(
                is_winner=(res_c.winner_idx == target_idx),
                vp=res_c.vp_scores[target_idx],
            )

            # Treatment Run
            t2 = time.perf_counter()
            res_t = _run_config_with_setup(
                config,
                {target_idx: SmartAgent()},
                GameRules(),
                max_turns=max_turns,
            )
            t3 = time.perf_counter()

            if res_t is None:
                tqdm.write(f"⚠️ Skipped Sim {games_processed}: Treatment timeout")
                # Note: We discard both if one fails to keep comparison fair?
                # Or keep control? Usually safe to discard both for strict A/B.
                # Since we already updated stats.control above, we should technically rollback or just accept the noise.
                # Given this is rare, noise is acceptable, but let's just log it.
                continue

            time_treatment += t3 - t2
            stats.update_treatment(
                is_winner=(res_t.winner_idx == target_idx),
                vp=res_t.vp_scores[target_idx],
            )

            games_processed += 1
            pbar.update(1)

    return ExperimentResult(
        run_id=run_id,
        timestamp=datetime.now(UTC),
        racer=target_racer,
        games_played=games_processed,
        winrate_control=stats.wins_control / games_processed if games_processed else 0,
        winrate_treatment=stats.wins_treatment / games_processed
        if games_processed
        else 0,
        vp_control=stats.total_vp_control / games_processed if games_processed else 0,
        vp_treatment=stats.total_vp_treatment / games_processed
        if games_processed
        else 0,
        max_vp_control=stats.max_vp_control,
        max_vp_treatment=stats.max_vp_treatment,
        elapsed_control=time_control,
        elapsed_treatment=time_treatment,
    )


def run_rule_comparison(
    rule_key: str,
    rule_value: Any,
    n_games: int = 1000,
    seed_offset: int = 0,
    max_turns: int = 200,
) -> list[ExperimentResult]:
    """Compare Default Rules vs Modified Rules across all racers."""

    all_racers = list(get_args(RacerName))
    all_boards = list(get_args(BoardName))

    gen = generate_combinations(
        eligible_racers=all_racers,
        racer_counts=[4, 5, 6],
        boards=all_boards,
        runs_per_combination=1,
        max_total_runs=n_games,
        seed_offset=seed_offset,
    )

    # Use a dictionary of typed objects instead of nested dicts
    racer_stats: dict[RacerName, RacerStats] = defaultdict(RacerStats)
    time_c_total = 0.0
    time_t_total = 0.0

    run_id = str(uuid.uuid4())[:8]
    timestamp = datetime.now(UTC)
    rules_c = GameRules()
    rules_t = GameRules(**{rule_key: rule_value})

    games_processed = 0

    with tqdm(total=n_games, desc=f"Testing {rule_key}={rule_value}") as pbar:
        for config in gen:
            tqdm.write(f"[Sim {games_processed + 1}] {config.repr}")

            t0 = time.perf_counter()
            res_c = _run_config_with_setup(config, {}, rules_c, max_turns=max_turns)
            t1 = time.perf_counter()

            if res_c is None:
                continue

            time_c_total += t1 - t0

            t2 = time.perf_counter()
            res_t = _run_config_with_setup(config, {}, rules_t, max_turns=max_turns)
            t3 = time.perf_counter()

            if res_t is None:
                continue

            time_t_total += t3 - t2

            for idx, r_name in enumerate(config.racers):
                stats = racer_stats[r_name]

                # Control Update
                stats.update_control(
                    is_winner=(res_c.winner_idx == idx),
                    vp=res_c.vp_scores[idx],
                )

                # Treatment Update
                stats.update_treatment(
                    is_winner=(res_t.winner_idx == idx),
                    vp=res_t.vp_scores[idx],
                )

            games_processed += 1
            pbar.update(1)

    results: list[ExperimentResult] = []
    if games_processed > 0:
        avg_time_c = time_c_total / games_processed
        avg_time_t = time_t_total / games_processed
    else:
        avg_time_c = 0.0
        avg_time_t = 0.0

    for r_name, stats in racer_stats.items():
        n = stats.games_played_control  # Or treatment, they are equal here
        if n == 0:
            continue

        results.append(
            ExperimentResult(
                run_id=run_id,
                timestamp=timestamp,
                racer=r_name,
                games_played=n,
                winrate_control=stats.wins_control / n,
                winrate_treatment=stats.wins_treatment / n,
                vp_control=stats.total_vp_control / n,
                vp_treatment=stats.total_vp_treatment / n,
                max_vp_control=stats.max_vp_control,
                max_vp_treatment=stats.max_vp_treatment,
                elapsed_control=avg_time_c * n,
                elapsed_treatment=avg_time_t * n,
            ),
        )

    return results


def run_racer_impact_comparison(
    target_racer: RacerName,
    n_games: int = 1000,
    seed_offset: int = 0,
    max_turns: int = 200,
) -> list[ExperimentResult]:
    """
    Measure the impact of a specific racer on the rest of the field.
    """
    all_racers = list(get_args(RacerName))
    all_boards = list(get_args(BoardName))
    racer_pool = set(all_racers)

    gen = generate_combinations(
        eligible_racers=all_racers,
        racer_counts=[4, 5, 6],
        boards=all_boards,
        runs_per_combination=1,
        max_total_runs=n_games * 10,
        seed_offset=seed_offset,
    )

    racer_stats: dict[RacerName, RacerStats] = defaultdict(RacerStats)
    time_c_total = 0.0
    time_t_total = 0.0

    run_id = str(uuid.uuid4())[:8]
    timestamp = datetime.now(UTC)

    games_processed = 0
    rng = random.Random(seed_offset)

    with tqdm(total=n_games, desc=f"Impact of {target_racer}") as pbar:
        for config in gen:
            if games_processed >= n_games:
                break

            if target_racer not in config.racers:
                continue

            # 1. Setup Treatment (With Target)
            tqdm.write(f"[Sim {games_processed + 1} Treatment] {config.repr}")

            t2 = time.perf_counter()
            res_t = _run_config_with_setup(config, {}, GameRules(), max_turns=max_turns)
            t3 = time.perf_counter()

            if res_t is None:
                continue

            time_t_total += t3 - t2

            for idx, r_name in enumerate(config.racers):
                racer_stats[r_name].update_treatment(
                    is_winner=(res_t.winner_idx == idx),
                    vp=res_t.vp_scores[idx],
                )

            # 2. Setup Control (Without Target)
            current_set = set(config.racers)
            available = list(racer_pool - current_set)
            if not available:
                continue

            replacement = rng.choice(available)
            target_idx = config.racers.index(target_racer)
            new_roster = list(config.racers)
            new_roster[target_idx] = replacement

            config_c = config.__class__(
                racers=tuple(new_roster),
                board=config.board,
                seed=config.seed,
                rules=config.rules,
            )

            tqdm.write(f"[Sim {games_processed + 1} Control] {config_c.repr}")

            t0 = time.perf_counter()
            res_c = _run_config_with_setup(
                config_c,
                {},
                GameRules(),
                max_turns=max_turns,
            )
            t1 = time.perf_counter()

            if res_c is None:
                continue

            time_c_total += t1 - t0

            for idx, r_name in enumerate(config_c.racers):
                racer_stats[r_name].update_control(
                    is_winner=(res_c.winner_idx == idx),
                    vp=res_c.vp_scores[idx],
                )

            games_processed += 1
            pbar.update(1)

    results: list[ExperimentResult] = []
    avg_time_c = time_c_total / games_processed if games_processed else 0
    avg_time_t = time_t_total / games_processed if games_processed else 0

    for r_name, stats in racer_stats.items():
        # In impact analysis, 'games_played' differs for control vs treatment for specific racers
        gc = stats.games_played_control
        gt = stats.games_played_treatment

        # Filter noise
        if gc < 5 and gt < 5:
            continue

        wc = stats.wins_control / gc if gc > 0 else 0.0
        wt = stats.wins_treatment / gt if gt > 0 else 0.0
        vc = stats.total_vp_control / gc if gc > 0 else 0.0
        vt = stats.total_vp_treatment / gt if gt > 0 else 0.0

        # Use whichever side had presence for "games_played" count in summary
        n = gt if gt > 0 else gc

        results.append(
            ExperimentResult(
                run_id=run_id,
                timestamp=timestamp,
                racer=r_name,
                games_played=n,
                winrate_control=wc,
                winrate_treatment=wt,
                vp_control=vc,
                vp_treatment=vt,
                max_vp_control=stats.max_vp_control,
                max_vp_treatment=stats.max_vp_treatment,
                elapsed_control=avg_time_c * gc,
                elapsed_treatment=avg_time_t * gt,
            ),
        )

    return results


def _run_config_with_setup(
    config: GameConfig,
    agent_overrides: dict[int, Agent],
    rules: GameRules,
    max_turns: int = 200,
) -> SimulationResult | None:
    """Unified internal runner. Returns None if max turns reached."""
    r_configs: list[RacerConfig] = []
    for i, name in enumerate(config.racers):
        agent = agent_overrides.get(i)
        r_configs.append(RacerConfig(idx=i, name=name, agent=agent))

    scenario = GameScenario(
        racers_config=r_configs,
        board=BOARD_DEFINITIONS[config.board](),
        seed=config.seed,
        rules=rules,
    )

    # Manual loop execution with safety break
    turn = 0
    try:
        while scenario.state.race_active:
            if turn >= max_turns:
                return None
            scenario.run_turn()
            turn += 1
    except Exception:  # noqa: BLE001
        # Log error if needed, but for baseline we often skip broken runs
        return None

    winner_idx = -1
    for r in scenario.state.racers:
        if r.finish_position == 1:
            winner_idx = r.idx
            break

    vps = {r.idx: r.victory_points for r in scenario.state.racers}
    return SimulationResult(winner_idx=winner_idx, vp_scores=vps)
