import torch
from typing import Optional, Tuple

from .libs.process import compute_chunk_size

# ============================================================
# Utility: safe clamp
# ============================================================

def _clamp01(x):
    return torch.clamp(x, 0.0, 1.0)

# ============================================================
# Core GIMP-like Levels transfer (per-channel)
# ============================================================

def _levels_transfer(
    v: torch.Tensor,
    in_low: float,
    in_high: float,
    gamma: float,
    out_low: float,
    out_high: float,
    clamp_input: bool,
    clamp_output: bool,
):
    denom = (in_high - in_low)
    if abs(denom) < 1e-8:
        denom = 1e-8

    v = (v - in_low) / denom

    if clamp_input:
        v = _clamp01(v)

    gamma = max(gamma, 1e-6)
    v = torch.pow(torch.clamp(v, min=0.0), 1.0 / gamma)

    v = out_low + v * (out_high - out_low)

    if clamp_output:
        v = _clamp01(v)

    return v

# ============================================================
# HSV conversion (needed for Value channel mode like GIMP)
# ============================================================

def _rgb_to_hsv(img: torch.Tensor):
    r, g, b = img[..., 0], img[..., 1], img[..., 2]
    maxc, _ = torch.max(img[..., :3], dim=-1)
    minc, _ = torch.min(img[..., :3], dim=-1)
    v = maxc
    delta = maxc - minc
    s = torch.where(maxc > 0, delta / (maxc + 1e-8), torch.zeros_like(maxc))
    hue = torch.zeros_like(maxc)
    mask = delta > 1e-8
    rc = (maxc - r) / (delta + 1e-8)
    gc = (maxc - g) / (delta + 1e-8)
    bc = (maxc - b) / (delta + 1e-8)
    hue = torch.where(
        (maxc == r) & mask,
        (bc - gc),
        hue
    )
    hue = torch.where(
        (maxc == g) & mask,
        2.0 + (rc - bc),
        hue
    )
    hue = torch.where(
        (maxc == b) & mask,
        4.0 + (gc - rc),
        hue
    )
    hue = (hue / 6.0) % 1.0
    return torch.stack([hue, s, v], dim=-1)

def _hsv_to_rgb(hsv: torch.Tensor):
    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
    h = (h % 1.0) * 6.0
    i = torch.floor(h)
    f = h - i
    p = v * (1.0 - s)
    q = v * (1.0 - s * f)
    t = v * (1.0 - s * (1.0 - f))
    i = i.long() % 6
    r = torch.zeros_like(v)
    g = torch.zeros_like(v)
    b = torch.zeros_like(v)
    cond = i == 0
    r = torch.where(cond, v, r)
    g = torch.where(cond, t, g)
    b = torch.where(cond, p, b)
    cond = i == 1
    r = torch.where(cond, q, r)
    g = torch.where(cond, v, g)
    b = torch.where(cond, p, b)
    cond = i == 2
    r = torch.where(cond, p, r)
    g = torch.where(cond, v, g)
    b = torch.where(cond, t, b)
    cond = i == 3
    r = torch.where(cond, p, r)
    g = torch.where(cond, q, g)
    b = torch.where(cond, v, b)
    cond = i == 4
    r = torch.where(cond, t, r)
    g = torch.where(cond, p, g)
    b = torch.where(cond, v, b)
    cond = i == 5
    r = torch.where(cond, v, r)
    g = torch.where(cond, p, g)
    b = torch.where(cond, q, b)
    return torch.stack([r, g, b], dim=-1)

# ============================================================
# Histogram (GIMP-style endpoint detection)
# ============================================================

def _histogram_bounds(channel: torch.Tensor):
    flat = channel.reshape(-1)
    flat = flat[torch.isfinite(flat)]
    if flat.numel() == 0:
        return 0.0, 1.0
    lo = torch.min(flat)
    hi = torch.max(flat)
    return lo.item(), hi.item()

# ============================================================
# Mask blending (GEGL-style)
# ============================================================

def _apply_mask(src: torch.Tensor, dst: torch.Tensor, mask: Optional[torch.Tensor]):
    if mask is None:
        return dst
    if mask.dim() == 3:
        mask = mask.unsqueeze(-1)
    mask = mask.to(dtype=src.dtype, device=src.device)
    return src * (1.0 - mask) + dst * mask

# ============================================================
# Apply per-channel Levels
# ============================================================

def _apply_levels(
    img: torch.Tensor,
    channel: str,
    in_low: float,
    in_high: float,
    gamma: float,
    out_low: float,
    out_high: float,
    clamp_input: bool,
    clamp_output: bool,
    auto_input: bool,
    mask: Optional[torch.Tensor],
):
    B, H, W, C = img.shape

    if auto_input:
        if channel == "Value":
            v = img[..., :3].max(dim=-1).values
            in_low, in_high = _histogram_bounds(v)
        elif channel == "Red":
            in_low, in_high = _histogram_bounds(img[..., 0])
        elif channel == "Green":
            in_low, in_high = _histogram_bounds(img[..., 1])
        elif channel == "Blue":
            in_low, in_high = _histogram_bounds(img[..., 2])
        elif channel == "Alpha" and C == 4:
            in_low, in_high = _histogram_bounds(img[..., 3])

    if channel == "Value":
        hsv = _rgb_to_hsv(img[..., :3])
        v = hsv[..., 2]
        v2 = _levels_transfer(
            v,
            in_low, in_high,
            gamma,
            out_low, out_high,
            clamp_input,
            clamp_output,
        )
        hsv[..., 2] = v2
        rgb = _hsv_to_rgb(hsv)
        out = img.clone()
        out[..., :3] = rgb
        if C == 4:
            out[..., 3] = img[..., 3]
        return out

    out = img.clone()
    channel_map = {
        "Red": 0,
        "Green": 1,
        "Blue": 2,
        "Alpha": 3,
    }
    if channel in channel_map:
        c = channel_map[channel]
        if c < C:
            processed = _levels_transfer(
                img[..., c],
                in_low, in_high,
                gamma,
                out_low, out_high,
                clamp_input,
                clamp_output,
            )
            out[..., c] = processed

    return _apply_mask(img, out, mask)

# ============================================================
# FINAL OPERATOR ENTRY
# ============================================================

def exec_operator_color_levels(
    image: torch.Tensor,
    channel: str,
    interpolation: str,  # kept for API compatibility (ignored like GIMP Levels)
    input_low: float,
    input_high: float,
    input_gamma: float,
    input_clamp: bool,
    output_low: float,
    output_high: float,
    output_clamp: bool,
    auto_input: bool,
    mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    Apply a GIMP‑style Levels adjustment to an image tensor.

    The adjustment remaps the selected channel (or the Value component of HSV)
    using input low/high, gamma, and output low/high. The operation supports
    batch processing with an optional per‑image mask. Input and output tensors
    must be in RGB or RGBA format with values in [0,1].

    Args:
        image: Input tensor of shape (B, H, W, C) with C=3 (RGB) or C=4 (RGBA).
        channel: Channel to modify: "Value", "Red", "Green", "Blue", or "Alpha".
        interpolation: Unused, kept for API compatibility.
        input_low, input_high: Input range in UI units (0‑255).
        input_gamma: Gamma correction factor (>0).
        input_clamp: Clamp normalized values to [0,1] before gamma.
        output_low, output_high: Output range in UI units (0‑255).
        output_clamp: Clamp final values to [0,1].
        auto_input: If True, compute input low/high from the image histogram.
        mask: Optional mask tensor of shape (B, H, W) or (B, H, W, 1) to blend
              the original and adjusted images.

    Returns:
        Adjusted image tensor with the same shape and dtype as the input.
    """
    # ------------------------------------------------------------------
    # Input validation
    # ------------------------------------------------------------------
    if not isinstance(image, torch.Tensor):
        raise TypeError("image must be a torch.Tensor")
    if image.dim() != 4:
        raise ValueError("image must have 4 dimensions [B, H, W, C]")
    B, H, W, C = image.shape
    if C not in (3, 4):
        raise ValueError("image must have 3 or 4 channels (RGB or RGBA)")

    if mask is not None:
        if not isinstance(mask, torch.Tensor):
            raise TypeError("mask must be a torch.Tensor")
        if mask.dim() == 3:
            mask = mask.unsqueeze(-1)          # [B, H, W, 1]
        elif mask.dim() == 4 and mask.shape[-1] == 1:
            pass
        else:
            raise ValueError("mask must have shape [B, H, W] or [B, H, W, 1]")
        if mask.shape[0] != B:
            raise ValueError(
                f"mask batch size {mask.shape[0]} does not match image batch size {B}"
            )
        mask = mask.to(device=image.device, dtype=image.dtype)

    # ------------------------------------------------------------------
    # Early exit: identity transformation
    # ------------------------------------------------------------------
    if (not auto_input and
        input_low == 0 and input_high == 255 and input_gamma == 1.0 and
        output_low == 0 and output_high == 255 and
        not input_clamp and not output_clamp):
        return image.clone()

    # ------------------------------------------------------------------
    # Normalize parameters to [0,1]
    # ------------------------------------------------------------------
    in_low = input_low / 255.0
    in_high = input_high / 255.0
    out_low = output_low / 255.0
    out_high = output_high / 255.0
    gamma = max(float(input_gamma), 1e-6)

    # ------------------------------------------------------------------
    # Auto‑input: compute global histogram bounds on the whole batch
    # ------------------------------------------------------------------
    if auto_input:
        if channel == "Value":
            v = image[..., :3].max(dim=-1).values
            in_low, in_high = _histogram_bounds(v)
        elif channel == "Red":
            in_low, in_high = _histogram_bounds(image[..., 0])
        elif channel == "Green":
            in_low, in_high = _histogram_bounds(image[..., 1])
        elif channel == "Blue":
            in_low, in_high = _histogram_bounds(image[..., 2])
        elif channel == "Alpha" and C == 4:
            in_low, in_high = _histogram_bounds(image[..., 3])
        else:
            # Channel not available: keep default [0,1] (no effect)
            in_low, in_high = 0.0, 1.0

    # ------------------------------------------------------------------
    # Chunked processing for VRAM management
    # ------------------------------------------------------------------
    chunk_size = compute_chunk_size(image, mask)
    if chunk_size is None or chunk_size <= 0:
        chunk_size = B

    outputs = []
    for start in range(0, B, chunk_size):
        end = min(start + chunk_size, B)
        img_chunk = image[start:end]
        mask_chunk = mask[start:end] if mask is not None else None

        out_chunk = _apply_levels(
            img_chunk,
            channel=channel,
            in_low=in_low,
            in_high=in_high,
            gamma=gamma,
            out_low=out_low,
            out_high=out_high,
            clamp_input=input_clamp,
            clamp_output=output_clamp,
            auto_input=False,          # already computed globally
            mask=mask_chunk,
        )
        outputs.append(out_chunk)

    if outputs:
        result = torch.cat(outputs, dim=0)
    else:
        result = image   # should not happen for non‑empty batch

    return result.to(image.dtype).clamp(0.0, 1.0)