from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Sequence, Tuple, Union

import pygame

from ..core.engine import OdometerConfig, OdometerEngine
from ..theme.aligned import OdometerTheme, TokenKind, total_string_width, visible_tokens_for_digits


@dataclass(frozen=True)
class BlitOp:
    surface: pygame.Surface
    dest: Tuple[int, int]
    area: Optional[pygame.Rect] = None
    clip: Optional[pygame.Rect] = None


Placement = Union[pygame.Rect, Tuple[int, int]]


class OdometerWidget:
    def __init__(
        self,
        theme: OdometerTheme,
        config: Optional[OdometerConfig] = None,
        initial_target_pennies: int = 0,
    ):
        self.theme = theme
        self.config = config or OdometerConfig()
        self.engine = OdometerEngine(
            self.config,
            initial_target_pennies=initial_target_pennies,
        )
        self._static_surface: Optional[pygame.Surface] = None
        self._static_surface_key: Optional[Tuple[int, ...]] = None
        self._has_received_live_target: bool = False

    @property
    def display_pennies(self) -> int:
        return self.engine.display_pennies

    @property
    def target_pennies(self) -> int:
        return self.engine.target_pennies

    def set_target_pennies(
        self,
        value: int,
        *,
        immediate: bool = False,
        first_value_immediate: bool = False,
        mark_live: bool = True,
    ) -> None:
        use_immediate = immediate or (
            first_value_immediate and not self._has_received_live_target
        )
        if use_immediate:
            self.engine.display_immediate_pennies(value)
        else:
            self.engine.set_target_pennies(value)
        if mark_live:
            self._has_received_live_target = True

    def display_immediate_pennies(self, value: int) -> None:
        self.set_target_pennies(value, immediate=True, mark_live=False)

    def clear_live_target(self) -> None:
        """No fresh CSV / no pennies — hide digits (do not paint $0.00)."""
        self._has_received_live_target = False

    def update(self, dt_seconds: float) -> None:
        self.engine.update(dt_seconds)

    def lag_pennies(self) -> int:
        return self.engine.lag_pennies()

    def is_animating(self) -> bool:
        return self.engine.is_animating()

    def draw(self, surface: pygame.Surface) -> None:
        self._apply_ops(surface, self.get_draw_ops())

    def draw_static_only(self, surface: pygame.Surface) -> None:
        """Blit background + frame/bezel only (no odometer digits)."""
        static = self._ensure_default_static_surface()
        if static is not None:
            surface.blit(static, (0, 0))
        else:
            container = self._resolve_container(None)
            static_op = self._build_static_op(None, container)
            if static_op is not None:
                self._apply_ops(surface, [static_op])

    def draw_digits_only(self, surface: pygame.Surface) -> None:
        """Redraw odometer digits on an existing static background."""
        self._apply_ops(surface, self.get_draw_ops(include_static=False))

    def draw_at(self, surface: pygame.Surface, rect_or_anchor: Placement) -> None:
        self._apply_ops(surface, self.get_draw_ops(rect_or_anchor))

    def get_draw_ops(
        self,
        rect_or_anchor: Optional[Placement] = None,
        *,
        include_static: bool = True,
    ) -> List[BlitOp]:
        ops: List[BlitOp] = []
        container = self._resolve_container(rect_or_anchor)

        if include_static:
            static_op = self._build_static_op(rect_or_anchor, container)
            if static_op is not None:
                ops.append(static_op)

        if not self._has_received_live_target:
            return ops

        digits = self.engine.get_display_digits()
        tokens = visible_tokens_for_digits(digits)
        string_w = total_string_width(self.theme.aligned, tokens)
        x_left = container.x + max(0, (container.w - string_w) // 2)
        x_right = x_left + string_w
        bottom_y = self._resolve_bottom_y(rect_or_anchor, container)

        for rtl_index, (kind, digit) in enumerate(tokens):
            spec = self.theme.aligned.spec_for(kind, digit)
            x = x_right - spec.w
            y = bottom_y - spec.h
            x_right = x

            if kind == "digit" and digit is not None:
                power = _digit_power_from_rtl_index(tokens, rtl_index)
                draw_digit, scroll = self.engine.get_column_draw(power)
                glyph = self.theme.aligned.spec_for("digit", draw_digit)
                if scroll > 1e-6:
                    ops.extend(_scrolling_digit_ops(x, y, glyph.w, glyph.h, self.theme, draw_digit, scroll))
                    continue
                ops.append(BlitOp(glyph.surface, (x, y)))
                continue

            ops.append(BlitOp(spec.surface, (x, y)))
        return ops

    def _build_static_op(
        self, rect_or_anchor: Optional[Placement], container: pygame.Rect
    ) -> Optional[BlitOp]:
        if rect_or_anchor is None:
            surface = self._ensure_default_static_surface()
            if surface is not None:
                return BlitOp(surface, (0, 0))
            return None

        parts: List[Tuple[pygame.Surface, Tuple[int, int]]] = []
        if self.theme.background is not None and self.theme.background.get_size() == container.size:
            parts.append((self.theme.background, (0, 0)))
        anchor = self.theme.layout.frame_anchor
        if self.theme.blocking is not None:
            parts.append((self.theme.blocking, anchor))
        if self.theme.frame is not None:
            parts.append((self.theme.frame, anchor))
        if not parts:
            return None
        if len(parts) == 1:
            return BlitOp(parts[0][0], container.topleft)

        key = (
            container.w,
            container.h,
            id(self.theme.background) if self.theme.background is not None else 0,
            id(self.theme.blocking) if self.theme.blocking is not None else 0,
            id(self.theme.frame) if self.theme.frame is not None else 0,
        )
        if self._static_surface is None or self._static_surface_key != key:
            composed = pygame.Surface(container.size, pygame.SRCALPHA)
            for surf, pos in parts:
                composed.blit(surf, pos)
            self._static_surface = composed
            self._static_surface_key = key
        return BlitOp(self._static_surface, container.topleft)

    def _ensure_default_static_surface(self) -> Optional[pygame.Surface]:
        if (
            self.theme.background is None
            and self.theme.blocking is None
            and self.theme.frame is None
        ):
            return None

        size = self._default_screen_size()
        key = self._static_layer_key_for_size(size)
        if self._static_surface is not None and self._static_surface_key == key:
            return self._static_surface

        layer = pygame.Surface(size, pygame.SRCALPHA)
        self._blit_background_into(layer, size)
        self._blit_frame_blocking_into(layer)
        self._static_surface = layer
        self._static_surface_key = key
        return layer

    def _default_screen_size(self) -> Tuple[int, int]:
        size = self.theme.layout.screen_size
        if size != (0, 0):
            return size
        if self.theme.background is not None:
            return self.theme.background.get_size()
        if self.theme.frame is not None:
            fx, fy = self.theme.layout.frame_anchor
            fw, fh = self.theme.frame.get_size()
            return (fx + fw, fy + fh)
        if self.theme.blocking is not None:
            fx, fy = self.theme.layout.frame_anchor
            bw, bh = self.theme.blocking.get_size()
            return (fx + bw, fy + bh)
        return (0, 0)

    def _static_layer_key_for_size(self, size: Tuple[int, int]) -> Tuple[int, ...]:
        return (
            size[0],
            size[1],
            self.theme.layout.frame_anchor[0],
            self.theme.layout.frame_anchor[1],
            id(self.theme.background) if self.theme.background is not None else 0,
            id(self.theme.blocking) if self.theme.blocking is not None else 0,
            id(self.theme.frame) if self.theme.frame is not None else 0,
        )

    def _blit_background_into(
        self, layer: pygame.Surface, size: Tuple[int, int]
    ) -> None:
        if self.theme.background is None:
            return
        if self.theme.background.get_size() == size:
            layer.blit(self.theme.background, (0, 0))
        else:
            layer.blit(
                pygame.transform.smoothscale(self.theme.background, size),
                (0, 0),
            )

    def _blit_frame_blocking_into(self, layer: pygame.Surface) -> None:
        anchor = self.theme.layout.frame_anchor
        if self.theme.blocking is not None:
            layer.blit(self.theme.blocking, anchor)
        if self.theme.frame is not None:
            layer.blit(self.theme.frame, anchor)

    def render_background_layer(self, output_size: Tuple[int, int]) -> pygame.Surface:
        layer = pygame.Surface(output_size, pygame.SRCALPHA)
        self._blit_background_into(layer, output_size)
        return layer

    def render_frame_blocking_layer(self, output_size: Tuple[int, int]) -> pygame.Surface:
        layer = pygame.Surface(output_size, pygame.SRCALPHA)
        self._blit_frame_blocking_into(layer)
        return layer

    def render_digits_layer(self, output_size: Tuple[int, int]) -> pygame.Surface:
        layer = pygame.Surface(output_size, pygame.SRCALPHA)
        self._apply_ops(layer, self.get_draw_ops(include_static=False))
        return layer

    def _resolve_container(self, rect_or_anchor: Optional[Placement]) -> pygame.Rect:
        if rect_or_anchor is None:
            if self.theme.frame is not None:
                return pygame.Rect(self.theme.layout.frame_anchor, self.theme.frame.get_size())
            if self.theme.layout.screen_size != (0, 0):
                return pygame.Rect((0, 0), self.theme.layout.screen_size)
            return pygame.Rect((0, 0), (0, 0))
        if isinstance(rect_or_anchor, pygame.Rect):
            return rect_or_anchor.copy()
        anchor = rect_or_anchor
        if self.theme.frame is not None:
            return pygame.Rect(anchor, self.theme.frame.get_size())
        if self.theme.layout.screen_size != (0, 0):
            return pygame.Rect(anchor, self.theme.layout.screen_size)
        return pygame.Rect(anchor, (0, 0))

    def _resolve_bottom_y(
        self, rect_or_anchor: Optional[Placement], container: pygame.Rect
    ) -> int:
        if rect_or_anchor is None and self.theme.layout.aligned_anchor_bottom is not None:
            return self.theme.layout.aligned_anchor_bottom
        if self.theme.anchor_bottom_offset != 0:
            return container.y + self.theme.anchor_bottom_offset
        return container.bottom

    @staticmethod
    def _apply_ops(surface: pygame.Surface, ops: Sequence[BlitOp]) -> None:
        previous_clip = surface.get_clip()
        try:
            for op in ops:
                if op.clip is not None:
                    surface.set_clip(op.clip)
                else:
                    surface.set_clip(previous_clip)
                surface.blit(op.surface, op.dest, op.area)
        finally:
            surface.set_clip(previous_clip)


def _digit_power_from_rtl_index(
    tokens: List[Tuple[TokenKind, int | None]], rtl_index: int
) -> int:
    power = 0
    for index in range(rtl_index):
        if tokens[index][0] == "digit":
            power += 1
    return power


def _scrolling_digit_ops(
    x: int,
    y: int,
    w: int,
    h: int,
    theme: OdometerTheme,
    digit: int,
    scroll: float,
) -> List[BlitOp]:
    current = theme.aligned.digits[digit].surface
    nxt = theme.aligned.digits[(digit + 1) % 10].surface
    step = max(current.get_height(), nxt.get_height(), h)
    offset_y = -int(round(scroll * step))
    clip = pygame.Rect(x, y, w, h)
    return [
        BlitOp(current, (x, y + offset_y), clip=clip),
        BlitOp(nxt, (x, y + offset_y + step), clip=clip),
    ]