import base64
import math
from io import BytesIO

import cv2
import numpy as np
from PIL import Image


LAPLACIAN_MASK = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]], dtype=np.float32)
SOBEL_GX = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=np.float32)
SOBEL_GY = SOBEL_GX.T
ROBERTS_GX = np.array([[1, 0], [0, -1]], dtype=np.float32)
ROBERTS_GY = np.array([[0, 1], [-1, 0]], dtype=np.float32)


class ProcessingError(ValueError):
    """Raised when an image operation receives invalid input or parameters."""

    pass


def ensure_uint8(image):
    """Clip image values to the display range [0, 255] and return uint8 data.

    This is used after arithmetic or filtering so the result can be displayed as a normal 8-bit image.
    """

    return np.clip(image, 0, 255).astype(np.uint8)


def normalize_to_uint8(image):
    """Linearly normalize any numeric image to the full 8-bit display range.

    This is useful for derivative, subtraction, and spectrum results that may contain negative or very large values.
    """

    arr = image.astype(np.float32)
    min_value = float(np.min(arr))
    max_value = float(np.max(arr))
    if math.isclose(min_value, max_value):
        return np.zeros(arr.shape, dtype=np.uint8)
    return np.round((arr - min_value) * 255.0 / (max_value - min_value)).astype(np.uint8)


def require_odd(value, name="size", minimum=3):
    """Validate that a mask size is an odd integer greater than or equal to minimum."""

    try:
        value = int(value)
    except (TypeError, ValueError) as exc:
        raise ProcessingError(f"{name} must be an odd integer.") from exc
    if value < minimum or value % 2 == 0:
        raise ProcessingError(f"{name} must be an odd integer >= {minimum}.")
    return value


def require_finite_positive(value, name):
    """Validate that a parameter is finite and strictly positive."""

    try:
        value = float(value)
    except (TypeError, ValueError) as exc:
        raise ProcessingError(f"{name} must be a finite positive number.") from exc
    if not np.isfinite(value) or value <= 0:
        raise ProcessingError(f"{name} must be a finite positive number.")
    return value


def to_gray(image):
    """Convert an RGB image to grayscale, leaving grayscale input unchanged.

    Many spatial-domain formulas work on intensity, so this gives them a single gray-level channel.
    """

    if image.ndim == 2:
        return image
    return cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)


def gray_to_rgb(gray):
    """Convert a single-channel grayscale image to RGB for consistent display.

    The frontend expects displayable RGB images even when the algorithm result is grayscale.
    """

    return cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)


def histogram(image):
    """Return the normalized intensity histogram p(r_k) for gray levels 0..255.

    Histograms are used to inspect contrast, brightness distribution, and equalization results.
    """

    gray = to_gray(image)
    counts = np.bincount(gray.ravel(), minlength=256).astype(np.float64)
    probabilities = counts / max(gray.size, 1)
    return probabilities.round(8).tolist()


def histogram_payload(image):
    """Return intensity histogram and, for RGB images, separate R/G/B histograms.

    This lets the UI explain both overall intensity and per-channel color behavior.
    """

    gray = to_gray(image)
    payload = {"intensity": histogram(gray)}
    if image.ndim == 3:
        payload["r"] = (np.bincount(image[:, :, 0].ravel(), minlength=256).astype(np.float64) / image[:, :, 0].size).round(8).tolist()
        payload["g"] = (np.bincount(image[:, :, 1].ravel(), minlength=256).astype(np.float64) / image[:, :, 1].size).round(8).tolist()
        payload["b"] = (np.bincount(image[:, :, 2].ravel(), minlength=256).astype(np.float64) / image[:, :, 2].size).round(8).tolist()
    return payload


def image_to_data_url(image):
    """Encode a uint8 image as a PNG data URL for API responses."""

    pil_image = Image.fromarray(ensure_uint8(image))
    buffer = BytesIO()
    pil_image.save(buffer, format="PNG")
    payload = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/png;base64,{payload}"


def data_url_to_bytes(value):
    """Decode a base64 data URL or raw base64 string into image bytes."""

    if "," in value:
        value = value.split(",", 1)[1]
    return base64.b64decode(value)


def decode_image(uploaded_file=None, base64_image=None):
    """Decode an uploaded file or base64 payload into a uint8 NumPy image.

    Grayscale inputs stay single-channel so intensity-only operations do not create fake RGB channels.
    """

    if uploaded_file is None and not base64_image:
        raise ProcessingError("Provide an image file or base64 image payload.")
    if uploaded_file is not None:
        raw = uploaded_file.read()
    else:
        raw = data_url_to_bytes(base64_image)

    image = Image.open(BytesIO(raw))
    if image.mode in {"1", "L", "I;16", "I", "F"}:
        return np.array(image.convert("L"), dtype=np.uint8)
    return np.array(image.convert("RGB"), dtype=np.uint8)


def negative(image, params):
    """Apply the image negative transform, s = 255 - r.

    Use it to invert bright and dark structures, which can make some details easier to see.
    """

    return 255 - image


def logarithmic(image, params):
    """Apply logarithmic intensity expansion, s = c log(1 + r), on normalized pixels.

    Use it to expand dark gray levels while compressing very bright regions.
    """

    c = require_finite_positive(params.get("c", 1.0 / math.log(2.0)), "c")
    normalized = image.astype(np.float32) / 255.0
    transformed = c * np.log1p(normalized)
    return ensure_uint8(np.round(np.clip(transformed, 0.0, 1.0) * 255.0))


def gamma(image, params):
    """Apply power-law correction, s = c r^gamma, on normalized pixels.

    Use it to brighten dark images with gamma < 1 or darken washed-out images with gamma > 1.
    """

    gamma_value = require_finite_positive(params.get("gamma", 1.0), "gamma")
    c = require_finite_positive(params.get("c", 1.0), "c")
    normalized = image.astype(np.float32) / 255.0
    transformed = c * np.power(normalized, gamma_value)
    return ensure_uint8(np.round(np.clip(transformed, 0.0, 1.0) * 255.0))


def contrast_stretch(image, params):
    """Stretch the selected gray-level interval [low, high] to the full [0, 255] range.

    Use it when useful image values occupy a narrow dynamic range and need stronger contrast.
    """

    low = int(params.get("low", 0))
    high = int(params.get("high", 255))
    if low < 0 or high > 255 or low >= high:
        raise ProcessingError("Contrast stretch requires 0 <= low < high <= 255.")
    stretched = (image.astype(np.float32) - low) * (255.0 / (high - low))
    return ensure_uint8(np.round(stretched))


def gray_slice(image, params):
    """Highlight pixels whose grayscale intensity lies inside [start, end].

    Use it to emphasize one gray-level band, such as a tissue, object, or intensity region of interest.
    """

    start = int(params.get("start", 96))
    end = int(params.get("end", 160))
    if start < 0 or end > 255 or start > end:
        raise ProcessingError("Gray-level slicing requires 0 <= start <= end <= 255.")
    preserve = bool(params.get("preserve_background", True))
    highlight = np.array(params.get("highlight", [255, 64, 64]), dtype=np.uint8)
    if highlight.shape != (3,):
        raise ProcessingError("highlight must be an RGB triplet.")
    gray = to_gray(image)
    mask = (gray >= start) & (gray <= end)
    base = image.copy() if image.ndim == 3 else gray_to_rgb(gray if preserve else np.zeros_like(gray))
    if not preserve:
        base = np.zeros((*gray.shape, 3), dtype=np.uint8)
    base[mask] = highlight
    return base


def bit_plane(image, params):
    """Extract one grayscale bit plane and display it as a binary image.

    Use it to study which bits carry the main visual information or fine/noisy details.
    """

    bit = int(params.get("bit", 7))
    if bit < 0 or bit > 7:
        raise ProcessingError("bit must be between 0 and 7.")
    plane = ((to_gray(image) >> bit) & 1) * 255
    return gray_to_rgb(plane.astype(np.uint8))


def histogram_equalization(image, params):
    """Equalize a grayscale image using the discrete cumulative distribution function.

    Use it to improve global contrast when the histogram is concentrated in a small intensity range.
    """

    grayscale_input = image.ndim == 2
    gray = to_gray(image)
    counts = np.bincount(gray.ravel(), minlength=256)
    cdf = counts.cumsum().astype(np.float64)
    nonzero = cdf[cdf > 0]
    if nonzero.size == 0:
        return gray if grayscale_input else gray_to_rgb(gray)
    cdf_min = nonzero[0]
    denom = gray.size - cdf_min
    if denom <= 0:
        equalized = np.zeros_like(gray)
    else:
        lut = np.round((cdf - cdf_min) / denom * 255.0).clip(0, 255).astype(np.uint8)
        equalized = lut[gray]
    return equalized if grayscale_input else gray_to_rgb(equalized)


def apply_kernel(image, kernel, normalize_derivative=False):
    """Apply a 2D convolution mask to each channel using reflected borders.

    This is the shared imfilter-style step behind smoothing and sharpening masks.
    """

    source = image.astype(np.float32)
    if image.ndim == 2:
        filtered = cv2.filter2D(source, cv2.CV_32F, kernel, borderType=cv2.BORDER_REFLECT)
    else:
        channels = [cv2.filter2D(source[:, :, idx], cv2.CV_32F, kernel, borderType=cv2.BORDER_REFLECT) for idx in range(source.shape[2])]
        filtered = np.stack(channels, axis=2)
    if normalize_derivative:
        return normalize_to_uint8(filtered)
    return ensure_uint8(np.round(filtered))


def filter_float(image, kernel):
    """Apply a 2D convolution mask and keep the float result for derivative math.

    Use it when intermediate negative edge/detail values must be preserved before display normalization.
    """

    source = image.astype(np.float32)
    if image.ndim == 2:
        return cv2.filter2D(source, cv2.CV_32F, kernel, borderType=cv2.BORDER_REFLECT)
    channels = [cv2.filter2D(source[:, :, idx], cv2.CV_32F, kernel, borderType=cv2.BORDER_REFLECT) for idx in range(source.shape[2])]
    return np.stack(channels, axis=2)


def box_filter(image, params):
    """Blur an image with a K x K average mask, equivalent to ones(K,K) / K^2.

    Use it for simple smoothing or reducing Gaussian-like noise, accepting that edges become softer.
    """

    size = require_odd(params.get("size", 3), "size")
    return cv2.blur(image, (size, size), borderType=cv2.BORDER_REFLECT)


def weighted_average(image, params):
    """Blur an image with a normalized weighted average mask.

    Use it for gentler smoothing that gives the center pixel more influence than a plain box filter.
    """

    size = require_odd(params.get("size", 3), "size")
    if "kernel" in params:
        kernel = np.array(params["kernel"], dtype=np.float32)
        if kernel.shape != (size, size):
            raise ProcessingError("kernel dimensions must match size.")
    elif size == 3:
        kernel = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=np.float32)
    else:
        sigma = max(size / 6.0, 0.1)
        ax = np.arange(-(size // 2), size // 2 + 1, dtype=np.float32)
        xx, yy = np.meshgrid(ax, ax)
        kernel = np.exp(-(xx**2 + yy**2) / (2.0 * sigma**2))
    total = float(np.sum(kernel))
    if math.isclose(total, 0.0):
        raise ProcessingError("weighted average kernel sum must not be zero.")
    return apply_kernel(image, kernel / total)


def median_filter(image, params):
    """Apply an order-statistic median filter for impulse-noise removal.

    Use it to remove salt-and-pepper noise while preserving edges better than linear averaging.
    """

    size = require_odd(params.get("size", 3), "size")
    return cv2.medianBlur(image, size)


def laplacian(image, params):
    """Apply a zero-sum Laplacian detail mask and optionally add it back for sharpening.

    Use it to reveal fine second-derivative detail or sharpen small structures.
    """

    mode = params.get("mode", "sharpen")
    lap = filter_float(image, LAPLACIAN_MASK)
    if mode == "edge":
        return normalize_to_uint8(lap)
    sign = params.get("sign", "add")
    source = image.astype(np.float32)
    sharpened = source + lap if sign == "add" else source - lap
    return ensure_uint8(np.round(sharpened))


def high_boost(image, params):
    """Apply high-boost filtering, f_hb = A f - blurred(f), with A >= 1.

    Use it to emphasize edges and details while retaining more of the original image than pure high-pass filtering.
    """

    amplification = float(params.get("amplification", 1.5))
    if not np.isfinite(amplification) or amplification < 1.0:
        raise ProcessingError("amplification must be >= 1.")
    size = require_odd(params.get("size", 3), "size")
    blurred = cv2.blur(image, (size, size), borderType=cv2.BORDER_REFLECT).astype(np.float32)
    boosted = amplification * image.astype(np.float32) - blurred
    return ensure_uint8(np.round(boosted))


def gradient_magnitude(image, gx_kernel, gy_kernel):
    """Compute gradient magnitude from Gx and Gy derivative masks.

    Use it to find strong first-derivative changes, which usually correspond to object edges.
    """

    gray = to_gray(image).astype(np.float32)
    gx = cv2.filter2D(gray, cv2.CV_32F, gx_kernel, borderType=cv2.BORDER_REFLECT)
    gy = cv2.filter2D(gray, cv2.CV_32F, gy_kernel, borderType=cv2.BORDER_REFLECT)
    magnitude = np.sqrt(gx**2 + gy**2)
    return gray_to_rgb(normalize_to_uint8(magnitude))


def sobel(image, params):
    """Detect edges with Sobel horizontal and vertical derivative masks.

    Use it for edge detection with some built-in smoothing from the larger 3 x 3 masks.
    """

    return gradient_magnitude(image, SOBEL_GX, SOBEL_GY)


def roberts(image, params):
    """Detect edges with Roberts cross-gradient masks.

    Use it for a simple 2 x 2 gradient operator that responds to diagonal intensity changes.
    """

    return gradient_magnitude(image, ROBERTS_GX, ROBERTS_GY)


OPERATIONS = {
    "negative": negative,
    "log": logarithmic,
    "gamma": gamma,
    "contrast_stretch": contrast_stretch,
    "gray_slice": gray_slice,
    "bit_plane": bit_plane,
    "hist_equalization": histogram_equalization,
    "box_filter": box_filter,
    "weighted_average": weighted_average,
    "median_filter": median_filter,
    "laplacian": laplacian,
    "high_boost": high_boost,
    "sobel": sobel,
    "roberts": roberts,
}


def process_image(image, operation, params=None):
    """Run one named legacy operation against an image."""

    params = params or {}
    if operation not in OPERATIONS:
        raise ProcessingError(f"Unsupported operation '{operation}'.")
    return ensure_uint8(OPERATIONS[operation](ensure_uint8(image), params))


def subtract_images(left, right):
    """Subtract two registered images and normalize the absolute difference.

    Use it for change detection between two aligned images or processing states.
    """

    verify_registration([left, right])
    diff = left.astype(np.float32) - right.astype(np.float32)
    return normalize_to_uint8(np.abs(diff))


def average_images(images):
    """Average a stack of registered images to reduce independent noise.

    Use it when multiple aligned captures of the same scene are available.
    """

    verify_registration(images)
    stack = np.stack([image.astype(np.float32) for image in images], axis=0)
    return ensure_uint8(np.round(np.mean(stack, axis=0)))


def verify_registration(images):
    """Ensure all images have identical dimensions and channel counts.

    This prevents invalid arithmetic between images that are not aligned pixel-for-pixel.
    """

    if len(images) < 2:
        raise ProcessingError("At least two registered images are required.")
    shape = images[0].shape
    if any(image.shape != shape for image in images[1:]):
        raise ProcessingError("Images must have identical width, height, and channel count.")