from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Self, override

from magsim.ai.evaluation import (
    get_benefit_at,
    get_hazard_at,
)
from magsim.core.abilities import Ability
from magsim.core.agent import (
    Agent,
    BooleanDecisionMixin,
    DecisionContext,
)
from magsim.core.events import (
    AbilityTriggeredEvent,
    AbilityTriggeredEventOrSkipped,
    GameEvent,
    RollModificationWindowEvent,
    TurnStartEvent,
)
from magsim.engine.roll import trigger_reroll

if TYPE_CHECKING:
    from magsim.core.state import ActiveRacerState
    from magsim.core.types import AbilityName, D6VAlueSet
    from magsim.engine.game_engine import GameEngine


@dataclass
class AbilityMagicalReroll(Ability, BooleanDecisionMixin):
    name: AbilityName = "MagicalReroll"
    triggers: tuple[type[GameEvent], ...] = (
        RollModificationWindowEvent,
        TurnStartEvent,
    )

    # Local State
    reroll_count: int = 0
    preferred_dice: D6VAlueSet = frozenset({4, 5, 6})

    @override
    def execute(
        self,
        event: GameEvent,
        owner: ActiveRacerState,
        engine: GameEngine,
        agent: Agent,
    ) -> AbilityTriggeredEventOrSkipped:
        # 1. Reset logic
        if isinstance(event, TurnStartEvent):
            if event.target_racer_idx == owner.idx:
                self.reroll_count = 0
            return "skip_trigger"

        # 2. Reroll Logic
        if not isinstance(event, RollModificationWindowEvent):
            return "skip_trigger"

        # 1. Eligibility Check
        if event.target_racer_idx != owner.idx or self.reroll_count >= 2:
            return "skip_trigger"

        should_reroll = agent.make_boolean_decision(
            engine,
            ctx=DecisionContext(
                source=self,
                event=event,
                game_state=engine.state,
                source_racer_idx=owner.idx,
            ),
        )

        if not should_reroll:
            return "skip_trigger"

        self.reroll_count += 1
        engine.push_event(
            AbilityTriggeredEvent(
                owner.idx,
                source=self.name,
                phase=event.phase,
                target_racer_idx=owner.idx,
            ),
        )
        trigger_reroll(engine, owner.idx, "MagicalReroll")
        # ability trigger handled by trigger_reroll
        return "skip_trigger"

    @override
    def get_baseline_boolean_decision(
        self,
        engine: GameEngine,
        ctx: DecisionContext[Self],
    ) -> bool:
        dice_val = ctx.game_state.roll_state.dice_value
        return dice_val is not None and dice_val not in self.preferred_dice

    @override
    def get_auto_boolean_decision(
        self,
        engine: GameEngine,
        ctx: DecisionContext[Self],
    ) -> bool:
        if (me := engine.get_active_racer(ctx.source_racer_idx)) is None or (
            raw_roll := ctx.game_state.roll_state.dice_value
        ) is None:
            return False

        # Use final_value to account for modifiers like Gunk
        dest = me.position + ctx.game_state.roll_state.final_value

        # 1. PRIORITY: Keep if Excellent (Win / VP / Boost)
        if benefit := get_benefit_at(engine, dest):
            engine.log_info(f"{me.repr} keeps roll to reach {benefit}!")
            return False

        # 2. AVOIDANCE: Reroll if Hazard (Trip / Backward)
        if hazard := get_hazard_at(engine, dest):
            engine.log_info(f"{me.repr} uses {self.name} to avoid {hazard}!")
            return True

        # 3. GREED: Threshold based on remaining rerolls
        threshold = 5 if self.reroll_count == 0 else 4

        if raw_roll < threshold:
            engine.log_info(f"{me.repr} uses {self.name} due to low roll ({raw_roll})!")
            return True

        return False
