from dataclasses import dataclass
from PIL import Image
import importlib


@dataclass(frozen=True)
class Color:
    RED: int
    GREEN: int
    BLUE: int


@dataclass(frozen=True)
class ColorMask:
    target_color: Color
    replacement_color: Color


def load_color_map(name: str):
    module_path = f"tools.recolor.color_maps.{name}"
    module = importlib.import_module(module_path)

    if not hasattr(module, "COLOR_MAP"):
        raise RuntimeError(f"color map '{name}' has no COLOR_MAP")

    return module.COLOR_MAP


def recolor_image(img: Image.Image, color_map: list[ColorMask]) -> Image.Image:
    img = img.convert("RGBA")
    pixels = img.load()

    lookup = {
        (m.target_color.RED, m.target_color.GREEN, m.target_color.BLUE):
        (m.replacement_color.RED, m.replacement_color.GREEN, m.replacement_color.BLUE)
        for m in color_map
    }

    width, height = img.size

    for y in range(height):
        for x in range(width):
            r, g, b, a = pixels[x, y]
            if (r, g, b) in lookup:
                nr, ng, nb = lookup[(r, g, b)]
                pixels[x, y] = (nr, ng, nb, a)

    return img


def extract_unique_colors(img: Image.Image) -> set[Color]:
    img = img.convert("RGBA")
    pixels = img.load()

    width, height = img.size
    colors: set[Color] = set()

    for y in range(height):
        for x in range(width):
            r, g, b, a = pixels[x, y]
            if a != 0:
                colors.add(Color(r, g, b))

    return colors
