from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Sequence

import pyray as rl

from grim.fonts.small import SmallFontData, measure_small_text_width
from grim.math import clamp

from ...gameplay import (
    GameplayState,
    PerkSelectionState,
    PlayerState,
    _CreatureForPerks,
    perk_selection_current_choices,
    perk_selection_pick,
)
from ...perks import PerkId, perk_display_description, perk_display_name
from ...ui.menu_panel import draw_classic_menu_panel
from ...ui.perk_menu import (
    PERK_MENU_TRANSITION_MS,
    PerkMenuAssets,
    PerkMenuLayout,
    UiButtonState,
    button_draw,
    button_update,
    button_width,
    draw_menu_item,
    draw_ui_text,
    menu_item_hit_rect,
    perk_menu_compute_layout,
    perk_menu_panel_slide_x,
)
from ...ui.layout import ui_origin, ui_scale

PlaySfxFn = Callable[[str], None]
OnCloseFn = Callable[[], None]
OnPickFn = Callable[[int], None]

UI_TEXT_COLOR = rl.Color(220, 220, 220, 255)
UI_SPONSOR_COLOR = rl.Color(255, 255, 255, int(255 * 0.5))


@dataclass(frozen=True, slots=True)
class PerkMenuContext:
    state: GameplayState
    perk_state: PerkSelectionState
    players: list[PlayerState]
    creatures: Sequence[_CreatureForPerks]
    player: PlayerState
    game_mode: int
    player_count: int
    fx_toggle: int

    font: SmallFontData | None
    assets: PerkMenuAssets | None
    mouse: rl.Vector2
    fx_detail: bool = False
    play_sfx: PlaySfxFn | None = None


class PerkMenuController:
    _DESC_WRAP_WIDTH_PX = 256.0

    def __init__(
        self,
        *,
        cancel_label: str = "Cancel",
        on_close: OnCloseFn | None = None,
        on_pick: OnPickFn | None = None,
    ) -> None:
        self._cancel_label = cancel_label
        self._on_close = on_close
        self._on_pick = on_pick
        self.reset()

    @property
    def open(self) -> bool:
        return bool(self._open)

    @open.setter
    def open(self, value: bool) -> None:
        if not value and self._open:
            self.close()
        else:
            self._open = bool(value)

    @property
    def selected_index(self) -> int:
        return int(self._selected_index)

    @selected_index.setter
    def selected_index(self, value: int) -> None:
        self._selected_index = int(value)

    @property
    def timeline_ms(self) -> float:
        return float(self._timeline_ms)

    @timeline_ms.setter
    def timeline_ms(self, value: float) -> None:
        self._timeline_ms = float(value)

    @property
    def active(self) -> bool:
        return bool(self._open) or self._timeline_ms > 1e-3

    def reset(self) -> None:
        self._layout = PerkMenuLayout()
        self._cancel_button = UiButtonState(self._cancel_label)
        self._open = False
        self._selected_index = 0
        self._timeline_ms = 0.0
        self._wrapped_desc_cache: dict[tuple[int, int], str] = {}

    def _prewrapped_perk_desc(self, perk_id: int, font: SmallFontData, *, fx_toggle: int) -> str:
        key = (int(perk_id), int(fx_toggle))
        cached = self._wrapped_desc_cache.get(key)
        if cached is not None:
            return cached
        desc = perk_display_description(int(perk_id), fx_toggle=int(fx_toggle))
        wrapped = self._wrap_small_text_native(
            font,
            desc,
            max_width_px=self._DESC_WRAP_WIDTH_PX,
            scale=1.0,
        )
        self._wrapped_desc_cache[key] = wrapped
        return wrapped

    @staticmethod
    def _wrap_small_text_native(font: SmallFontData, text: str, max_width_px: float, *, scale: float) -> str:
        wrapped = list(str(text))
        if not wrapped:
            return ""

        max_width = float(max_width_px)
        remaining = max_width
        i = 0
        while i < len(wrapped):
            ch = wrapped[i]
            if ch == "\r":
                i += 1
                continue
            if ch == "\n":
                remaining = max_width
                i += 1
                continue

            remaining -= measure_small_text_width(font, ch, float(scale))
            if remaining < 0.0:
                j = i
                while j > 0 and wrapped[j] not in {" ", "\n"}:
                    j -= 1
                if wrapped[j] == " ":
                    wrapped[j] = "\n"
                    i = j
                remaining = max_width
            i += 1

        return "".join(wrapped)

    def close(self) -> None:
        if not self._open:
            return
        self._open = False
        if self._on_close is not None:
            self._on_close()

    def open_if_available(self, ctx: PerkMenuContext) -> bool:
        if self._open:
            return True
        if ctx.assets is None:
            return False
        choices = perk_selection_current_choices(
            ctx.state,
            ctx.players,
            ctx.perk_state,
            game_mode=int(ctx.game_mode),
            player_count=int(ctx.player_count),
        )
        if not choices:
            self._open = False
            return False
        if ctx.play_sfx is not None:
            ctx.play_sfx("sfx_ui_panelclick")
        self._open = True
        self._selected_index = 0
        return True

    def tick_timeline(self, dt_ui_ms: float) -> None:
        if self._open:
            self._timeline_ms = clamp(self._timeline_ms + float(dt_ui_ms), 0.0, PERK_MENU_TRANSITION_MS)
        else:
            self._timeline_ms = clamp(self._timeline_ms - float(dt_ui_ms), 0.0, PERK_MENU_TRANSITION_MS)

    def handle_input(self, ctx: PerkMenuContext, *, dt_frame: float, dt_ui_ms: float) -> None:
        if ctx.assets is None:
            self.close()
            return

        choices = perk_selection_current_choices(
            ctx.state,
            ctx.players,
            ctx.perk_state,
            game_mode=int(ctx.game_mode),
            player_count=int(ctx.player_count),
        )
        if not choices:
            self.close()
            return

        if self._selected_index >= len(choices):
            self._selected_index = 0

        if rl.is_key_pressed(rl.KeyboardKey.KEY_DOWN):
            self._selected_index = (self._selected_index + 1) % len(choices)
        if rl.is_key_pressed(rl.KeyboardKey.KEY_UP):
            self._selected_index = (self._selected_index - 1) % len(choices)

        screen_w = float(rl.get_screen_width())
        screen_h = float(rl.get_screen_height())
        scale = ui_scale(screen_w, screen_h)
        origin = ui_origin(screen_w, screen_h, scale)
        slide_x = perk_menu_panel_slide_x(self._timeline_ms, width=self._layout.panel_size.x)

        click = rl.is_mouse_button_pressed(rl.MouseButton.MOUSE_BUTTON_LEFT)

        master_owned = int(ctx.player.perk_counts[int(PerkId.PERK_MASTER)]) > 0
        expert_owned = int(ctx.player.perk_counts[int(PerkId.PERK_EXPERT)]) > 0
        computed = perk_menu_compute_layout(
            self._layout,
            screen_w=screen_w,
            origin=origin,
            scale=scale,
            choice_count=len(choices),
            expert_owned=expert_owned,
            master_owned=master_owned,
            panel_slide_x=slide_x,
        )

        for idx, perk_id in enumerate(choices):
            label = perk_display_name(int(perk_id), fx_toggle=int(ctx.fx_toggle))
            item_pos = computed.list_pos.offset(dy=float(idx) * computed.list_step_y)
            rect = menu_item_hit_rect(ctx.font, label, pos=item_pos, scale=scale)
            if rect.contains(ctx.mouse):
                self._selected_index = idx
                if click:
                    if ctx.play_sfx is not None:
                        ctx.play_sfx("sfx_ui_buttonclick")
                    picked = perk_selection_pick(
                        ctx.state,
                        ctx.players,
                        ctx.perk_state,
                        idx,
                        game_mode=int(ctx.game_mode),
                        player_count=int(ctx.player_count),
                        dt=float(dt_frame),
                        creatures=ctx.creatures,
                    )
                    if picked is not None and ctx.play_sfx is not None:
                        ctx.play_sfx("sfx_ui_bonus")
                    if picked is not None and self._on_pick is not None:
                        self._on_pick(int(idx))
                    self.close()
                    return
                break

        cancel_w = button_width(
            ctx.font,
            self._cancel_button.label,
            scale=scale,
            force_wide=self._cancel_button.force_wide,
        )
        if button_update(
            self._cancel_button,
            pos=computed.cancel_pos,
            width=cancel_w,
            dt_ms=float(dt_ui_ms),
            mouse=ctx.mouse,
            click=click,
        ):
            if ctx.play_sfx is not None:
                ctx.play_sfx("sfx_ui_buttonclick")
            self.close()
            return

        if rl.is_key_pressed(rl.KeyboardKey.KEY_ENTER) or rl.is_key_pressed(rl.KeyboardKey.KEY_SPACE):
            if ctx.play_sfx is not None:
                ctx.play_sfx("sfx_ui_buttonclick")
            picked = perk_selection_pick(
                ctx.state,
                ctx.players,
                ctx.perk_state,
                self._selected_index,
                game_mode=int(ctx.game_mode),
                player_count=int(ctx.player_count),
                dt=float(dt_frame),
                creatures=ctx.creatures,
            )
            if picked is not None and ctx.play_sfx is not None:
                ctx.play_sfx("sfx_ui_bonus")
            if picked is not None and self._on_pick is not None:
                self._on_pick(int(self._selected_index))
            self.close()

    def draw(self, ctx: PerkMenuContext) -> None:
        menu_t = clamp(self._timeline_ms / PERK_MENU_TRANSITION_MS, 0.0, 1.0)
        if menu_t <= 1e-3:
            return
        if ctx.assets is None:
            return

        choices = perk_selection_current_choices(
            ctx.state,
            ctx.players,
            ctx.perk_state,
            game_mode=int(ctx.game_mode),
            player_count=int(ctx.player_count),
        )
        if not choices:
            return

        screen_w = float(rl.get_screen_width())
        screen_h = float(rl.get_screen_height())
        scale = ui_scale(screen_w, screen_h)
        origin = ui_origin(screen_w, screen_h, scale)
        slide_x = perk_menu_panel_slide_x(self._timeline_ms, width=self._layout.panel_size.x)

        master_owned = int(ctx.player.perk_counts[int(PerkId.PERK_MASTER)]) > 0
        expert_owned = int(ctx.player.perk_counts[int(PerkId.PERK_EXPERT)]) > 0
        computed = perk_menu_compute_layout(
            self._layout,
            screen_w=screen_w,
            origin=origin,
            scale=scale,
            choice_count=len(choices),
            expert_owned=expert_owned,
            master_owned=master_owned,
            panel_slide_x=slide_x,
        )

        panel_tex = ctx.assets.menu_panel
        if panel_tex is not None:
            draw_classic_menu_panel(panel_tex, dst=computed.panel.to_rl(), shadow=bool(ctx.fx_detail))

        title_tex = ctx.assets.title_pick_perk
        if title_tex is not None:
            src = rl.Rectangle(0.0, 0.0, float(title_tex.width), float(title_tex.height))
            rl.draw_texture_pro(
                title_tex,
                src,
                computed.title.to_rl(),
                rl.Vector2(0.0, 0.0),
                0.0,
                rl.WHITE,
            )

        sponsor = None
        if master_owned:
            sponsor = "extra perks sponsored by the Perk Master"
        elif expert_owned:
            sponsor = "extra perk sponsored by the Perk Expert"
        if sponsor:
            draw_ui_text(ctx.font, sponsor, computed.sponsor_pos, scale=scale, color=UI_SPONSOR_COLOR)

        for idx, perk_id in enumerate(choices):
            label = perk_display_name(int(perk_id), fx_toggle=int(ctx.fx_toggle))
            item_pos = computed.list_pos.offset(dy=float(idx) * computed.list_step_y)
            rect = menu_item_hit_rect(ctx.font, label, pos=item_pos, scale=scale)
            hovered = rect.contains(ctx.mouse) or (idx == self._selected_index)
            draw_menu_item(ctx.font, label, pos=item_pos, scale=scale, hovered=hovered)

        selected = choices[self._selected_index]
        desc = perk_display_description(int(selected), fx_toggle=int(ctx.fx_toggle))
        if ctx.font is not None:
            desc = self._prewrapped_perk_desc(int(selected), ctx.font, fx_toggle=int(ctx.fx_toggle))
        draw_ui_text(
            ctx.font,
            desc,
            computed.desc.top_left,
            scale=scale,
            color=UI_TEXT_COLOR,
        )

        cancel_w = button_width(
            ctx.font, self._cancel_button.label, scale=scale, force_wide=self._cancel_button.force_wide
        )
        button_draw(
            ctx.assets,
            ctx.font,
            self._cancel_button,
            pos=computed.cancel_pos,
            width=cancel_w,
            scale=scale,
        )
