import torch
from typing import Tuple, Optional

from .libs.process import compute_chunk_size


def exec_operator_color_threshold(
    image: torch.Tensor,
    channel: str,
    low: float,
    high: float,
    mask: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor]:
    """
    Apply a threshold to the input image, producing a binary output.

    Supported channels:
        Value, Red, Green, Blue, Alpha, Luminance, RGB

    Image:
        [B,H,W,C] float32 RGB/RGBA in [0,1]

    Mask:
        [H,W] or [B,H,W]
    """
    try:
        if not isinstance(image, torch.Tensor):
            raise TypeError("image must be a torch.Tensor")

        if image.ndim != 4:
            raise ValueError("image must be [B,H,W,C]")

        B, H, W, C = image.shape

        if C not in (3, 4):
            raise ValueError("image must be RGB or RGBA")

        low = max(0.0, min(255.0, float(low))) / 255.0
        high = max(0.0, min(255.0, float(high))) / 255.0

        if low > high:
            low, high = high, low

        if mask is not None:
            if not isinstance(mask, torch.Tensor):
                raise TypeError("mask must be a torch.Tensor")

            if mask.ndim == 2:
                mask = mask.unsqueeze(0)
            elif mask.ndim != 3:
                raise ValueError("mask must be [H,W] or [B,H,W]")

            if mask.shape[-2:] != (H, W):
                raise ValueError("mask size mismatch")

            if mask.shape[0] not in (1, B):
                raise ValueError("mask batch mismatch")

            mask = mask.to(dtype=image.dtype, device=image.device).clamp_(0, 1)

            if torch.count_nonzero(mask) == 0:
                return (image,)

            if torch.all(mask == 1):
                mask = None

        device = image.device

        chunk_size = (
            max(1, compute_chunk_size(image, mask))
            if device.type == "cuda"
            else B
        )

        def threshold_chunk(
            img: torch.Tensor,
            msk: Optional[torch.Tensor],
        ) -> torch.Tensor:

            rgb = img[..., :3]

            if channel == "RGB":
                inside = (
                    (rgb[..., 0] >= low)
                    & (rgb[..., 0] <= high)
                    & (rgb[..., 1] >= low)
                    & (rgb[..., 1] <= high)
                    & (rgb[..., 2] >= low)
                    & (rgb[..., 2] <= high)
                )

            else:
                if channel == "Value":
                    intensity = rgb.amax(dim=-1)

                elif channel == "Red":
                    intensity = rgb[..., 0]

                elif channel == "Green":
                    intensity = rgb[..., 1]

                elif channel == "Blue":
                    intensity = rgb[..., 2]

                elif channel == "Alpha":
                    if C == 4:
                        intensity = img[..., 3]
                    else:
                        intensity = torch.ones(
                            img.shape[:-1],
                            device=img.device,
                            dtype=img.dtype,
                        )

                elif channel == "Luminance":
                    intensity = (
                        rgb[..., 0] * 0.299
                        + rgb[..., 1] * 0.587
                        + rgb[..., 2] * 0.114
                    )

                else:
                    raise ValueError(f"Unknown channel '{channel}'")

                inside = (intensity >= low) & (intensity <= high)

            bw = inside.to(img.dtype)[..., None]

            out_rgb = bw.expand(-1, -1, -1, 3).clone()

            if C == 4:
                out = torch.cat((out_rgb, img[..., 3:4]), dim=-1)
            else:
                out = out_rgb

            if msk is not None:
                out = out * msk[..., None] + img * (1.0 - msk[..., None])

            return out

        chunks = []

        for start in range(0, B, chunk_size):

            end = min(start + chunk_size, B)

            img_chunk = image[start:end]

            if mask is None:
                mask_chunk = None
            elif mask.shape[0] == 1:
                mask_chunk = mask
            else:
                mask_chunk = mask[start:end]

            chunks.append(threshold_chunk(img_chunk, mask_chunk))

        return (torch.cat(chunks, dim=0),)

    except Exception as e:
        raise RuntimeError(
            f"Color threshold operation failed: {e}"
        ) from e