"""Minimal image processing pipeline adapted for the Qt frontend."""

from __future__ import annotations

import colorsys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, Tuple

import numpy as np
from PIL import Image, ImageDraw
from PySide6 import QtGui

from app.logic import PREVIEW_MAX_SIZE


@dataclass
class Stats:
    matches_all: int = 0
    total_all: int = 0
    matches_keep: int = 0
    total_keep: int = 0
    matches_excl: int = 0
    total_excl: int = 0

    def summary(self, translate) -> str:
        if self.total_all == 0:
            return translate("stats.placeholder")
        with_pct = (self.matches_keep / self.total_keep * 100) if self.total_keep else 0.0
        without_pct = (self.matches_all / self.total_all * 100) if self.total_all else 0.0
        excluded_pct = (self.total_excl / self.total_all * 100) if self.total_all else 0.0
        excluded_match_pct = (self.matches_excl / self.total_excl * 100) if self.total_excl else 0.0
        return translate(
            "stats.summary",
            with_pct=with_pct,
            without_pct=without_pct,
            excluded_pct=excluded_pct,
            excluded_match_pct=excluded_match_pct,
        )


def _rgb_to_hsv_numpy(arr: np.ndarray) -> np.ndarray:
    """Vectorized RGB→HSV conversion. arr shape: (H, W, 3), dtype float32, range [0,1].
    Returns array of same shape with channels [H(0-360), S(0-100), V(0-100)].
    """
    r = arr[..., 0]
    g = arr[..., 1]
    b = arr[..., 2]

    cmax = np.maximum(np.maximum(r, g), b)
    cmin = np.minimum(np.minimum(r, g), b)
    delta = cmax - cmin

    # Value
    v = cmax

    # Saturation
    s = np.zeros_like(r)
    np.divide(delta, cmax, out=s, where=cmax > 0)

    # Hue
    h = np.zeros_like(r)
    mask_r = (delta > 0) & (cmax == r)
    mask_g = (delta > 0) & (cmax == g)
    mask_b = (delta > 0) & (cmax == b)
    h[mask_r] = (60.0 * ((g[mask_r] - b[mask_r]) / delta[mask_r])) % 360.0
    h[mask_g] = (60.0 * ((b[mask_g] - r[mask_g]) / delta[mask_g]) + 120.0) % 360.0
    h[mask_b] = (60.0 * ((r[mask_b] - g[mask_b]) / delta[mask_b]) + 240.0) % 360.0

    return np.stack([h, s * 100.0, v * 100.0], axis=-1)


class QtImageProcessor:
    """Process images and build overlays for the Qt UI."""

    def __init__(self) -> None:
        self.orig_img: Image.Image | None = None
        self.preview_img: Image.Image | None = None
        self.overlay_img: Image.Image | None = None
        self.preview_paths: list[Path] = []
        self.current_index: int = -1
        self.stats = Stats()

        # Overlay tint color
        self.overlay_r = 255
        self.overlay_g = 0
        self.overlay_b = 0

        self.defaults: Dict[str, int] = {
            "hue_min": 0,
            "hue_max": 360,
            "sat_min": 25,
            "val_min": 15,
            "val_max": 100,
            "alpha": 120,
        }
        self.hue_min = self.defaults["hue_min"]
        self.hue_max = self.defaults["hue_max"]
        self.sat_min = self.defaults["sat_min"]
        self.val_min = self.defaults["val_min"]
        self.val_max = self.defaults["val_max"]
        self.alpha = self.defaults["alpha"]

        self.exclude_shapes: list[dict[str, object]] = []
        self.reset_exclusions_on_switch: bool = False

        # Mask caching
        self._cached_mask: np.ndarray | None = None
        self._cached_mask_size: Tuple[int, int] | None = None
        self.exclude_ref_size: Tuple[int, int] | None = None

    def set_defaults(self, defaults: dict) -> None:
        for key in self.defaults:
            if key in defaults:
                self.defaults[key] = int(defaults[key])
        for key, value in self.defaults.items():
            setattr(self, key, value)
        self._rebuild_overlay()

    # thresholds -------------------------------------------------------------

    def set_threshold(self, key: str, value: int) -> None:
        setattr(self, key, value)
        if self.preview_img is not None:
            self._rebuild_overlay()

    # image handling --------------------------------------------------------

    def load_single_image(self, path: Path, *, reset_collection: bool = True) -> Path:
        image = Image.open(path).convert("RGBA")
        self.orig_img = image
        if reset_collection:
            self.preview_paths = [path]
            self.current_index = 0
        self._build_preview()
        self._rebuild_overlay()
        return path

    def load_folder(self, paths: Iterable[Path], start_index: int = 0) -> Path:
        self.preview_paths = list(paths)
        if not self.preview_paths:
            raise ValueError("No images in folder.")
        self.current_index = max(0, min(start_index, len(self.preview_paths) - 1))
        return self._load_image_at_current()

    def next_image(self) -> Path | None:
        if not self.preview_paths:
            return None
        self.current_index = (self.current_index + 1) % len(self.preview_paths)
        return self._load_image_at_current()

    def previous_image(self) -> Path | None:
        if not self.preview_paths:
            return None
        self.current_index = (self.current_index - 1) % len(self.preview_paths)
        return self._load_image_at_current()

    def _load_image_at_current(self) -> Path:
        path = self.preview_paths[self.current_index]
        return self.load_single_image(path, reset_collection=False)

    # preview/overlay -------------------------------------------------------

    def _build_preview(self) -> None:
        if self.orig_img is None:
            self.preview_img = None
            return
        width, height = self.orig_img.size
        max_w, max_h = PREVIEW_MAX_SIZE
        scale = min(max_w / width, max_h / height)
        if scale <= 0:
            scale = 1.0
        size = (max(1, int(width * scale)), max(1, int(height * scale)))
        self.preview_img = self.orig_img.resize(size, Image.LANCZOS)

    def _rebuild_overlay(self) -> None:
        """Build color-match overlay using vectorized NumPy operations."""
        if self.preview_img is None:
            self.overlay_img = None
            self.stats = Stats()
            return

        base = self.preview_img.convert("RGBA")
        arr = np.asarray(base, dtype=np.float32)  # (H, W, 4)

        rgb = arr[..., :3] / 255.0
        alpha_ch = arr[..., 3]                     # alpha channel of the image

        hsv = _rgb_to_hsv_numpy(rgb)               # (H, W, 3): H°, S%, V%

        hue = hsv[..., 0]
        sat = hsv[..., 1]
        val = hsv[..., 2]

        hue_min = float(self.hue_min)
        hue_max = float(self.hue_max)
        if hue_min <= hue_max:
            hue_ok = (hue >= hue_min) & (hue <= hue_max)
        else:
            hue_ok = (hue >= hue_min) | (hue <= hue_max)

        match_mask = (
            hue_ok
            & (sat >= float(self.sat_min))
            & (val >= float(self.val_min))
            & (val <= float(self.val_max))
            & (alpha_ch > 0)
        )

        # Exclusion mask (same pixel space as preview)
        excl_mask = self._build_exclusion_mask_numpy(base.size)  # bool (H,W)

        keep_match = match_mask & ~excl_mask
        excl_match = match_mask & excl_mask
        visible = alpha_ch > 0

        matches_all = int(match_mask[visible].sum())
        total_all   = int(visible.sum())
        matches_keep = int(keep_match[visible].sum())
        total_keep   = int((visible & ~excl_mask).sum())
        matches_excl = int(excl_match[visible].sum())
        total_excl   = int((visible & excl_mask).sum())

        # Build overlay image
        overlay_arr = np.zeros((base.height, base.width, 4), dtype=np.uint8)
        overlay_arr[keep_match, 0] = self.overlay_r
        overlay_arr[keep_match, 1] = self.overlay_g
        overlay_arr[keep_match, 2] = self.overlay_b
        overlay_arr[keep_match, 3] = int(self.alpha)

        self.overlay_img = Image.fromarray(overlay_arr, "RGBA")
        self.stats = Stats(
            matches_all=matches_all,
            total_all=total_all,
            matches_keep=matches_keep,
            total_keep=total_keep,
            matches_excl=matches_excl,
            total_excl=total_excl,
        )

    def get_stats_headless(self, image: Image.Image) -> Stats:
        """Calculate color-match statistics natively without building UI elements or scaling."""
        base = image.convert("RGBA")
        arr = np.asarray(base, dtype=np.float32)

        rgb = arr[..., :3] / 255.0
        alpha_ch = arr[..., 3]                     

        hsv = _rgb_to_hsv_numpy(rgb)               

        hue = hsv[..., 0]
        sat = hsv[..., 1]
        val = hsv[..., 2]

        hue_min = float(self.hue_min)
        hue_max = float(self.hue_max)
        if hue_min <= hue_max:
            hue_ok = (hue >= hue_min) & (hue <= hue_max)
        else:
            hue_ok = (hue >= hue_min) | (hue <= hue_max)

        match_mask = (
            hue_ok
            & (sat >= float(self.sat_min))
            & (val >= float(self.val_min))
            & (val <= float(self.val_max))
            & (alpha_ch > 0)
        )

        excl_mask = self._build_exclusion_mask_numpy(base.size)  

        keep_match = match_mask & ~excl_mask
        excl_match = match_mask & excl_mask
        visible = alpha_ch > 0

        return Stats(
            matches_all=int(match_mask[visible].sum()),
            total_all=int(visible.sum()),
            matches_keep=int(keep_match[visible].sum()),
            total_keep=int((visible & ~excl_mask).sum()),
            matches_excl=int(excl_match[visible].sum()),
            total_excl=int((visible & excl_mask).sum()),
        )

    # helpers ----------------------------------------------------------------

    def _matches(self, r: int, g: int, b: int) -> bool:
        """Single-pixel match — kept for compatibility / eyedropper use."""
        h, s, v = colorsys.rgb_to_hsv(r / 255.0, g / 255.0, b / 255.0)
        hue = (h * 360.0) % 360.0
        if self.hue_min <= self.hue_max:
            hue_ok = self.hue_min <= hue <= self.hue_max
        else:
            hue_ok = hue >= self.hue_min or hue <= self.hue_max
        sat_ok = s * 100.0 >= self.sat_min
        val_ok = self.val_min <= v * 100.0 <= self.val_max
        return hue_ok and sat_ok and val_ok

    def pick_color(self, x: int, y: int) -> Tuple[float, float, float] | None:
        """Return (hue°, sat%, val%) of the preview pixel at (x, y), or None."""
        if self.preview_img is None:
            return None
        img = self.preview_img.convert("RGBA")
        try:
            r, g, b, a = img.getpixel((x, y))
        except IndexError:
            return None
        if a == 0:
            return None
        h, s, v = colorsys.rgb_to_hsv(r / 255.0, g / 255.0, b / 255.0)
        return (h * 360.0) % 360.0, s * 100.0, v * 100.0

    # exported data ----------------------------------------------------------

    def preview_pixmap(self) -> QtGui.QPixmap:
        return self._to_pixmap(self.preview_img)

    def overlay_pixmap(self) -> QtGui.QPixmap:
        if self.preview_img is None or self.overlay_img is None:
            return QtGui.QPixmap()
        merged = Image.alpha_composite(self.preview_img.convert("RGBA"), self.overlay_img)
        return self._to_pixmap(merged)

    @staticmethod
    def _to_pixmap(image: Image.Image | None) -> QtGui.QPixmap:
        if image is None:
            return QtGui.QPixmap()
        buffer = image.tobytes("raw", "RGBA")
        qt_image = QtGui.QImage(buffer, image.width, image.height, QtGui.QImage.Format_RGBA8888)
        return QtGui.QPixmap.fromImage(qt_image)

    # exclusions -------------------------------------------------------------

    def set_exclusions(self, shapes: list[dict[str, object]]) -> None:
        copied: list[dict[str, object]] = []
        for shape in shapes:
            kind = shape.get("kind")
            if kind == "rect":
                coords = tuple(shape.get("coords", (0, 0, 0, 0)))  # type: ignore[assignment]
                copied.append({"kind": "rect", "coords": tuple(int(c) for c in coords)})
            elif kind == "polygon":
                pts = shape.get("points", [])
                copied.append({"kind": "polygon", "points": [(int(x), int(y)) for x, y in pts]})
        self.exclude_shapes = copied
        if self.preview_img:
            self.exclude_ref_size = self.preview_img.size
        else:
            self.exclude_ref_size = None
            
        self._cached_mask = None  # Invalidate cache
        self._cached_mask_size = None
        self._rebuild_overlay()

    def _build_exclusion_mask(self, size: Tuple[int, int]) -> Image.Image | None:
        if not self.exclude_shapes:
            return None
        
        target_w, target_h = size
        ref_w, ref_h = self.exclude_ref_size or size
        sx = target_w / ref_w if ref_w > 0 else 1.0
        sy = target_h / ref_h if ref_h > 0 else 1.0

        mask = Image.new("L", size, 0)
        draw = ImageDraw.Draw(mask)
        for shape in self.exclude_shapes:
            kind = shape.get("kind")
            if kind == "rect":
                x0, y0, x1, y1 = shape["coords"]  # type: ignore[index]
                draw.rectangle([x0 * sx, y0 * sy, x1 * sx, y1 * sy], fill=255)
            elif kind == "polygon":
                points = shape.get("points", [])
                if len(points) >= 3:
                    scaled_pts = [(int(x * sx), int(y * sy)) for x, y in points]
                    draw.polygon(scaled_pts, fill=255)
        return mask

    def set_overlay_color(self, hex_code: str) -> None:
        """Set the RGB channels for the match overlay from a hex string."""
        if not hex_code.startswith("#") or len(hex_code) not in (7, 9):
            return
        try:
            self.overlay_r = int(hex_code[1:3], 16)
            self.overlay_g = int(hex_code[3:5], 16)
            self.overlay_b = int(hex_code[5:7], 16)
        except ValueError:
            pass

    def _build_exclusion_mask_numpy(self, size: Tuple[int, int]) -> np.ndarray:
        """Return a boolean (H, W) mask — True where pixels are excluded."""
        if self._cached_mask is not None and self._cached_mask_size == size:
            return self._cached_mask

        w, h = size
        if not self.exclude_shapes:
            mask = np.zeros((h, w), dtype=bool)
        else:
            pil_mask = self._build_exclusion_mask(size)
            if pil_mask is None:
                mask = np.zeros((h, w), dtype=bool)
            else:
                mask = np.asarray(pil_mask, dtype=bool)

        self._cached_mask = mask
        self._cached_mask_size = size
        return mask