from __future__ import annotations
import json
from pathlib import Path
from typing import List, Union, Self

import numpy as np

from .symbols import Symbol, ANSI_PALETTE, CSS_PALETTE
from .errors import ArcError
from PIL import Image, ImageDraw, PngImagePlugin
from IPython.display import FileLink
from hashlib import sha256


class Grid:
    """
    Represents a 2D grid of an ARC task, with basic utilities.

    The Grid class is the core data structure for representing ARC tasks' input and output grids.
    It provides methods for loading and saving grids in various formats (JSON, NPY, PNG), as well as
    utilities for visualization and comparison.

    Methods:
        load_json(path): Load a grid from a JSON file.
        load_npy(path): Load a grid from a NumPy .npy file.
        load_png(path): Load a grid from a PNG file with embedded metadata.
        save_json(path): Save the grid to a JSON file.
        save_npy(path): Save the grid to a NumPy .npy file.
        save_png(path, cell_size, border_width): Save the grid as a PNG image with metadata.
        to_list(): Convert the grid to a nested Python list.
        to_numpy(): Return a copy of the underlying NumPy array.
        to_json(): Serialize the grid as a JSON string.
        shape: Return the (width, height) of the grid.
        __eq__(other): Check equality with another grid.
        __sub__(other): Count differing cells between two grids.
        __hash__(): Compute a hash of the grid's contents.
        __repr__(): Return an ANSI-colored string representation.
        __str__(): Return a plain string representation.
        _repr_html_(): Return an HTML representation for Jupyter/IPython.
    """

    def __init__(self, array: Union[List[List[int]], np.ndarray, None] = None) -> None:
        """
        Initialize a Grid instance.

        Args:
            array (List[List[int]] | np.ndarray | None): 2D list or NumPy array of symbol values.
                If None, the grid is uninitialized.

        Raises:
            ArcError: If the array is not 2D or contains invalid symbol values.
        """
        if array is None:
            self._array = None
        else:
            np_array = np.asarray(array, dtype=np.uint8)
            if np_array.ndim != 2:
                raise ArcError("Grid array must be 2D")
            symbols = {s.value for s in Symbol}
            if not set(np_array.flat).issubset(symbols):
                raise ArcError(
                    f"Grid values must be one of the symbols: {sorted(symbols)}"
                )
            self._array = np_array

    @classmethod
    def load_json(cls, path: Union[str, Path]) -> Self:
        """
        Load a Grid from a JSON file.

        Args:
            path (str | Path): Path to the JSON file.

        Returns:
            Grid: The loaded Grid instance.

        Raises:
            ArcError: If loading or parsing fails.
        """
        try:
            data = json.loads(Path(path).read_text())
        except Exception as e:
            raise ArcError(f"Failed to load JSON from {path}: {e}") from e
        return cls(data)

    @classmethod
    def load_npy(cls, path: Union[str, Path]) -> Self:
        """
        Load a Grid from a NumPy .npy file.

        Args:
            path (str | Path): Path to the .npy file.

        Returns:
            Grid: The loaded Grid instance.

        Raises:
            ArcError: If loading fails.
        """
        try:
            arr = np.load(path)
        except Exception as e:
            raise ArcError(f"Failed to load .npy from {path}: {e}") from e
        return cls(arr)

    @classmethod
    def load_png(cls, path: Union[str, Path]) -> Self:
        """
        Load a Grid from a PNG file with embedded JSON metadata. The loaded .png file should be generated by the `.save_png` method.

        Args:
            path (str | Path): Path to the PNG file.

        Returns:
            Grid: The loaded Grid instance.

        Raises:
            ArcError: If loading or hash validation fails.
        """
        try:
            img = Image.open(Path(path))
            data = json.loads(img.info["json"])
            grid = cls(data)
            if str(grid.__hash__()) != img.info["hash"]:
                raise ArcError(
                    f"PNG loaded from {path} of corrupted: hash, json mismatch"
                )
        except Exception as e:
            raise ArcError(f"Failed to load .png from {path}: {e}") from e
        return grid

    def save_json(self, path: Union[str, Path]) -> FileLink:
        """
        Save the grid to a JSON file.

        Args:
            path (str | Path): Destination file path.

        Returns:
            FileLink: IPython FileLink to the saved file.

        Raises:
            ArcError: If saving fails.
        """
        try:
            path = Path(path)
            path.write_text(self.to_json())
            return FileLink(path)
        except Exception as e:
            raise ArcError(f"Failed to save JSON to {path}: {e}") from e

    def save_npy(self, path: Union[str, Path]) -> FileLink:
        """
        Save the grid to a NumPy .npy file.

        Args:
            path (str | Path): Destination file path.

        Returns:
            FileLink: IPython FileLink to the saved file.

        Raises:
            ArcError: If saving fails.
        """
        try:
            path = Path(path)
            np.save(path, self._array)
            return FileLink(path)
        except Exception as e:
            raise ArcError(f"Failed to save .npy to {path}: {e}") from e

    def save_png(
        self, path: Union[str, Path], cell_size: int = 200, border_width: int = 5
    ) -> FileLink:
        """
        Save the grid as a PNG image with embedded metadata.

        Args:
            path (str | Path): Destination file path.
            cell_size (int): Size of each cell in pixels.
            border_width (int): Width of cell borders in pixels.

        Returns:
            FileLink: IPython FileLink to the saved file.

        Raises:
            ArcError: If saving fails.
        """
        try:
            img = Image.new(
                "RGB",
                (
                    self.shape[0] * (cell_size + 2 * border_width),
                    self.shape[1] * (cell_size + 2 * border_width),
                ),
            )
            draw = ImageDraw.Draw(img)
            for y, row in enumerate(self.to_list()):
                for x, value in enumerate(row):
                    draw.rectangle(
                        [
                            (
                                x * (cell_size + 2 * border_width),
                                y * (cell_size + 2 * border_width),
                            ),
                            (
                                (x + 1) * (cell_size + 2 * border_width),
                                (y + 1) * (cell_size + 2 * border_width),
                            ),
                        ],
                        fill=CSS_PALETTE[Symbol(value)],
                        outline=(50, 50, 50),
                        width=border_width,
                    )
            metadata = PngImagePlugin.PngInfo()
            metadata.add_text("hash", str(self.__hash__()))
            metadata.add_text("width", str(self.shape[0]))
            metadata.add_text("height", str(self.shape[1]))
            metadata.add_text("json", self.to_json())
            path = Path(path)
            img.save(path, format="PNG", pnginfo=metadata)
            return FileLink(path)
        except Exception as e:
            raise ArcError(f"Failed to save .png to {path}: {e}") from e

    def to_list(self) -> List[List[int]]:
        """
        Convert the grid to a nested Python list.

        Returns:
            List[List[int]]: The grid as a list of lists.
        """
        return self._array.tolist()

    def to_numpy(self) -> np.ndarray:
        """
        Return a copy of the underlying NumPy array.

        Returns:
            np.ndarray: The grid as a NumPy array.
        """
        return self._array.copy()

    def to_json(self) -> str:
        """
        Serialize the grid as a JSON string.

        Returns:
            str: The grid in JSON format.
        """
        return json.dumps(self.to_list())

    @property
    def shape(self) -> tuple[int, int]:
        """
        Get the (width, height) of the grid.

        Returns:
            tuple[int, int]: (width, height) of the grid.
        """
        return (self._array.shape[1], self._array.shape[0])

    def __eq__(self, other: object) -> bool:
        """
        Check equality with another Grid.

        Args:
            other (object): Another Grid instance.

        Returns:
            bool: True if grids are equal, False otherwise.

        Raises:
            NotImplementedError: If other is not a Grid.
        """
        if not isinstance(other, Grid):
            raise NotImplementedError
        return np.array_equal(self._array, other._array)

    def __sub__(self, other: object) -> int:
        """
        Count the number of differing cells between two grids.

        Args:
            other (object): Another Grid instance.

        Returns:
            int: Number of differing cells.

        Raises:
            NotImplementedError: If other is not a Grid or shapes differ.
        """
        if not isinstance(other, Grid):
            raise NotImplementedError
        if self.shape != other.shape:
            raise NotImplementedError
        return int((self._array != other._array).sum())

    def __hash__(self) -> int:
        """
        Compute a hash of the grid's contents.

        Returns:
            int: Hash value.
        """
        h = sha256()
        h.update(repr(self._array.shape).encode())
        h.update(str(self._array.dtype).encode())
        h.update(self._array.tobytes())
        return int(h.hexdigest(), 16)

    def __repr__(self) -> str:
        """
        Return an ANSI-colored string representation of the grid.

        Returns:
            str: ANSI-colored grid. `<Grid (Not intalized)>` if the grid is not initialized(initialized with `None`).
        """
        if self._array is None:
            return "<Grid (Not initialized)>"
        return "\n".join(
            "".join(ANSI_PALETTE[Symbol(v)] for v in r) for r in self.to_list()
        )

    def __str__(self) -> str:
        """
        Return a plain string representation of the grid.

        Returns:
            str: Plain grid as space-separated values.
        """
        return "\n".join(" ".join(map(str, row)) for row in self.to_list())

    def _repr_html_(self) -> str:
        """
        Return an HTML representation of the grid for Jupyter/IPython.

        Returns:
            str: HTML string.
        """
        return (
            f'<div class="grid" style="display: grid; grid-template-columns: repeat({self.shape[0]}, 1rem); grid-template-rows: repeat({self.shape[1]}, 1rem); width: fit-content; height: fit-content; margin: auto;">'
            f"{
                ' '.join(
                    f'<div class="cell" data-x="{x}" data-y="{y}" data-symbol="{value}" style="width: 1rem; height: 1rem; background-color: {CSS_PALETTE[Symbol(value)]}; border: solid 1px dimgrey"></div>'
                    for y, row in enumerate(self.to_list())
                    for x, value in enumerate(row)
                )
            }"
            "</div>"
        )
