#-----------------------------------------------------------------------------------------------------------------------
# Aligned odometer cells under `Assets/Shared/{resolution}/` (shared across all games).
# Filenames encode cell size: 0_175x320.png, $_185x320.png or 185x320_$.png, comma_66x320.png, etc.
# String is built right-to-left from a fixed lower-right anchor (no leading zeros).
#-----------------------------------------------------------------------------------------------------------------------
from __future__ import annotations

import os
import re
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union

import pygame

from ohd_layout import COMMA_LEFT_POWER

TokenKind = str  # "digit" | "dollar" | "comma" | "decimal"


@dataclass
class CellSpec:
    kind: TokenKind
    w: int
    h: int
    surface: pygame.Surface
    digit: Optional[int] = None


@dataclass
class AlignedOdometerLayout:
    anchor_right: int
    anchor_bottom: int
    shared_dir: str


@dataclass
class AlignedCellLibrary:
    dollar: CellSpec
    digits: Dict[int, CellSpec]
    comma: CellSpec
    decimal: CellSpec
    layout: AlignedOdometerLayout

    def spec_for(self, kind: TokenKind, digit: Optional[int] = None) -> CellSpec:
        if kind == "dollar":
            return self.dollar
        if kind == "comma":
            return self.comma
        if kind == "decimal":
            return self.decimal
        if kind == "digit" and digit is not None:
            return self.digits[digit]
        raise KeyError(f"unknown token {kind!r} digit={digit!r}")


_FILENAME_DIGIT = re.compile(r"^(\d)_(\d+)x(\d+)\.png$", re.I)
_FILENAME_DOLLAR_SUFFIX = re.compile(r"^(\d+)x(\d+)_\$\.png$", re.I)
_FILENAME_DOLLAR_PREFIX = re.compile(r"^\$_(\d+)x(\d+)\.png$", re.I)
_FILENAME_PUNCT = re.compile(r"^(comma|decimal)_(\d+)x(\d+)\.png$", re.I)

# Lower-right of odometer row (frame-centered full string; tune in shared `layout.json`).
_DEFAULT_ANCHOR_RIGHT_1920 = 1798
_DEFAULT_ANCHOR_BOTTOM_1920 = 962


def default_shared_cells_dir(resolution: str = "1920x1080") -> str:
    from ohd_layout import default_shared_resolution_dir

    return default_shared_resolution_dir(resolution)


def _load_png(path: str) -> pygame.Surface:
    return pygame.image.load(path)


def load_aligned_cell_library(
    shared_dir: Optional[str] = None,
    *,
    anchor_right: int = _DEFAULT_ANCHOR_RIGHT_1920,
    anchor_bottom: int = _DEFAULT_ANCHOR_BOTTOM_1920,
) -> Optional[AlignedCellLibrary]:
    folder = shared_dir or default_shared_cells_dir()
    if not os.path.isdir(folder):
        return None

    digit_specs: Dict[int, Tuple[int, int, str]] = {}
    dollar_path: Optional[str] = None
    dollar_size: Optional[Tuple[int, int]] = None
    punct: Dict[str, Tuple[int, int, str]] = {}

    for name in os.listdir(folder):
        low = name.lower()
        if not low.endswith(".png"):
            continue
        path = os.path.join(folder, name)
        m = _FILENAME_DIGIT.match(name)
        if m:
            d = int(m.group(1))
            digit_specs[d] = (int(m.group(2)), int(m.group(3)), path)
            continue
        m = _FILENAME_DOLLAR_SUFFIX.match(name) or _FILENAME_DOLLAR_PREFIX.match(name)
        if m:
            dollar_path = path
            dollar_size = (int(m.group(1)), int(m.group(2)))
            continue
        m = _FILENAME_PUNCT.match(name)
        if m:
            punct[m.group(1)] = (int(m.group(2)), int(m.group(3)), path)

    if len(digit_specs) != 10 or dollar_path is None or "comma" not in punct or "decimal" not in punct:
        return None

    digits: Dict[int, CellSpec] = {}
    for d in range(10):
        w, h, path = digit_specs[d]
        digits[d] = CellSpec("digit", w, h, _load_png(path), digit=d)

    dw, dh = dollar_size or (0, 0)
    dollar = CellSpec("dollar", dw, dh, _load_png(dollar_path))
    cw, ch, cp = punct["comma"]
    comma = CellSpec("comma", cw, ch, _load_png(cp))
    dw, dh, dp = punct["decimal"]
    decimal = CellSpec("decimal", dw, dh, _load_png(dp))

    return AlignedCellLibrary(
        dollar=dollar,
        digits=digits,
        comma=comma,
        decimal=decimal,
        layout=AlignedOdometerLayout(anchor_right, anchor_bottom, folder),
    )


def build_rtl_tokens(
    digits: List[int],
    leftmost: int,
    *,
    show_comma: bool,
) -> List[Tuple[TokenKind, Optional[int]]]:
    """
    Rightmost token first (pennies), stepping left.
    Decimal is always between the dollars column (power 2) and dimes (power 1):
    RTL order is … dollars, decimal, dimes, pennies …
    """
    tokens: List[Tuple[TokenKind, Optional[int]]] = []
    for power in range(0, leftmost + 1):
        tokens.append(("digit", digits[power]))
        # After dimes digit in RTL list → decimal sits left of dimes, right of dollars.
        if power == 1 and leftmost >= 2:
            tokens.append(("decimal", None))
        if power == 4 and show_comma:
            tokens.append(("comma", None))
    tokens.append(("dollar", None))
    return tokens


def total_string_width(library: AlignedCellLibrary, tokens: List[Tuple[TokenKind, Optional[int]]]) -> int:
    w = 0
    for kind, digit in tokens:
        spec = library.spec_for(kind, digit)
        w += spec.w
    return w


def aligned_anchor_from_layout(layout) -> Tuple[int, int]:
    ac = layout.aligned_cells if getattr(layout, "aligned_cells", None) else None
    if isinstance(ac, dict):
        ar = ac.get("anchor_right")
        ab = ac.get("anchor_bottom")
        if ar is not None and ab is not None:
            return int(ar), int(ab)
    if layout.digit_baseline_y is not None:
        bottom = int(layout.digit_baseline_y)
    elif layout.digit_slots:
        bottom = layout.digit_slots[0].y + layout.digit_slots[0].h
    else:
        bottom = _DEFAULT_ANCHOR_BOTTOM_1920
    right = _DEFAULT_ANCHOR_RIGHT_1920
    if layout.digit_slots:
        pennies = next((s for s in layout.digit_slots if s.power == 0), None)
        if pennies is not None:
            right = pennies.x + pennies.w
    return right, bottom