"""V4: movement uses the best of land and fly speed + 3D flight over terrain.

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). A creature
whose budget comes from its fly speed (fly >= land) is AIRBORNE
(``_is_airborne``): it moves on ``Grid.airborne_view``, where ground difficult
terrain costs 1 and cells impassable to walkers (pits, water, rubble) are
crossable — SRD "Movement, Position, and Distance": "Flying and incorporeal
creatures are not hampered by difficult terrain"; "Flying and incorporeal
creatures are able to avoid most obstacles". Cells that block line of effect
(``blocks_los`` — solid barriers, assumed full height) still stop flight and
still block attacks. 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)


_LEGEND = {
    ".": TerrainType(type="floor", move_cost=1),
    "T": TerrainType(type="rubble", move_cost=2),
    "P": TerrainType(type="pit", move_cost=None),
    "#": TerrainType(type="wall", move_cost=None, blocks_los=True),
}


def _terrain_grid(rows: tuple[str, ...]) -> Grid:
    """A grid from ASCII rows over the shared ``_LEGEND`` terrain types."""
    spec = MapSpec(name="terrain", terrain=rows, legend=_LEGEND)
    return Grid.from_spec(spec)


def _run_move_on(
    *, land: int, fly: int | None, destination: tuple[int, int], grid: Grid,
) -> CombatantState:
    """Hero acts first and takes one ``move`` to ``destination`` on ``grid``.

    The foe idles far to the right (never threatens). Returns the hero state.
    """
    cols = grid.width
    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]), grid, [hero, foe], policy=policy, round_cap=1,
    )
    eng.run()
    return eng._states[0]


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)


# ── 3D flight over ground terrain (airborne creatures) ────────────────────────
#
# SRD "Movement, Position, and Distance" (Terrain and Obstacles): "Flying and
# incorporeal creatures are not hampered by difficult terrain" and "Flying and
# incorporeal creatures are able to avoid most obstacles." An airborne creature
# (fly >= land, ``_is_airborne``) moves on ``Grid.airborne_view``: ground
# difficult terrain costs 1 and cells impassable to walkers (pits) are
# crossable, EXCEPT cells that block line of effect (solid barriers, assumed
# full height) which still stop flight. Walkers are bit-identical. The foe sits
# two rows below the hero's lane so it never threatens the row being crossed.

_LEGEND = {
    ".": TerrainType(type="floor", move_cost=1),
    "T": TerrainType(type="rubble", move_cost=2),
    "P": TerrainType(type="pit", move_cost=None),
    "#": TerrainType(type="wall", move_cost=None, blocks_los=True),
}


def _lane_grid(lane: str) -> Grid:
    """A 3-row grid: the terrain ``lane`` on row 0, floor rows below."""
    width = len(lane)
    rows = (lane, "." * width, "." * width)
    spec = MapSpec(name="lane", terrain=rows, legend=_LEGEND)
    return Grid.from_spec(spec)


def _move_on_lane(*, land: int, fly: int | None, lane: str,
                  destination: tuple[int, int]) -> tuple[int, int]:
    """Hero (row 0, col 0) acts first and takes one ``move`` to ``destination``.

    The foe idles at (2, 0), two rows below and never threatening. Returns the
    hero's final position.
    """
    hero = make_state(
        _combatant("hero", land=land, fly=fly, initiative_mod=2), "players", (0, 0),
    )
    foe = make_state(_combatant("foe", land=30), "monsters", (2, 0))

    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]), _lane_grid(lane), [hero, foe],
        policy=policy, round_cap=1,
    )
    eng.run()
    return eng._states[0].pos


def test_flyer_crosses_rubble_at_normal_cost() -> None:
    # Rubble at cols 4-5: the flyer prices every cell at 1, so (0, 9) costs
    # 9 cells and is within the 16-cell (fly 80) budget.
    assert _move_on_lane(land=30, fly=80, lane="....TT....", destination=(0, 9)) == (0, 9)


def test_walker_still_pays_for_rubble() -> None:
    # Same lane, same 10-cell span: a land-30 walker pays 3 + 2*2 + 4 = 11
    # cells for (0, 9), beyond its 6-cell budget, so it is rejected in place.
    assert _move_on_lane(land=30, fly=None, lane="....TT....", destination=(0, 9)) == (0, 0)


def test_flyer_crosses_pit_impassable_to_walkers() -> None:
    # A pit (move_cost None, no blocks_los) at col 3: the flyer avoids the
    # obstacle and reaches (0, 7); the walker cannot enter the pit cell.
    assert _move_on_lane(land=30, fly=80, lane="...P....", destination=(0, 7)) == (0, 7)
    assert _move_on_lane(land=30, fly=None, lane="...P....", destination=(0, 7)) == (0, 0)


def test_flyer_still_stopped_by_wall() -> None:
    # A wall column (blocks_los) is a solid barrier assumed full height: it
    # stops flight just like walking — no crossing and no detour.
    hero = make_state(
        _combatant("hero", land=30, fly=80, initiative_mod=2), "players", (0, 0),
    )
    foe = make_state(_combatant("foe", land=30), "monsters", (2, 7))

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

    spec = MapSpec(
        name="wall", terrain=("...#....", "...#....", "...#...."), legend=_LEGEND,
    )
    eng = CombatEngine(
        ScriptedRng([10, 5]), Grid.from_spec(spec), [hero, foe],
        policy=policy, round_cap=1,
    )
    eng.run()
    assert eng._states[0].pos == (0, 0)


def test_flyer_runs_over_rubble_walker_run_stops() -> None:
    # Run is a straight line: an airborne runner is not hampered by the
    # rubble; a walker's run stops in front of it (CRB "Difficult Terrain":
    # "You can't run or charge across difficult terrain").
    def _run(land: int, fly: int | None) -> tuple[int, int]:
        hero = make_state(
            _combatant("hero", land=land, fly=fly, initiative_mod=2), "players", (0, 0),
        )
        foe = make_state(_combatant("foe", land=30), "monsters", (2, 0))

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

        eng = CombatEngine(
            ScriptedRng([10, 5]), _lane_grid("....TT...."),
            [hero, foe], policy=policy, round_cap=1,
        )
        eng.run()
        return eng._states[0].pos

    assert _run(30, 80) == (0, 9)
    # The walker's run destination lies beyond difficult terrain: the
    # reachability verdict and the walked line agree, so the run is rejected
    # outright (unreachable_destination) and the walker never starts.
    assert _run(30, None) == (0, 0)


def test_flyer_5foot_step_into_rubble() -> None:
    # SRD "Take 5-Foot Step": the step is legal only if movement isn't
    # hampered by difficult terrain — a flyer's isn't, a walker's is.
    def _step(land: int, fly: int | None) -> tuple[int, int]:
        hero = make_state(
            _combatant("hero", land=land, fly=fly, initiative_mod=2), "players", (0, 3),
        )
        foe = make_state(_combatant("foe", land=30), "monsters", (2, 0))

        def policy(_e: CombatEngine, s: CombatantState) -> tuple[Action, ...]:
            if s.combatant.id == "hero":
                return (Action(kind="5foot_step", destination=(0, 4)),)
            return _WAIT

        eng = CombatEngine(
            ScriptedRng([10, 5]), _lane_grid("...T...."),
            [hero, foe], policy=policy, round_cap=1,
        )
        eng.run()
        return eng._states[0].pos

    assert _step(30, 80) == (0, 4)
    assert _step(30, None) == (0, 3)


def test_movement_cost_ft_uses_the_air_view() -> None:
    # Cost from the hero's cell to the target's anchor: rubble prices at 1 for
    # the flyer (8 cells = 40 ft) and at 2 for the walker (13 cells = 65 ft,
    # beyond its 30 ft speed so None).
    grid = _lane_grid("..TTTTT..")
    flyer = make_state(_combatant("flyer", land=30, fly=80), "players", (0, 0))
    walker = make_state(_combatant("walker", land=30), "players", (0, 0))
    target_foe = make_state(_combatant("target", land=30), "monsters", (0, 8))
    eng = CombatEngine(ScriptedRng([]), grid, [flyer, target_foe])
    assert eng.movement_cost_ft(flyer, target_foe) == 40
    target_walker = make_state(_combatant("target", land=30), "monsters", (0, 8))
    eng2 = CombatEngine(ScriptedRng([]), grid, [walker, target_walker])
    assert eng2.movement_cost_ft(walker, target_walker) is None


def test_fly_equal_to_land_is_airborne() -> None:
    # fly >= land: the budget comes from the fly speed, so the creature flies
    # (6-cell budget, rubble priced at 1, (0, 6) reachable).
    assert _move_on_lane(land=30, fly=30, lane="....TT....", destination=(0, 6)) == (0, 6)


def test_fly_slower_than_land_walks_and_pays() -> None:
    # fly < land: the creature walks (budget = land) and pays the rubble cost
    # like any walker — (0, 9) costs 11 cells > the 6-cell budget.
    assert _move_on_lane(land=30, fly=10, lane="....TT....", destination=(0, 9)) == (0, 0)


def test_flyer_still_pays_elevation_climb() -> None:
    # The elevation layer stores STEP heights, not wall heights, so the
    # climbing cost still applies to flight (documented simplification): with
    # a 2-cell budget (fly 10 = land 10) the step up onto the higher ledge
    # costs 2 and only the flat neighbour is reachable.
    legend = {".": TerrainType(type="floor", move_cost=1)}
    spec = MapSpec(
        name="ledge", terrain=("....", "...."), legend=legend,
        elevation=("..1.", "...."),
    )
    grid = Grid.from_spec(spec)

    def _go(destination: tuple[int, int]) -> tuple[int, int]:
        hero = make_state(
            _combatant("hero", land=10, fly=10, initiative_mod=2), "players", (0, 0),
        )
        foe = make_state(_combatant("foe", land=30), "monsters", (1, 3))

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

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

    assert _go((0, 2)) == (0, 0)
    assert _go((0, 1)) == (0, 1)


def test_airborne_view_preserves_combat_flags() -> None:
    # The air view normalizes only the walking cost: solid barriers stay
    # impassable (so LoE/cover geometry is untouched by flight), cover flags
    # and the elevation layer are preserved.
    legend = dict(_LEGEND)
    legend["C"] = TerrainType(type="pillar", move_cost=None, cover=True)
    spec = MapSpec(
        name="flags", terrain=(".TP#C",), legend=legend, elevation=(".1...",),
    )
    grid = Grid.from_spec(spec)
    view = grid.airborne_view()
    wall = view.terrain_at((0, 3))
    assert wall.blocks_los
    assert wall.move_cost is None
    pit = view.terrain_at((0, 2))
    assert pit.move_cost == 1
    assert not pit.blocks_los
    assert view.terrain_at((0, 1)).move_cost == 1
    pillar = view.terrain_at((0, 4))
    assert pillar.move_cost == 1
    assert pillar.cover
    assert view.elevation_at((0, 1)) == grid.elevation_at((0, 1)) == 1