#-----------------------------------------------------------------------------------------------------------------------
# Load per-game OHD assets from `Assets/{NN}/{resolution}/`.
# Background uses `NN-BG`, blocking mask `NN-BL` (1920x1080 topboxes only; not used at 3840/1920x360),
# digits from Shared aligned cells; frame from `Assets/Shared/{resolution}/FR.png`.
#-----------------------------------------------------------------------------------------------------------------------
import os
import re
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import pygame

from ethergold_reader import PROGRESSIVE_GAME_IDS
from ohd_aligned_cells import AlignedCellLibrary, load_aligned_cell_library
from ohd_layout import (
    OhdLayout,
    RESOLUTION_1920,
    RESOLUTION_3840,
    StripCell,
    blocking_asset_applies,
    default_assets_root,
    default_shared_resolution_dir,
    digit_strip_cells,
    format_game_index,
    infer_uniform_digit_size,
    load_layout,
    symbol_strip_cells,
)

OHD_ASSET_INDICES = PROGRESSIVE_GAME_IDS

_BLANK_STRIP: Optional[pygame.Surface] = None


def _blank_digit_strip() -> pygame.Surface:
    global _BLANK_STRIP
    if _BLANK_STRIP is None:
        _BLANK_STRIP = pygame.Surface((1, 1), pygame.SRCALPHA)
    return _BLANK_STRIP


@dataclass
class OhdGameAssets:
    game_id: int
    resolution: str
    background: pygame.Surface
    blocking: Optional[pygame.Surface]
    frame: pygame.Surface
    digit_strip: pygame.Surface
    layout: OhdLayout
    digit_files: Optional[Dict[int, pygame.Surface]] = None
    aligned: Optional[AlignedCellLibrary] = None
    background_path: str = ""
    blocking_path: str = ""
    frame_path: str = ""
    digit_path: str = ""


def _pick_file(folder: str, patterns: List[str], exclude: Optional[List[str]] = None) -> Optional[str]:
    if not os.path.isdir(folder):
        return None
    exclude = exclude or ["preview", "sample", ".psd"]
    names = sorted(os.listdir(folder))
    for pattern in patterns:
        rx = re.compile(pattern, re.IGNORECASE)
        for name in names:
            low = name.lower()
            if not rx.search(name):
                continue
            if any(x in low for x in exclude):
                continue
            if low.endswith((".png", ".jpg", ".jpeg", ".webp")):
                return os.path.join(folder, name)
    return None


def _load_surface(path: str) -> pygame.Surface:
    """Load image; convert only after pygame.display.set_mode() (required by pygame)."""
    surf = pygame.image.load(path)
    return _optimize_surface(surf)


def _optimize_surface(surf: pygame.Surface) -> pygame.Surface:
    if pygame.display.get_surface() is None:
        return surf
    if surf.get_flags() & pygame.SRCALPHA:
        return surf.convert_alpha()
    return surf.convert()


def optimize_game_assets(assets: OhdGameAssets) -> None:
    """Call once after set_mode() for faster blits (optional if loaded after set_mode)."""
    assets.background = _optimize_surface(assets.background)
    if assets.blocking is not None:
        assets.blocking = _optimize_surface(assets.blocking)
    assets.frame = _optimize_surface(assets.frame)
    if assets.digit_strip.get_width() > 1 or assets.digit_strip.get_height() > 1:
        assets.digit_strip = _optimize_surface(assets.digit_strip)
    if assets.digit_files:
        for d in range(10):
            assets.digit_files[d] = _optimize_surface(assets.digit_files[d])
    if assets.aligned:
        assets.aligned.dollar.surface = _optimize_surface(assets.aligned.dollar.surface)
        assets.aligned.comma.surface = _optimize_surface(assets.aligned.comma.surface)
        assets.aligned.decimal.surface = _optimize_surface(assets.aligned.decimal.surface)
        for spec in assets.aligned.digits.values():
            spec.surface = _optimize_surface(spec.surface)


def _load_digit_files(folder: str) -> Optional[Dict[int, pygame.Surface]]:
    digits: Dict[int, pygame.Surface] = {}
    for d in range(10):
        path = os.path.join(folder, f"digit_{d}.png")
        if not os.path.isfile(path):
            return None
        digits[d] = _load_surface(path)
    return digits


def load_game_assets(
    game_id: int,
    resolution: str = RESOLUTION_1920,
    assets_root: Optional[str] = None,
) -> Optional[OhdGameAssets]:
    if game_id not in OHD_ASSET_INDICES:
        return None
    root = assets_root or default_assets_root()
    game_token = format_game_index(game_id)
    folder = os.path.join(root, game_token, resolution)
    if not os.path.isdir(folder):
        print(f"ohd_assets: missing folder {folder}")
        return None

    background_path = _pick_file(
        folder, [rf"^{game_token}-BG(?:\.[A-Za-z0-9]+)?$", r"theme", r"32x9.*theme"]
    )
    shared_folder = default_shared_resolution_dir(resolution)
    frame_path = _pick_file(
        shared_folder,
        [
            r"^FR(?:\.[A-Za-z0-9]+)?$",
            r"^Shared-FR(?:\.[A-Za-z0-9]+)?$",
            r"^[0-9]{2}-FR(?:\.[A-Za-z0-9]+)?$",
            r"frame_01",
            r"32x9_frame",
        ],
    )
    if frame_path is None:
        frame_path = _pick_file(
            folder, [rf"^{game_token}-FR(?:\.[A-Za-z0-9]+)?$", r"frame_01", r"32x9_frame"]
        )

    if not background_path or not frame_path:
        print(f"ohd_assets: incomplete assets in {folder} (need BG + shared FR)")
        return None

    blocking_path = None
    if blocking_asset_applies(resolution):
        blocking_path = _pick_file(
            folder, [rf"^{game_token}-BL(?:\.[A-Za-z0-9]+)?$"]
        )
    blocking = _load_surface(blocking_path) if blocking_path else None

    frame = _load_surface(frame_path)
    layout = load_layout(
        resolution,
        game_id,
        frame_size=(frame.get_width(), frame.get_height()),
        assets_root=root,
    )

    ac = getattr(layout, "aligned_cells", None) or {}
    shared_dir = default_shared_resolution_dir(resolution)
    if isinstance(ac, dict) and ac.get("shared_dir"):
        shared_dir = str(ac["shared_dir"])
    kw: dict = {}
    if isinstance(ac, dict):
        if ac.get("anchor_right") is not None:
            kw["anchor_right"] = int(ac["anchor_right"])
        if ac.get("anchor_bottom") is not None:
            kw["anchor_bottom"] = int(ac["anchor_bottom"])
    aligned = load_aligned_cell_library(shared_dir, **kw)

    digit_path = _pick_file(folder, [r"gold_numbers"])
    digit_files = _load_digit_files(folder) if not aligned else None
    if aligned:
        strip = _blank_digit_strip()
    elif digit_path:
        strip = _load_surface(digit_path)
    else:
        print(
            f"ohd_assets: no shared cells in {shared_dir} and no gold_numbers strip in {folder}"
        )
        return None

    return OhdGameAssets(
        game_id=game_id,
        resolution=resolution,
        background=_load_surface(background_path),
        blocking=blocking,
        frame=frame,
        digit_strip=strip,
        layout=layout,
        digit_files=digit_files,
        aligned=aligned,
        background_path=background_path,
        blocking_path=blocking_path or "",
        frame_path=frame_path,
        digit_path=digit_path or "",
    )


def build_glyph_surfaces(assets: OhdGameAssets) -> List[pygame.Surface]:
    """One pygame surface per digit 0-9, scaled to default slot size if available."""
    if assets.aligned:
        return [assets.aligned.digits[d].surface for d in range(10)]
    if assets.digit_files:
        return [assets.digit_files[d] for d in range(10)]

    glyphs: List[pygame.Surface] = []
    cells = digit_strip_cells(assets.layout)
    if not cells:
        w, h = assets.digit_strip.get_width(), assets.digit_strip.get_height()
        cell_h = max(1, h // 10)
        cells = [StripCell("digit", 0, d * cell_h, w, cell_h, digit=d) for d in range(10)]

    target_h = assets.layout.digit_cell_h or 120
    if assets.layout.digit_slots:
        target_h = assets.layout.digit_slots[0].h

    order = list(cells)
    if assets.layout.digit_order == "0_at_bottom":
        order = sorted(cells, key=lambda c: -c.y)

    by_digit = {c.digit: c for c in order if c.digit is not None}
    for d in range(10):
        cell = by_digit.get(d)
        if cell is None:
            nw = max(1, int(round(target_h * 196 / 296)))
            glyphs.append(pygame.Surface((nw, target_h)))
            continue
        rect = pygame.Rect(cell.x, cell.y, cell.w, cell.h)
        glyph = assets.digit_strip.subsurface(rect).copy()
        scale = target_h / max(1, cell.h)
        nw = max(1, int(round(cell.w * scale)))
        nh = target_h
        if (nw, nh) != glyph.get_size():
            glyph = pygame.transform.smoothscale(glyph, (nw, nh))
        glyphs.append(glyph)
    return glyphs


def build_symbol_glyphs(assets: OhdGameAssets) -> Dict[str, pygame.Surface]:
    """Crop $ , . from the strip; scale relative to digit slot height (baseline draw in renderer)."""
    if assets.aligned:
        return {
            "dollar": assets.aligned.dollar.surface,
            "comma": assets.aligned.comma.surface,
            "decimal": assets.aligned.decimal.surface,
        }
    symbols: Dict[str, pygame.Surface] = {}
    cells = symbol_strip_cells(assets.layout)
    if not cells or assets.digit_files:
        return symbols

    slot_h = 120
    if assets.layout.digit_slots:
        slot_h = assets.layout.digit_slots[0].h

    _, digit_h19, _ = infer_uniform_digit_size(assets.layout)
    if digit_h19 <= 0:
        digit_h19 = max(c.h for c in digit_strip_cells(assets.layout)) or 296

    decimal_cell = next((c for c in cells if c.name == "decimal"), None)
    punct_strip_h = decimal_cell.h if decimal_cell else 115
    punct_target_h = max(1, int(round(slot_h * punct_strip_h / digit_h19)))

    for cell in cells:
        if not cell.name:
            continue
        rect = pygame.Rect(cell.x, cell.y, cell.w, cell.h)
        glyph = assets.digit_strip.subsurface(rect).copy()
        if cell.name in ("comma", "decimal"):
            target_h = punct_target_h
        else:
            target_h = max(1, int(round(slot_h * cell.h / digit_h19)))
        scale = target_h / max(1, glyph.get_height())
        target_w = max(1, int(round(glyph.get_width() * scale)))
        symbols[cell.name] = pygame.transform.smoothscale(glyph, (target_w, target_h))
    return symbols


def slice_digit_strip(strip: pygame.Surface, digit_count: int = 10) -> Tuple[int, int]:
    w, h = strip.get_width(), strip.get_height()
    return w, max(1, h // digit_count)