from __future__ import annotations

import math

from grim.color import RGBA
from grim.geom import Vec2
from grim.math import clamp
from grim.raylib_api import rd, rl

from ...projectiles.types import ProjectileTemplateId
from ...sim.world_defs import KNOWN_PROJ_FRAMES
from .common import proj_origin
from .types import ProjectileDrawCtx


def _begin_darken_src_zero_blend() -> None:
    # Native projectile_render switches Plague Spreader to D3D8 SRC=ZERO / DST=INVSRCALPHA.
    rl.rl_set_blend_factors_separate(
        rd.RL_ZERO,
        rd.RL_ONE_MINUS_SRC_ALPHA,
        rd.RL_ZERO,
        rd.RL_ONE,
        rd.RL_FUNC_ADD,
        rd.RL_FUNC_ADD,
    )
    rl.begin_blend_mode(rl.BlendMode.BLEND_CUSTOM_SEPARATE)
    rl.rl_set_blend_factors_separate(
        rd.RL_ZERO,
        rd.RL_ONE_MINUS_SRC_ALPHA,
        rd.RL_ZERO,
        rd.RL_ONE,
        rd.RL_FUNC_ADD,
        rd.RL_FUNC_ADD,
    )


def draw_pulse_gun(ctx: ProjectileDrawCtx) -> bool:
    renderer = ctx.renderer
    if ctx.type_id != ProjectileTemplateId.PULSE_GUN:
        return False
    if ctx.texture is None:
        return False

    mapping = KNOWN_PROJ_FRAMES.get(ctx.type_id)
    if mapping is None:
        return True
    grid, frame = mapping
    cell_w = float(ctx.texture.width) / float(grid)

    alpha = float(ctx.alpha)
    life = float(ctx.life)
    if life >= 0.4:
        origin = proj_origin(ctx.proj, ctx.pos)
        dist = origin.distance_to(ctx.pos)

        desired_size = dist * 0.16 * ctx.scale
        if desired_size <= 1e-3:
            return True
        sprite_scale = desired_size / cell_w if cell_w > 1e-6 else 0.0
        if sprite_scale <= 1e-6:
            return True

        tint = RGBA(0.1, 0.6, 0.2, alpha * 0.7).to_rl()
        rl.begin_blend_mode(rl.BlendMode.BLEND_ADDITIVE)
        renderer._draw_atlas_sprite(
            ctx.texture,
            grid=grid,
            frame=frame,
            pos=ctx.screen_pos,
            scale=sprite_scale,
            rotation_rad=ctx.angle,
            tint=tint,
        )
        rl.end_blend_mode()
        return True

    fade = clamp(life * 2.5, 0.0, 1.0)
    fade_alpha = fade * alpha
    if fade_alpha <= 1e-3:
        return True

    desired_size = 56.0 * ctx.scale
    sprite_scale = desired_size / cell_w if cell_w > 1e-6 else 0.0
    if sprite_scale <= 1e-6:
        return True

    tint = RGBA(1.0, 1.0, 1.0, fade_alpha).to_rl()
    rl.begin_blend_mode(rl.BlendMode.BLEND_ADDITIVE)
    renderer._draw_atlas_sprite(
        ctx.texture,
        grid=grid,
        frame=frame,
        pos=ctx.screen_pos,
        scale=sprite_scale,
        rotation_rad=ctx.angle,
        tint=tint,
    )
    rl.end_blend_mode()
    return True


def draw_splitter_or_blade(ctx: ProjectileDrawCtx) -> bool:
    renderer = ctx.renderer
    type_id = ctx.type_id
    if type_id not in (ProjectileTemplateId.SPLITTER_GUN, ProjectileTemplateId.BLADE_GUN):
        return False
    if ctx.texture is None:
        return False

    mapping = KNOWN_PROJ_FRAMES.get(type_id)
    if mapping is None:
        return True
    grid, frame = mapping
    cell_w = float(ctx.texture.width) / float(grid)

    if float(ctx.life) < 0.4:
        return True

    origin = proj_origin(ctx.proj, ctx.pos)
    dist = origin.distance_to(ctx.pos)

    desired_size = min(dist, 20.0) * ctx.scale
    if desired_size <= 1e-3:
        return True

    sprite_scale = desired_size / cell_w if cell_w > 1e-6 else 0.0
    if sprite_scale <= 1e-6:
        return True

    rotation_rad = ctx.angle
    rgb = (1.0, 1.0, 1.0)
    if type_id == ProjectileTemplateId.BLADE_GUN:
        rotation_rad = float(int(ctx.proj_index)) * 0.1 - float(renderer.frame.elapsed_ms) * 0.1
        rgb = (0.8, 0.8, 0.8)

    tint = RGBA(rgb[0], rgb[1], rgb[2], float(ctx.alpha)).to_rl()
    renderer._draw_atlas_sprite(
        ctx.texture,
        grid=grid,
        frame=frame,
        pos=ctx.screen_pos,
        scale=sprite_scale,
        rotation_rad=rotation_rad,
        tint=tint,
    )
    return True


def draw_plague_spreader(ctx: ProjectileDrawCtx) -> bool:
    renderer = ctx.renderer
    if ctx.type_id != ProjectileTemplateId.PLAGUE_SPREADER:
        return False
    texture = ctx.texture
    if texture is None:
        return False

    grid = 4
    frame = 2
    cell_w = float(texture.width) / float(grid)

    alpha = float(ctx.alpha)
    life = float(ctx.life)
    if life >= 0.4:
        tint = RGBA(1.0, 1.0, 1.0, alpha).to_rl()

        def draw_plague_quad(*, pos: Vec2, size: float) -> None:
            size = float(size)
            if size <= 1e-3:
                return
            desired_size = size * ctx.scale
            sprite_scale = desired_size / cell_w if cell_w > 1e-6 else 0.0
            if sprite_scale <= 1e-6:
                return
            pos_screen = renderer.world_to_screen(pos)
            renderer._draw_atlas_sprite(
                texture,
                grid=grid,
                frame=frame,
                pos=pos_screen,
                scale=sprite_scale,
                rotation_rad=0.0,
                tint=tint,
            )

        _begin_darken_src_zero_blend()
        try:
            draw_plague_quad(pos=ctx.pos, size=60.0)

            offset = Vec2.from_heading(ctx.angle + math.pi) * 15.0
            draw_plague_quad(
                pos=ctx.pos + offset,
                size=60.0,
            )

            phase = float(int(ctx.proj_index)) + float(renderer.frame.elapsed_ms) * 0.01
            cos_phase = math.cos(phase)
            sin_phase = math.sin(phase)
            draw_plague_quad(
                pos=ctx.pos.offset(dx=cos_phase * cos_phase - 5.0, dy=sin_phase * 11.0 - 5.0),
                size=52.0,
            )

            phase_120 = phase + 2.0943952
            sin_phase_120 = math.sin(phase_120)
            draw_plague_quad(
                pos=ctx.pos + Vec2.from_polar(phase_120, 10.0),
                size=62.0,
            )

            phase_240 = phase + 4.1887903
            draw_plague_quad(
                pos=ctx.pos + Vec2(math.cos(phase_240) * 10.0, math.sin(phase_240) * sin_phase_120),
                size=62.0,
            )
        finally:
            rl.end_blend_mode()
        return True

    fade = clamp(life * 2.5, 0.0, 1.0)
    fade_alpha = fade * alpha
    if fade_alpha <= 1e-3:
        return True

    desired_size = (fade * 40.0 + 32.0) * ctx.scale
    sprite_scale = desired_size / cell_w if cell_w > 1e-6 else 0.0
    if sprite_scale <= 1e-6:
        return True

    tint = RGBA(1.0, 1.0, 1.0, fade_alpha).to_rl()
    _begin_darken_src_zero_blend()
    try:
        renderer._draw_atlas_sprite(
            texture,
            grid=grid,
            frame=frame,
            pos=ctx.screen_pos,
            scale=sprite_scale,
            rotation_rad=0.0,
            tint=tint,
        )
    finally:
        rl.end_blend_mode()
    return True


__all__ = ["draw_plague_spreader", "draw_pulse_gun", "draw_splitter_or_blade"]
