"""图像管线：调整 → 抖动量化 → 1bpp 平面打包。
算法与参数证据：analysis/web/timedjs.js（见 docs/protocol.md §5）。
所有实现与 Web 端 JS 语义对齐：同样的灰度系数、误差核、BCD 无 —— 像素级等价为目标。
"""
from __future__ import annotations

from PIL import Image

GRAY = (0.299, 0.587, 0.114)

BAYER4 = (
    (15, 135, 45, 165),
    (195, 75, 225, 105),
    (60, 180, 30, 150),
    (240, 120, 210, 90),
)

# 误差扩散核：(dx, dy, 权重分子)；除数见 _KERNELS 键名
_KERNELS = {
    "floydsteinberg": ([(1, 0, 7), (-1, 1, 3), (0, 1, 5), (1, 1, 1)], 16),
    "atkinson": (
        [(1, 0, 1), (2, 0, 1), (-1, 1, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1)],
        8,
    ),
    "stucki": (
        [
            (1, 0, 8), (2, 0, 4),
            (-2, 1, 2), (-1, 1, 4), (0, 1, 8), (1, 1, 4), (2, 1, 2),
            (-2, 2, 1), (-1, 2, 2), (0, 2, 4), (1, 2, 2), (2, 2, 1),
        ],
        42,
    ),
    "jarvis": (
        [
            (1, 0, 7), (2, 0, 5),
            (-2, 1, 3), (-1, 1, 5), (0, 1, 7), (1, 1, 5), (2, 1, 3),
            (-2, 2, 1), (-1, 2, 3), (0, 2, 5), (1, 2, 3), (2, 2, 1),
        ],
        48,
    ),
}

ALGORITHMS = ["none", "floydsteinberg", "atkinson", "bayer", "stucki", "jarvis"]


def _clamp(v: float) -> int:
    return 0 if v < 0 else (255 if v > 255 else round(v))


def _gray_of(r: float, g: float, b: float) -> float:
    return GRAY[0] * r + GRAY[1] * g + GRAY[2] * b


def adjust_pixels(pixels, brightness=0, contrast=0, saturation=100):
    """JS applyImageAdjustments 语义（timedjs.js:1-41）。
    brightness: ±100 → v±255*b/100；contrast: ±100 → v*f+(1-f)*128, f=(c+100)/100；
    saturation: 0–100+ → gray+(v-gray)*s/100。"""
    out = []
    bf = 255 * brightness / 100
    cf = (contrast + 100) / 100
    ca = (1 - cf) * 128
    sf = saturation / 100
    for r, g, b in pixels:
        if brightness:
            r, g, b = r + bf, g + bf, b + bf
        if contrast:
            r, g, b = r * cf + ca, g * cf + ca, b * cf + ca
        if saturation != 100:
            gray = _gray_of(r, g, b)
            r, g, b = (gray + (v - gray) * sf for v in (r, g, b))
        out.append((_clamp(r), _clamp(g), _clamp(b)))
    return out


def _is_red(r, g, b, threshold):
    return r > g * 1.5 and r > b * 1.5 and r > threshold


def dither_pixels(pixels, width, height, algorithm, threshold=125,
                  diffusion=1.0, tricolor=False):
    """量化到黑/白（+红）。返回与输入等长的 (r,g,b) 列表。
    与 JS 一致：红判定优先且红像素不扩散误差；误差按通道扩散；写入即截断取整。"""
    if algorithm not in ALGORITHMS:
        raise ValueError(f"未知抖动算法 {algorithm!r}，可选：{ALGORITHMS}")
    buf = [[float(r), float(g), float(b)] for r, g, b in pixels]
    out = [(0, 0, 0)] * (width * height)

    def quantize(i, x, y):
        r, g, b = buf[i]
        if tricolor and _is_red(r, g, b, threshold):
            out[i] = (255, 0, 0)
            return None  # 红像素无误差扩散
        gray = _gray_of(r, g, b)
        if algorithm == "bayer":
            limit = BAYER4[y % 4][x % 4]
        else:
            limit = threshold
        new = 0 if gray < limit else 255
        out[i] = (new, new, new)
        return (r - new) * diffusion, (g - new) * diffusion, (b - new) * diffusion

    if algorithm in ("none", "bayer"):
        for y in range(height):
            for x in range(width):
                quantize(y * width + x, x, y)
        return out

    kernel, div = _KERNELS[algorithm]
    for y in range(height):
        for x in range(width):
            i = y * width + x
            errs = quantize(i, x, y)
            if errs is None:
                continue
            er, eg, eb = errs
            for dx, dy, w in kernel:
                nx, ny = x + dx, y + dy
                if 0 <= nx < width and ny < height:
                    ni = ny * width + nx
                    # JS：写入 Uint8ClampedArray 即截断取整（round）
                    buf[ni][0] = _clamp(buf[ni][0] + er * w / div)
                    buf[ni][1] = _clamp(buf[ni][1] + eg * w / div)
                    buf[ni][2] = _clamp(buf[ni][2] + eb * w / div)
    return out


def pack_plane(pixels, plane):
    """(r,g,b) 列表 → 1bpp MSB-first 字节流。
    bw: R,G,B 全>0 → 1；red: R>0 且 G=0 且 B=0 → 1（timedjs.js:564-584）。"""
    if len(pixels) % 8:
        raise ValueError("像素数必须是 8 的倍数")
    out = bytearray()
    acc = 0
    nbits = 0
    for r, g, b in pixels:
        if plane == "bw":
            bit = 1 if (r > 0 and g > 0 and b > 0) else 0
        elif plane == "red":
            bit = 1 if (r > 0 and g == 0 and b == 0) else 0
        else:
            raise ValueError("plane 必须是 bw 或 red")
        acc = (acc << 1) | bit
        nbits += 1
        if nbits == 8:
            out.append(acc)
            acc = 0
            nbits = 0
    return bytes(out)


def process_image(img: Image.Image, *, size=(400, 300), algorithm="atkinson",
                  tricolor=True, threshold=125, diffusion=1.0,
                  brightness=0, contrast=0, saturation=100, rotate=0):
    """完整管线：缩放/旋转 → 调整 → 抖动 → 打包。返回 (bw_bytes, red_bytes|None)。"""
    img = img.convert("RGB").resize(size, Image.LANCZOS)
    if rotate:
        img = img.rotate(-rotate, expand=False)  # 与 JS 顺时针旋转语义一致
    w, h = size
    px = list(img.getdata())
    px = adjust_pixels(px, brightness, contrast, saturation)
    px = dither_pixels(px, w, h, algorithm, threshold, diffusion, tricolor)
    bw = pack_plane(px, "bw")
    red = pack_plane(px, "red") if tricolor else None
    return bw, red