"""V4: movement uses the best of land and fly speed.

The engine budgets every movement action (move, 5-foot step toward a target,
charge at 2x, withdraw at 2x, run at 4x) from ``_speed_ft(state)`` — the
maximum of the creature's ``speed_land_ft`` and ``speed_fly_ft`` (a missing
fly speed counts as 0, so a land-only creature is unchanged). Simplification:
a creature that has a fly speed is assumed to use it, but flying OVER
obstacles or difficult terrain is NOT modeled — the grid's passability and
terrain costs still apply to the whole budget. The observation snapshot adds
an additive ``speed_effective_ft`` key; ``speed_land_ft``/``speed_fly_ft`` are
kept as-is.
"""
from __future__ import annotations

from collections.abc import Callable

from pf1e_simulator.combat import (
    Action,
    CombatantState,
    CombatEngine,
)
from pf1e_simulator.dice import parse_dice
from pf1e_simulator.grid import Grid
from pf1e_simulator.map import MapSpec, TerrainType
from pf1e_simulator.models import (
    AbilityScores,
    ACProfile,
    AttackSpec,
    Combatant,
    DamageComponent,
    Saves,
)
from pf1e_simulator.rng import ScriptedRng
from tests.test_observation_enriched import make_state

PolicyFn = Callable[[CombatEngine, CombatantState], tuple[Action, ...]]

_WAIT: tuple[Action, ...] = (Action(kind="wait"),)


def _combatant(
    id_: str, *, land: int, fly: int | None = None, initiative_mod: int = 0
) -> Combatant:
    return Combatant(
        id=id_,
        name=id_,
        level=1,
        abilities=AbilityScores(
            str_score=10, dex_score=10, con_score=10,
            int_score=10, wis_score=10, cha_score=10,
        ),
        ac=ACProfile(total=10, touch=10, flat_footed=10),
        bab=1,
        initiative_mod=initiative_mod,
        speed_land_ft=land,
        speed_fly_ft=fly,
        hp_max=100,
        saves=Saves(fort=0, ref=0, will=0),
        attacks=[
            AttackSpec(
                id=f"{id_}-sword",
                name="short sword",
                kind="melee",
                attack_bonus=5,
                damage=[
                    DamageComponent(formula=parse_dice("1d6"), types=["slashing"])
                ],
            )
        ],
    )


def _corridor(cols: int) -> Grid:
    """A single-row open floor of ``cols`` cells (all move_cost 1)."""
    legend = {".": TerrainType(type="floor", move_cost=1)}
    spec = MapSpec(name="corridor", terrain=("." * cols,), legend=legend)
    return Grid.from_spec(spec)


def _run_move(
    *, land: int, fly: int | None, destination: tuple[int, int], cols: int,
) -> CombatantState:
    """Hero acts first, takes a single ``move`` to ``destination``; foe idles
    far down the corridor (never threatens). Returns the hero state."""
    hero = make_state(
        _combatant("hero", land=land, fly=fly, initiative_mod=2),
        "players",
        (0, 0),
    )
    foe = make_state(_combatant("foe", land=30), "monsters", (0, cols - 1))

    def policy(_engine: CombatEngine, state: CombatantState) -> tuple[Action, ...]:
        if state.combatant.id == "hero":
            return (Action(kind="move", destination=destination),)
        return _WAIT

    eng = CombatEngine(
        ScriptedRng([10, 5]), _corridor(cols), [hero, foe],
        policy=policy, round_cap=1,
    )
    eng.run()
    return eng._states[0]


# ── (a) a land-only creature moves at its land speed (unchanged) ──────────────


def test_land_only_creature_uses_land_speed() -> None:
    # speed 30 ft = 6 cells; (0, 6) is exactly reachable, (0, 7) is not.
    hero = _run_move(land=30, fly=None, destination=(0, 6), cols=12)
    assert hero.pos == (0, 6)

    hero = _run_move(land=30, fly=None, destination=(0, 7), cols=12)
    assert hero.pos == (0, 0)


# ── (b) a fly 80 / land 30 creature moves up to 80 ft (16 cells) ──────────────


def test_flyer_moves_farther_than_land_only() -> None:
    # fly 80 ft = 16 cells: (0, 16) is within the move budget.
    flyer = _run_move(land=30, fly=80, destination=(0, 16), cols=20)
    assert flyer.pos == (0, 16)

    # A land-only creature cannot reach the same cell this turn.
    crawler = _run_move(land=30, fly=None, destination=(0, 16), cols=20)
    assert crawler.pos == (0, 0)
    assert _corridor(20).distance(flyer.pos, (0, 0)) > _corridor(20).distance(
        crawler.pos, (0, 0)
    )


# ── (c) charge uses 2x the effective speed; run uses 4x ──────────────────────


def test_charge_budget_is_double_the_fly_speed() -> None:
    # fly 80 -> charge up to 2x16 = 32 cells; the foe at (0, 20) is chargeable.
    flyer = make_state(
        _combatant("hero", land=30, fly=80), "players", (0, 0)
    )
    eng = CombatEngine(
        ScriptedRng([]), _corridor(24),
        [flyer, make_state(_combatant("foe", land=30), "monsters", (0, 20))],
    )
    result = eng.can_charge(flyer, eng._states[1])
    assert result is not None
    assert result.weapon is not None
    assert result.reason is None

    # A land-only creature (max 2x6 = 12 cells) cannot reach (0, 20).
    crawler = make_state(_combatant("hero", land=30), "players", (0, 0))
    foe2 = make_state(_combatant("foe", land=30), "monsters", (0, 20))
    eng2 = CombatEngine(ScriptedRng([]), _corridor(24), [crawler, foe2])
    result2 = eng2.can_charge(crawler, foe2)
    assert result2 is not None
    assert result2.reason == "no_charge_path"


def test_run_budget_is_four_times_the_fly_speed() -> None:
    # fly 80 -> run 4x16 = 64 cells; (0, 40) is reachable by running.
    hero = make_state(
        _combatant("hero", land=30, fly=80, initiative_mod=2), "players", (0, 0)
    )
    foe = make_state(_combatant("foe", land=30), "monsters", (0, 63))

    def policy(_engine: CombatEngine, state: CombatantState) -> tuple[Action, ...]:
        if state.combatant.id == "hero":
            return (Action(kind="run", destination=(0, 40)),)
        return _WAIT

    eng = CombatEngine(
        ScriptedRng([10, 5]), _corridor(64), [hero, foe],
        policy=policy, round_cap=1,
    )
    eng.run()
    assert eng._states[0].pos == (0, 40)

    # A land-only creature (run 4x6 = 24 cells) cannot reach (0, 40).
    crawler = make_state(
        _combatant("hero", land=30, fly=None, initiative_mod=2), "players", (0, 0)
    )
    foe2 = make_state(_combatant("foe", land=30), "monsters", (0, 63))
    eng2 = CombatEngine(
        ScriptedRng([10, 5]), _corridor(64), [crawler, foe2],
        policy=policy, round_cap=1,
    )
    eng2.run()
    assert eng2._states[0].pos == (0, 0)


# ── (d) the snapshot exposes speed_effective_ft ───────────────────────────────


def test_snapshot_exposes_speed_effective_ft() -> None:
    flyer = make_state(_combatant("flyer", land=30, fly=80), "players", (0, 0))
    land_only = make_state(_combatant("walker", land=30), "monsters", (0, 3))
    eng = CombatEngine(ScriptedRng([]), _corridor(8), [flyer, land_only])
    obs = eng._build_observation(flyer).to_dict()
    by_id = {c["id"]: c for c in obs["combatants"]}
    assert by_id["flyer"]["speed_effective_ft"] == 80
    assert by_id["flyer"]["speed_land_ft"] == 30
    assert by_id["flyer"]["speed_fly_ft"] == 80
    assert by_id["walker"]["speed_effective_ft"] == 30
    assert by_id["walker"]["speed_fly_ft"] is None


# ── (e) a fly speed lower than land falls back to land (max) ──────────────────


def test_fly_lower_than_land_uses_land() -> None:
    hero = make_state(_combatant("hero", land=30, fly=10), "players", (0, 0))
    foe = make_state(_combatant("foe", land=30), "monsters", (0, 7))
    eng = CombatEngine(ScriptedRng([]), _corridor(8), [hero, foe])
    obs = eng._build_observation(hero).to_dict()
    by_id = {c["id"]: c for c in obs["combatants"]}
    assert by_id["hero"]["speed_effective_ft"] == 30

    # Movement still uses the 30 ft land budget (6 cells), not the 10 ft fly.
    mover = _run_move(land=30, fly=10, destination=(0, 6), cols=12)
    assert mover.pos == (0, 6)
    mover = _run_move(land=30, fly=10, destination=(0, 7), cols=12)
    assert mover.pos == (0, 0)