"""
duplicate_detector.py
----------------------
Two layers of duplicate detection:

  1. Exact duplicates -- SHA256 of the raw PNG-encoded pixel bytes.
  2. Visual (near) duplicates -- a perceptual hash (dHash, 64-bit) with
     Hamming-distance comparison. dHash is used instead of average-hash
     because it is far more robust to the flat color blocks typical of
     pixel art (aHash tends to false-positive heavily on sprites that
     share a dominant background/body color).

No external `imagehash` dependency is required -- the hash is implemented
directly with numpy/Pillow (both already required elsewhere in this
project), which also keeps the Space's dependency footprint smaller.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from PIL import Image

from config import PHASH_DUPLICATE_DISTANCE_FRACTION, PHASH_SIZE
from image_utils import sha256_bytes

TOTAL_HASH_BITS = PHASH_SIZE * PHASH_SIZE + 4 * 4 * 3  # structural + color bits


@dataclass
class DuplicateMatch:
    sprite_id: str
    exact: bool
    hamming_distance: int
    similarity_percent: float


def compute_sha256(png_bytes: bytes) -> str:
    return sha256_bytes(png_bytes)


def compute_dhash(rgba: np.ndarray, hash_size: int = PHASH_SIZE) -> str:
    """
    Perceptual hash combining two signals so that visually distinct
    sprites of the same *shape* (e.g. two same-silhouette plants in
    different colors) are not conflated:

      1. Structural dHash (grayscale gradient hash, `hash_size x hash_size`
         bits): captures the silhouette/shape/light-dark edge pattern.
      2. Coarse color hash (4x4 average-color grid, 3 bits per cell --
         one bit per RGB channel vs. that channel's sprite-wide median):
         captures roughly *which* colors dominate the sprite, cheaply.

    Both resizes are for HASH COMPUTATION ONLY -- never written to disk or
    used as the exported sprite -- so this does not violate the
    pixel-preservation rule, which applies to dataset *output*.

    Returns a hex string encoding the combined bit string.
    """
    if rgba.shape[0] == 0 or rgba.shape[1] == 0:
        return "0" * ((hash_size * hash_size + 4 * 4 * 3) // 4 + 1)

    img = Image.fromarray(rgba, mode="RGBA")
    background = Image.new("RGBA", img.size, (128, 128, 128, 255))
    composited = Image.alpha_composite(background, img)

    # --- structural component (grayscale gradient hash) ---
    gray = composited.convert("L")
    small_gray = gray.resize((hash_size + 1, hash_size), resample=Image.NEAREST)
    gray_pixels = np.asarray(small_gray, dtype=np.int16)
    structural_bits = (gray_pixels[:, 1:] > gray_pixels[:, :-1]).flatten()

    # --- color component (coarse 4x4 average-color grid, absolute-referenced) ---
    # Compared against a fixed midpoint (128) rather than the sprite's own
    # median: a self-referenced median captures "brighter/darker than this
    # sprite's own average", which duplicates the structural hash's job and
    # collapses to the identical bit pattern for e.g. a green sprite and a
    # magenta sprite of the same silhouette on the same background (their
    # own two-tone medians make every foreground pixel "above median" and
    # every background pixel "below" regardless of which color it actually
    # is). An absolute reference keeps the hash sensitive to hue and channel
    # balance, which is exactly the visual duplicate signal color is meant
    # to add on top of the structural hash.
    rgb_img = composited.convert("RGB")
    grid = 4
    small_color = rgb_img.resize((grid, grid), resample=Image.NEAREST)
    color_pixels = np.asarray(small_color, dtype=np.int16).reshape(-1, 3)  # (16, 3)
    color_bits = (color_pixels > 128).flatten()  # 16*3 = 48 bits

    all_bits = np.concatenate([structural_bits, color_bits])

    value = 0
    for bit in all_bits:
        value = (value << 1) | int(bit)
    hex_digits = (len(all_bits) + 3) // 4
    return format(value, f"0{hex_digits}x")


def hamming_distance(hash_a: str, hash_b: str) -> int:
    int_a = int(hash_a, 16)
    int_b = int(hash_b, 16)
    return bin(int_a ^ int_b).count("1")


def find_duplicates(
    sha256_hash: str,
    dhash: str,
    existing: list[dict],
    distance_threshold: int | None = None,
) -> list[DuplicateMatch]:
    """
    `existing` is a list of dicts with at least {"id", "sha256", "phash"}
    representing already-accepted or already-staged sprites to compare
    against. `distance_threshold` defaults to
    PHASH_DUPLICATE_DISTANCE_FRACTION * TOTAL_HASH_BITS when not given.
    """
    if distance_threshold is None:
        distance_threshold = round(PHASH_DUPLICATE_DISTANCE_FRACTION * TOTAL_HASH_BITS)

    matches: list[DuplicateMatch] = []
    for item in existing:
        if item.get("sha256") == sha256_hash:
            matches.append(DuplicateMatch(sprite_id=item["id"], exact=True, hamming_distance=0, similarity_percent=100.0))
            continue
        other_hash = item.get("phash")
        if not other_hash:
            continue
        dist = hamming_distance(dhash, other_hash)
        if dist <= distance_threshold:
            similarity = 100.0 * (1 - dist / max(1, TOTAL_HASH_BITS))
            matches.append(
                DuplicateMatch(sprite_id=item["id"], exact=False, hamming_distance=dist, similarity_percent=similarity)
            )
    matches.sort(key=lambda m: (not m.exact, m.hamming_distance))
    return matches