from __future__ import annotations

from vsdeband import f3k_deband
from vsmasktools import dre_edgemask
from vstools import (
    UnsupportedVideoFormatError,
    VariableFormatError,
    check_variable,
    core,
    vs,
)

__all__ = [
    "retinex_deband",
]


def retinex_deband(
    clip: vs.VideoNode,
    threshold: int = 24,
    showmask: bool = False,
) -> vs.VideoNode:
    """Debanding using contrast-adaptive edge masking.

    Args:
        clip: Input video (8-16bit YUV required).
        threshold: Debanding strength (0-255). A reasonable range is around 16-64.
            According to the original f3kdb documentation (https://f3kdb.readthedocs.io/en/stable/presets.html):
            - Low = 32
            - Medium = 48
            - High = 64
            - Very high = 80
        showmask: If True, return edge mask instead of debanded clip.

    Returns:
        Debanded video clip or edge mask.

    Note:
        Does not add grain. Use vsdeband.AddNoise for post-denoising.
    """
    if threshold < 0 or threshold > 255:
        raise ValueError(f"threshold must be between 0-255, got {threshold}")

    if not check_variable(clip, retinex_deband):
        raise VariableFormatError("clip must have constant format and fps")

    if (
        clip.format.color_family != vs.YUV
        or clip.format.sample_type != vs.INTEGER
        or clip.format.bits_per_sample > 16
    ):
        raise UnsupportedVideoFormatError(
            retinex_deband,
            clip.format,
            "The format {format.name} is not supported! It must be an 8-16bit integer YUV bit format!",
        )

    mask: vs.VideoNode = dre_edgemask.CLAHE(clip)

    if showmask:
        return mask

    # The bitshift here is to adjust threshold values from neo_f3kdb's
    # `scale=False` behavior to match vs-jetpack's `scale=True` behavior.
    # It is not related to bit depth. The threshold value is independent of bit depth.
    deband = f3k_deband(clip, thr=(threshold << 2))
    return core.std.MaskedMerge(deband, clip, mask)
