"""
spritesheet.py — SpriteSheet helper for Powernet.

Loads a PNG sprite sheet and provides clean methods to extract
individual frames, rows, grids, animation strips, and 9-slice tiles.

All filenames are resolved relative to  src/assets/textures/.

Known sheets
------------
ui_panels.png         (736 × 288)  — flat UI panel tiles in 3 colour states
ui_animated.png       (128 × 128)  — animated button / panel tiles
ui_pins.png           (192 × 144)  — electrical pin / node icons (8 × 6, 24 px)
wasteland_props.png   (416 × 448)  — post-apoc prop tiles (~32 px grid)
IconsShadow-32.png    (320 × 320)  — 100 shadow-style UI icons (32 px, JSON atlas)
"""

from __future__ import annotations

import pygame
from pathlib import Path

_TEXTURES = Path(__file__).parent.parent / "assets" / "textures"


class SpriteSheet:
    """Load a sprite sheet image and slice it into pygame Surfaces."""

    def __init__(self, filename: str) -> None:
        """
        Parameters
        ----------
        filename : str
            PNG filename inside ``src/assets/textures/``.
        """
        path = _TEXTURES / filename
        self._sheet: pygame.Surface = pygame.image.load(str(path)).convert_alpha()
        self.width, self.height = self._sheet.get_size()

    # ------------------------------------------------------------------
    # Core extraction
    # ------------------------------------------------------------------

    def clip(
        self,
        x: int,
        y: int,
        w: int,
        h: int,
        scale: float = 1.0,
    ) -> pygame.Surface:
        """
        Extract a rectangular region from the sheet.

        Parameters
        ----------
        x, y    Top-left pixel position in the source image.
        w, h    Width and height of the region in source pixels.
        scale   Uniform scale multiplier applied after extraction.
        """
        surf = pygame.Surface((w, h), pygame.SRCALPHA)
        surf.blit(self._sheet, (0, 0), (x, y, w, h))
        if scale != 1.0:
            surf = pygame.transform.scale_by(surf, scale)
        return surf

    def clip_to(
        self,
        x: int,
        y: int,
        w: int,
        h: int,
        target_w: int,
        target_h: int,
    ) -> pygame.Surface:
        """Extract and resize to an explicit pixel size."""
        surf = self.clip(x, y, w, h)
        return pygame.transform.scale(surf, (target_w, target_h))

    # ------------------------------------------------------------------
    # Row / grid / animation helpers
    # ------------------------------------------------------------------

    def row(
        self,
        y: int,
        cell_w: int,
        cell_h: int,
        count: int,
        *,
        offset_x: int = 0,
        gap: int = 0,
        scale: float = 1.0,
    ) -> list[pygame.Surface]:
        """
        Extract a horizontal strip of equal-sized cells.

        Parameters
        ----------
        y        Top pixel position of the strip.
        cell_w   Width of each cell.
        cell_h   Height of each cell.
        count    Number of cells to extract.
        offset_x Horizontal pixel offset before the first cell.
        gap      Pixel gap between consecutive cells.
        scale    Uniform scale applied to each extracted cell.
        """
        return [
            self.clip(offset_x + i * (cell_w + gap), y, cell_w, cell_h, scale)
            for i in range(count)
        ]

    def grid(
        self,
        cell_w: int,
        cell_h: int,
        cols: int,
        rows: int,
        *,
        offset_x: int = 0,
        offset_y: int = 0,
        gap_x: int = 0,
        gap_y: int = 0,
        scale: float = 1.0,
    ) -> list[list[pygame.Surface]]:
        """
        Parse a rectangular grid of uniformly-sized sprites.

        Returns
        -------
        list[list[Surface]]
            Indexed as ``result[row][col]``.
        """
        return [
            [
                self.clip(
                    offset_x + c * (cell_w + gap_x),
                    offset_y + r * (cell_h + gap_y),
                    cell_w,
                    cell_h,
                    scale,
                )
                for c in range(cols)
            ]
            for r in range(rows)
        ]

    def animation_strip(
        self,
        row_index: int,
        cell_w: int,
        cell_h: int,
        frame_count: int,
        *,
        offset_x: int = 0,
        offset_y: int = 0,
        gap: int = 0,
        scale: float = 1.0,
    ) -> list[pygame.Surface]:
        """
        Return a list of frames from a single animation strip (one row).

        Parameters
        ----------
        row_index   Zero-based row index.
        cell_w/h    Size of each frame in source pixels.
        frame_count How many frames to read.
        offset_x/y  Pixel offset before the grid starts.
        gap         Pixel gap between frames (and between rows if stacked).
        scale       Uniform scale applied to each frame.
        """
        y = offset_y + row_index * (cell_h + gap)
        return self.row(y, cell_w, cell_h, frame_count, offset_x=offset_x, gap=gap, scale=scale)

    # ------------------------------------------------------------------
    # 9-slice helper for scalable panels / buttons
    # ------------------------------------------------------------------

    def nine_slice(
        self,
        x: int,
        y: int,
        src_w: int,
        src_h: int,
        border: int,
        target_w: int,
        target_h: int,
    ) -> pygame.Surface:
        """
        Build a scalable surface via 9-slice from a source rectangle.

        The source rectangle is divided into a 3×3 grid.  Corners keep
        their size; edges are stretched along one axis; the centre is
        stretched along both.

        Parameters
        ----------
        x, y          Top-left of source rectangle in the sheet.
        src_w, src_h  Size of source rectangle in pixels.
        border        Corner/edge size in source pixels.
        target_w/h    Desired output size in pixels.
        """
        out = pygame.Surface((target_w, target_h), pygame.SRCALPHA)

        b = border
        inner_sw = src_w - 2 * b
        inner_sh = src_h - 2 * b
        inner_dw = target_w - 2 * b
        inner_dh = target_h - 2 * b

        # Helper: extract a slice and stretch it to (dw, dh)
        def _blit(sx: int, sy: int, sw: int, sh: int, dx: int, dy: int, dw: int, dh: int) -> None:
            piece = self.clip(x + sx, y + sy, sw, sh)
            if dw != sw or dh != sh:
                piece = pygame.transform.scale(piece, (dw, dh))
            out.blit(piece, (dx, dy))

        # Corners
        _blit(0,              0,              b,        b,        0,            0,            b,        b)
        _blit(src_w - b,      0,              b,        b,        target_w - b, 0,            b,        b)
        _blit(0,              src_h - b,      b,        b,        0,            target_h - b, b,        b)
        _blit(src_w - b,      src_h - b,      b,        b,        target_w - b, target_h - b, b,        b)

        # Edges
        _blit(b,              0,              inner_sw, b,        b,            0,            inner_dw, b)         # top
        _blit(b,              src_h - b,      inner_sw, b,        b,            target_h - b, inner_dw, b)         # bottom
        _blit(0,              b,              b,        inner_sh, 0,            b,            b,        inner_dh)  # left
        _blit(src_w - b,      b,              b,        inner_sh, target_w - b, b,            b,        inner_dh)  # right

        # Centre
        _blit(b, b, inner_sw, inner_sh, b, b, inner_dw, inner_dh)

        return out


# ---------------------------------------------------------------------------
# IconAtlas  (Aseprite JSON packed atlas)
# ---------------------------------------------------------------------------


class IconAtlas:
    """
    Loads an Aseprite JSON-packed sprite atlas and provides icons by index.

    Frame names are expected to follow the pattern ``"Name N.ase"`` where
    *N* is the integer index (as exported by Aseprite).

    Usage::

        atlas = IconAtlas("IconsShadow-32.png", "IconsShadow-32.json")
        surf  = atlas.icon(12, size=24)   # returns a 24×24 Surface
    """

    def __init__(self, image_name: str, json_name: str) -> None:
        import json as _json

        self._sheet = pygame.image.load(str(_TEXTURES / image_name)).convert_alpha()

        with open(_TEXTURES / json_name) as fh:
            data = _json.load(fh)

        self._frames: dict[int, pygame.Rect] = {}
        for name, info in data["frames"].items():
            try:
                idx = int(name.rsplit(" ", 1)[-1].replace(".ase", ""))
                fr = info["frame"]
                self._frames[idx] = pygame.Rect(fr["x"], fr["y"], fr["w"], fr["h"])
            except (ValueError, KeyError):
                continue

        self._cache: dict[tuple[int, int], pygame.Surface] = {}

    def icon(self, index: int, size: int = 24) -> pygame.Surface:
        """Return icon *index* scaled to *size*×*size* pixels."""
        key = (index, size)
        if key in self._cache:
            return self._cache[key]
        rect = self._frames.get(index)
        if rect is None:
            return pygame.Surface((size, size), pygame.SRCALPHA)
        src = self._sheet.subsurface(rect)
        scaled = pygame.transform.scale(src, (size, size)) if size != rect.w else src.copy()
        self._cache[key] = scaled
        return scaled

    @property
    def count(self) -> int:
        return len(self._frames)
