from __future__ import annotations

import math
from collections.abc import Callable, Sequence
from typing import TYPE_CHECKING

import msgspec

from grim.color import RGBA
from grim.geom import Vec2

from ..math_parity import NATIVE_TAU, f32, heading_from_delta_f32
from ..perks import PerkId
from ..perks.helpers import perk_active
from ..projectiles.runtime import SecondarySpawnSpec
from ..projectiles.types import ProjectileTemplateId, SecondaryProjectileTypeId
from ..sim.input import PlayerInput
from ..sim.state_types import GameplayState, PlayerState
from ..weapons import WEAPON_TABLE, WeaponId, weapon_entry_for_projectile_type_id
from .assign import player_start_reload, weapon_entry
from .fire_recipes import (
    MaskCenteredJitter,
    ModuloCenteredJitter,
    ModuloSpeedScale,
    MultiPlasmaFanMode,
    NoJitter,
    NoSpeedScale,
    ParticleStreamMode,
    PrimaryPelletsMode,
    SecondaryShotMode,
    SwarmerDumpMode,
    UseAimTargetHint,
    resolve_fire_recipe,
)
from .spawn import owner_ref_for_player, owner_ref_for_player_projectiles, travel_budget_for_type_id

if TYPE_CHECKING:
    from ..creatures.runtime import CreatureState

WEAPON_COUNT_SIZE = max(int(entry.weapon_id) for entry in WEAPON_TABLE) + 1

_NATIVE_FIRE_MUZZLE_SPRITES: dict[int, tuple[tuple[float, float, float], ...]] = {
    WeaponId.PISTOL: ((25.0, 1.0, 0.23), (15.0, 2.0, 0.213)),
    WeaponId.ASSAULT_RIFLE: ((25.0, 1.0, 0.23), (15.0, 2.0, 0.213)),
    WeaponId.SHOTGUN: ((25.0, 1.0, 0.25), (15.0, 2.0, 0.223)),
    WeaponId.SAWED_OFF_SHOTGUN: ((25.0, 1.0, 0.26), (15.0, 2.0, 0.233)),
    WeaponId.SUBMACHINE_GUN: ((25.0, 1.0, 0.23), (15.0, 2.0, 0.213)),
    WeaponId.GAUSS_GUN: ((25.0, 1.0, 0.33), (15.0, 2.0, 0.263)),
    WeaponId.ROCKET_LAUNCHER: ((25.0, 1.0, 0.34), (15.0, 2.0, 0.283)),
    WeaponId.SEEKER_ROCKETS: ((25.0, 1.0, 0.31), (15.0, 2.0, 0.243)),
    WeaponId.MINI_ROCKET_SWARMERS: ((25.0, 1.0, 0.34), (15.0, 2.0, 0.283)),
    WeaponId.ROCKET_MINIGUN: ((25.0, 1.0, 0.34),),
    WeaponId.JACKHAMMER: ((15.0, 2.0, 0.223),),
    WeaponId.SHRINKIFIER_5K: ((25.0, 1.0, 0.23), (15.0, 2.0, 0.213)),
    WeaponId.GAUSS_SHOTGUN: ((25.0, 1.0, 0.33), (15.0, 2.0, 0.263)),
}

_NATIVE_FIRE_MUZZLE_AFTER_PROJECTILE: frozenset[int] = frozenset(
    {
        WeaponId.PISTOL,
        WeaponId.SHRINKIFIER_5K,
    },
)


class WeaponFireCtx(msgspec.Struct):
    player: PlayerState
    input_state: PlayerInput
    dt: float
    state: GameplayState
    detail_preset: int = 5
    creatures: Sequence[CreatureState] | None = None
    players: Sequence[PlayerState] | None = None
    force_pre_swap_fire_gate: bool = False
    on_player_lethal: Callable[[PlayerState], None] | None = None


class WeaponFireResult(msgspec.Struct, frozen=True):
    fired: bool
    shot_count: int = 0
    ammo_cost: float = 0.0


def _spawn_native_fire_muzzle_sprites(
    *,
    state: GameplayState,
    weapon_id: int,
    muzzle: Vec2,
    aim_heading: float,
    fire_bullets_active: bool,
) -> None:
    if fire_bullets_active:
        specs: tuple[tuple[float, float, float], ...] = ((25.0, 1.0, 0.413),)
    else:
        specs = _NATIVE_FIRE_MUZZLE_SPRITES.get(int(weapon_id), ())
    if not specs:
        return

    for speed, scale, alpha in specs:
        state.sprite_effects.spawn(
            pos=muzzle,
            vel=Vec2.from_heading(aim_heading) * float(speed),
            scale=float(scale),
            color=RGBA(0.5, 0.5, 0.5, float(alpha)),
        )


def _native_shot_angle_with_jitter(
    *,
    aim: Vec2,
    player_pos: Vec2,
    spread_heat: float,
    rand: Callable[[], int],
) -> float:
    # `player_fire_weapon` computes jitter in float locals before `to_heading`.
    aim_dx = float(f32(float(aim.x) - float(player_pos.x)))
    aim_dy = float(f32(float(aim.y) - float(player_pos.y)))
    dist_sq = float(f32(float(f32(float(aim_dx) * float(aim_dx))) + float(f32(float(aim_dy) * float(aim_dy)))))
    dist = float(f32(math.sqrt(float(dist_sq))))
    max_offset = float(f32(float(f32(float(dist) * float(spread_heat))) * 0.5))

    dir_angle = float(f32(float(int(rand()) & 0x1FF) * (float(NATIVE_TAU) / 512.0)))
    mag = float(f32(float(int(rand()) & 0x1FF) * (1.0 / 512.0)))
    offset = float(f32(float(max_offset) * float(mag)))

    dir_x = float(f32(math.cos(float(dir_angle))))
    dir_y = float(f32(math.sin(float(dir_angle))))
    aim_jitter_x = float(f32(float(aim.x) + float(f32(float(dir_x) * float(offset)))))
    aim_jitter_y = float(f32(float(aim.y) + float(f32(float(dir_y) * float(offset)))))

    shot_dx = float(f32(float(aim_jitter_x) - float(player_pos.x)))
    shot_dy = float(f32(float(aim_jitter_y) - float(player_pos.y)))
    return float(heading_from_delta_f32(dx=float(shot_dx), dy=float(shot_dy)))


def _apply_pellet_jitter(
    *,
    shot_angle: float,
    rand: Callable[[], int],
    jitter_rule: NoJitter | ModuloCenteredJitter | MaskCenteredJitter,
) -> float:
    match jitter_rule:
        case NoJitter():
            return float(shot_angle)
        case ModuloCenteredJitter(modulo=modulo, center=center, step=step):
            return float(shot_angle) + float(int(rand()) % int(modulo) - int(center)) * float(step)
        case MaskCenteredJitter(mask=mask, center=center, step=step):
            return float(shot_angle) + float((int(rand()) & int(mask)) - int(center)) * float(step)


def _apply_speed_scale_rule(
    *,
    state: GameplayState,
    proj_id: int,
    speed_rule: NoSpeedScale | ModuloSpeedScale,
) -> None:
    match speed_rule:
        case NoSpeedScale():
            return
        case ModuloSpeedScale(base=base, modulo=modulo, step=step):
            state.projectiles.entries[int(proj_id)].speed_scale = float(base) + float(int(state.rng.rand()) % int(modulo)) * float(step)


def fire_weapon(ctx: WeaponFireCtx) -> WeaponFireResult:
    player = ctx.player
    input_state = ctx.input_state
    dt = float(ctx.dt)
    state = ctx.state
    detail_preset = int(ctx.detail_preset)
    creatures = ctx.creatures
    players = ctx.players
    force_pre_swap_fire_gate = bool(ctx.force_pre_swap_fire_gate)
    on_player_lethal = ctx.on_player_lethal

    weapon_id = player.weapon.weapon_id
    weapon = weapon_entry(weapon_id)

    if (not force_pre_swap_fire_gate) and player.weapon.shot_cooldown > 0.0:
        return WeaponFireResult(fired=False)
    if not input_state.fire_down:
        return WeaponFireResult(fired=False)

    ammo_cost = 1.0
    is_fire_bullets = float(player.fire_bullets_timer) > 0.0
    if (not force_pre_swap_fire_gate) and player.weapon.reload_timer > 0.0:
        if player.experience <= 0:
            return WeaponFireResult(fired=False)
        if perk_active(player, PerkId.REGRESSION_BULLETS):
            ammo_class = int(weapon.ammo_class) if weapon.ammo_class is not None else 0

            reload_time = float(weapon.reload_time)
            factor = 4.0 if ammo_class == 1 else 200.0
            player.experience = int(float(player.experience) - reload_time * factor)
            if player.experience < 0:
                player.experience = 0
        elif perk_active(player, PerkId.AMMUNITION_WITHIN):
            ammo_class = int(weapon.ammo_class) if weapon.ammo_class is not None else 0

            from ..player_damage import player_take_damage

            cost = 0.15 if ammo_class == 1 else 1.0
            player_take_damage(
                state,
                player,
                cost,
                dt=dt,
                rand=state.rng.rand,
                players=players,
                on_lethal=(lambda: on_player_lethal(player)) if on_player_lethal is not None else None,
            )
        else:
            return WeaponFireResult(fired=False)

    pellet_count = int(weapon.pellet_count)
    fire_bullets_weapon = weapon_entry_for_projectile_type_id(ProjectileTemplateId.FIRE_BULLETS)

    shot_cooldown = float(f32(float(weapon.shot_cooldown)))
    weapon_spread_heat = float(weapon.spread_heat_inc)
    fire_bullets_spread_heat = float(fire_bullets_weapon.spread_heat_inc)

    if is_fire_bullets and pellet_count == 1:
        shot_cooldown = float(f32(float(fire_bullets_weapon.shot_cooldown)))

    spread_heat_base = fire_bullets_spread_heat if is_fire_bullets else weapon_spread_heat
    spread_inc = spread_heat_base * 1.3

    if perk_active(player, PerkId.FASTSHOT):
        shot_cooldown = float(f32(float(shot_cooldown) * 0.88))
    if perk_active(player, PerkId.SHARPSHOOTER):
        shot_cooldown = float(f32(float(shot_cooldown) * 1.05))
    player.weapon.shot_cooldown = max(0.0, float(f32(float(shot_cooldown))))

    aim = input_state.aim
    aim_delta = aim - player.pos
    aim_heading = float(heading_from_delta_f32(dx=float(aim_delta.x), dy=float(aim_delta.y)))

    muzzle = player.pos + Vec2.from_heading(aim_heading).rotated(-0.150915) * 16.0
    state.effects.spawn_shell_casing(
        pos=muzzle,
        aim_heading=aim_heading,
        weapon_flags=int(weapon.flags or 0),
        rand=state.rng.rand,
        detail_preset=detail_preset,
    )

    shot_angle = _native_shot_angle_with_jitter(
        aim=aim,
        player_pos=player.pos,
        spread_heat=float(player.spread_heat),
        rand=state.rng.rand,
    )
    particle_angle = Vec2.from_heading(shot_angle).to_angle()
    if weapon_id in (WeaponId.FLAMETHROWER, WeaponId.BLOW_TORCH, WeaponId.HR_FLAMER):
        particle_angle = Vec2.from_heading(aim_heading).to_angle()

    # Native `player_fire_weapon` consumes one RNG draw for shot SFX variant
    # selection on every non-Fire-Bullets shot.
    if not is_fire_bullets:
        state.rng.rand()

    owner = owner_ref_for_player(player.index)
    projectile_owner = owner_ref_for_player_projectiles(state, player.index)
    shot_count = 1
    spawn_muzzle_after_projectile = bool(is_fire_bullets) or int(weapon_id) in _NATIVE_FIRE_MUZZLE_AFTER_PROJECTILE
    if not spawn_muzzle_after_projectile:
        _spawn_native_fire_muzzle_sprites(
            state=state,
            weapon_id=int(weapon_id),
            muzzle=muzzle,
            aim_heading=float(aim_heading),
            fire_bullets_active=bool(is_fire_bullets),
        )

    recipe = resolve_fire_recipe(
        weapon_id=weapon_id,
        pellet_count=pellet_count,
        fire_bullets_active=is_fire_bullets,
    )
    ammo_cost = float(recipe.ammo_cost)

    match recipe.mode:
        case PrimaryPelletsMode(type_id=type_id, count=count, jitter=jitter_rule, speed_scale=speed_rule):
            if type_id is None:
                raise ValueError(f"missing projectile type in recipe for weapon {int(weapon_id)}")
            pellets = max(0, int(count if count is not None else 0))
            shot_count = pellets
            meta = travel_budget_for_type_id(type_id)
            for _ in range(pellets):
                angle = _apply_pellet_jitter(
                    shot_angle=float(shot_angle),
                    rand=state.rng.rand,
                    jitter_rule=jitter_rule,
                )
                proj_id = state.projectiles.spawn(
                    pos=muzzle,
                    angle=angle,
                    type_id=type_id,
                    owner=projectile_owner,
                    travel_budget=meta,
                )
                _apply_speed_scale_rule(
                    state=state,
                    proj_id=int(proj_id),
                    speed_rule=speed_rule,
                )
        case SecondaryShotMode(type_id=type_id, targeting=targeting):
            target_hint = None
            spawn_creatures = None
            match targeting:
                case UseAimTargetHint():
                    target_hint = aim
                    spawn_creatures = creatures
                case _:
                    pass
            state.secondary_projectiles.spawn_from_spec(
                SecondarySpawnSpec(
                    pos=muzzle,
                    angle=shot_angle,
                    type_id=type_id,
                    owner=owner,
                    target_hint=target_hint,
                    creatures=spawn_creatures,
                ),
            )
        case ParticleStreamMode(style=style, slow=slow):
            if slow:
                state.particles.spawn_particle_slow(
                    pos=muzzle,
                    angle=Vec2.from_heading(shot_angle).to_angle(),
                    owner=owner,
                )
            else:
                particle_id = state.particles.spawn_particle(
                    pos=muzzle,
                    angle=particle_angle,
                    intensity=1.0,
                    owner=owner,
                )
                if style is not None:
                    state.particles.entries[particle_id].style_id = style
        case MultiPlasmaFanMode():
            # Multi-Plasma: 5-shot fixed spread using type 0x09 and 0x0B.
            shot_count = 5
            spread_small = math.pi / 10
            spread_large = math.pi / 6
            patterns: tuple[tuple[float, ProjectileTemplateId], ...] = (
                (-spread_small, ProjectileTemplateId.PLASMA_RIFLE),
                (-spread_large, ProjectileTemplateId.PLASMA_MINIGUN),
                (0.0, ProjectileTemplateId.PLASMA_RIFLE),
                (spread_large, ProjectileTemplateId.PLASMA_MINIGUN),
                (spread_small, ProjectileTemplateId.PLASMA_RIFLE),
            )
            for angle_offset, type_id in patterns:
                state.projectiles.spawn(
                    pos=muzzle,
                    angle=shot_angle + angle_offset,
                    type_id=type_id,
                    owner=projectile_owner,
                    travel_budget=travel_budget_for_type_id(type_id),
                )
        case SwarmerDumpMode():
            # Mini-Rocket Swarmers -> secondary type 2 (fires the full clip in a spread).
            rocket_count = max(1, int(player.weapon.ammo))
            if bool(state.preserve_bugs):
                # Native bug: step scales by ammo (`ammo * pi/3`), which aliases to identical headings
                # for some clip sizes (e.g. 6 rockets), causing visible clumping.
                step = float(rocket_count) * (math.pi / 3.0)
                angle = (shot_angle - math.pi) - step * float(rocket_count) * 0.5
            else:
                spread = math.pi * (2.0 / 3.0)
                step = 0.0 if rocket_count <= 1 else spread / float(rocket_count - 1)
                angle = shot_angle - spread * 0.5
            for _ in range(rocket_count):
                state.secondary_projectiles.spawn_from_spec(
                    SecondarySpawnSpec(
                        pos=muzzle,
                        angle=angle,
                        type_id=SecondaryProjectileTypeId.HOMING_ROCKET,
                        owner=owner,
                        target_hint=aim,
                        creatures=creatures,
                    ),
                )
                angle += step
            ammo_cost = float(rocket_count)
            shot_count = rocket_count

    if 0 <= int(player.index) < len(state.shots_fired):
        state.shots_fired[int(player.index)] += int(shot_count)
        if 0 <= weapon_id < WEAPON_COUNT_SIZE:
            state.weapon_shots_fired[int(player.index)][weapon_id] += int(shot_count)

    if spawn_muzzle_after_projectile:
        _spawn_native_fire_muzzle_sprites(
            state=state,
            weapon_id=int(weapon_id),
            muzzle=muzzle,
            aim_heading=float(aim_heading),
            fire_bullets_active=bool(is_fire_bullets),
        )

    if not perk_active(player, PerkId.SHARPSHOOTER):
        player.spread_heat = min(0.48, max(0.0, player.spread_heat + spread_inc))

    muzzle_inc = weapon_spread_heat
    if is_fire_bullets and pellet_count == 1:
        muzzle_inc = fire_bullets_spread_heat
    player.muzzle_flash_alpha = min(1.0, player.muzzle_flash_alpha)
    player.muzzle_flash_alpha = min(1.0, player.muzzle_flash_alpha + muzzle_inc)
    player.muzzle_flash_alpha = min(0.8, player.muzzle_flash_alpha)

    player.shot_seq += 1
    if state.bonuses.reflex_boost <= 0.0 and not is_fire_bullets:
        # Native allows ammo to cross below zero for reload-time firing paths
        # (for example Regression Bullets), and replay checkpoints rely on that.
        player.weapon.ammo = float(player.weapon.ammo) - float(ammo_cost)
    reload_start_gate_open = bool(player.weapon.reload_timer <= 0.0)
    if force_pre_swap_fire_gate:
        # Alt-weapon same-tick fire uses the pre-swap gate (reload_timer==0) for
        # reload restart eligibility after ammo drains below zero.
        reload_start_gate_open = True
    if player.weapon.ammo <= 0.0 and reload_start_gate_open:
        player_start_reload(player, state)
    return WeaponFireResult(fired=True, shot_count=int(shot_count), ammo_cost=float(ammo_cost))
