import numpy as np
import os
from PIL import Image
import cv2
from typing import Dict, Tuple

from logic.validation import ValidationError, ensure_file_exists

_TILE_SOURCE_CACHE: Dict[str, list] = {}
_TILE_RESIZE_CACHE: Dict[Tuple[str, int], Tuple[np.ndarray, np.ndarray]] = {}

def calculate_average_color(image: np.ndarray) -> np.ndarray:
    """Calculate the average BGR color of an image.

    Args:
        image (np.ndarray): Tile image in BGR format.

    Returns:
        np.ndarray: Average color ``[b, g, r]``.
    """
    return np.mean(image.reshape(-1, 3), axis=0)

def color_distance(color1: np.ndarray, color2: np.ndarray) -> float:
    """Return the Euclidean distance between two colors.

    Args:
        color1 (np.ndarray): First color vector.
        color2 (np.ndarray): Second color vector.

    Returns:
        float: Euclidean distance.
    """
    return float(np.linalg.norm(color1 - color2))

def _ensure_tile_sources(tile_dir: str):
    """Load base BGR tiles from disk once per directory.

    Args:
        tile_dir (str): Directory containing tile imagery.

    Returns:
        list[np.ndarray]: List of base-resolution BGR tiles.
    """
    if tile_dir in _TILE_SOURCE_CACHE:
        return _TILE_SOURCE_CACHE[tile_dir]
    base_tiles = []
    ensure_file_exists(tile_dir, "Tile directory")
    for f in sorted(os.listdir(tile_dir)):
        if not f.lower().endswith(('.png', '.jpg', '.jpeg', '.png', '.webp', '.bmp')):
            continue
        path = os.path.join(tile_dir, f)
        try:
            img = Image.open(path).convert('RGB')
            tile_array = np.array(img)
            tile_bgr = cv2.cvtColor(tile_array, cv2.COLOR_RGB2BGR)
            base_tiles.append(tile_bgr)
        except Exception as exc:
            print(f"Error loading {path}: {exc}")
            continue
    _TILE_SOURCE_CACHE[tile_dir] = base_tiles
    return base_tiles


def _ensure_resized_tiles(tile_dir: str, cell_size: int):
    """Resize cached tiles to a particular cell size.

    Args:
        tile_dir (str): Source directory.
        cell_size (int): Desired tile size.

    Returns:
        tuple[np.ndarray, np.ndarray]: Resized tiles and their colors.
    """
    key = (tile_dir, cell_size)
    if key in _TILE_RESIZE_CACHE:
        return _TILE_RESIZE_CACHE[key]
    base_tiles = _ensure_tile_sources(tile_dir)
    resized_tiles = []
    tile_colors = []
    for tile in base_tiles:
        interpolation = cv2.INTER_AREA if tile.shape[0] >= cell_size else cv2.INTER_CUBIC
        resized = cv2.resize(tile, (cell_size, cell_size), interpolation=interpolation)
        resized_tiles.append(resized)
        tile_colors.append(calculate_average_color(resized))
    if not resized_tiles:
        # fallback synthetic tile if directory empty
        color = np.array([128, 128, 128], dtype=np.uint8)
        dummy = np.full((cell_size, cell_size, 3), color, dtype=np.uint8)
        resized_tiles.append(dummy)
        tile_colors.append(color.astype(float))
    cache_value = (np.stack(resized_tiles, axis=0), np.vstack(tile_colors))
    _TILE_RESIZE_CACHE[key] = cache_value
    return cache_value


def load_tile_images(tile_dir: str, n_colors: int, cell_size: int):
    """Load tiles, raising informative errors when assets are missing.

    Args:
        tile_dir (str): Directory containing tile images.
        n_colors (int): Number of representative colors required.
        cell_size (int): Target cell dimension in pixels.

    Returns:
        tuple[np.ndarray, np.ndarray]: Arrays of tiles and their average colors.

    Raises:
        ValidationError: If no tiles could be produced.
    """
    tiles_arr, tile_colors = _ensure_resized_tiles(tile_dir, cell_size)
    if len(tiles_arr) < n_colors:
        deficit = n_colors - len(tiles_arr)
        if deficit > 0:
            random_tiles = []
            random_colors = []
            for _ in range(deficit):
                color = np.random.randint(0, 255, 3, dtype=np.uint8)
                tile = np.full((cell_size, cell_size, 3), color, dtype=np.uint8)
                random_tiles.append(tile)
                random_colors.append(color.astype(float))
            tiles_arr = np.concatenate([tiles_arr, np.stack(random_tiles)], axis=0)
            tile_colors = np.concatenate([tile_colors, np.vstack(random_colors)], axis=0)
    return tiles_arr, tile_colors

def map_tiles_to_grid(
    grid_labels: np.ndarray,
    tiles: np.ndarray,
    cell_size: int,
    color_centers: np.ndarray | None = None,
    tile_colors: np.ndarray | None = None,
) -> np.ndarray:
    """Assemble the mosaic image by replacing each cell with a tile.

    Args:
        grid_labels (np.ndarray): ``grid_size x grid_size`` indices.
        tiles (np.ndarray): Stack of tile images.
        cell_size (int): Tile dimension in pixels.
        color_centers (np.ndarray | None): Optional color palette for matching.
        tile_colors (np.ndarray | None): Average colors per tile.

    Returns:
        np.ndarray: Mosaic image in BGR format.
    """
    grid_size = grid_labels.shape[0]
    tiles_arr = np.asarray(tiles)

    if color_centers is not None and tile_colors is not None:
        color_centers = np.asarray(color_centers)
        tile_colors = np.asarray(tile_colors)
        distances = np.linalg.norm(
            color_centers[:, np.newaxis, :] - tile_colors[np.newaxis, :, :], axis=2
        )
        color_to_tile = np.argmin(distances, axis=1)
        tile_indices = color_to_tile[grid_labels]
    else:
        max_index = len(tiles_arr) - 1
        tile_indices = np.clip(grid_labels, 0, max_index)

    flat_tiles = tiles_arr[tile_indices.reshape(-1)]
    tile_grid = flat_tiles.reshape(grid_size, grid_size, cell_size, cell_size, 3)
    mosaic = tile_grid.transpose(0, 2, 1, 3, 4).reshape(grid_size * cell_size, grid_size * cell_size, 3)
    return mosaic