"""PBC3 bitstream framing, separate from the image algorithm."""

import numpy as np

import pbc3_ops as ops
from pbc3_types import BitReader, BitWriter


def write_grid(bw: BitWriter, values, bitcount: int) -> None:
    """Writes a flat grid of palette indices."""
    acc = bw.acc
    nbits = bw.nbits
    data = bw.data
    bitmask = (1 << bitcount) - 1
    for value in values:
        acc = (acc << bitcount) | (int(value) & bitmask)
        nbits += bitcount
        while nbits >= 8:
            shift = nbits - 8
            data.append((acc >> shift) & 255)
            acc &= (1 << shift) - 1
            nbits = shift
    bw.acc = acc
    bw.nbits = nbits


def read_grid(br: BitReader, count: int, bitcount: int) -> np.ndarray:
    """Reads a flat grid of palette indices."""
    values = np.empty(count, dtype=np.uint16)
    data = br.data
    data_len = len(data)
    data_index = br.i
    acc = br.acc
    nbits = br.nbits
    bitmask = (1 << bitcount) - 1
    for i in range(count):
        while nbits < bitcount:
            if data_index >= data_len:
                raise EOFError("bitstream ended early")
            acc = (acc << 8) | data[data_index]
            data_index += 1
            nbits += 8
        shift = nbits - bitcount
        values[i] = (acc >> shift) & bitmask
        acc &= (1 << shift) - 1
        nbits = shift
    br.i = data_index
    br.acc = acc
    br.nbits = nbits
    return values


def write_header(bw, w, h, original_w, original_h, downsampled, color_id, channels,
                 channel_bits, positive_bias, has_alpha, patch_count, base_values, warmup=None):
    """Writes the image-level stream header."""
    bw.write(int(downsampled), 1)
    if downsampled:
        bw.write(original_w, 16)
        bw.write(original_h, 16)
    bw.write(w, 16)
    bw.write(h, 16)
    bw.write(color_id, 2)
    bw.write(channels, 8)
    bw.write(channel_bits, 4)
    bw.write(int(positive_bias), 1)
    bw.write(int(has_alpha), 1)
    bw.write(patch_count, 32)
    for base in base_values:
        bw.write(base, 8)
    bw.write(int(warmup is not None), 1)
    if warmup is not None:
        warm_w, warm_h, warm_split = warmup
        bw.write(warm_w, 16)
        bw.write(warm_h, 16)
        bw.write(warm_split, 32)


def read_header(br, color_space_names):
    """Reads the image-level stream header."""
    downsampled = bool(br.read(1))
    original_w = br.read(16) if downsampled else None
    original_h = br.read(16) if downsampled else None
    w = br.read(16)
    h = br.read(16)
    color_id = br.read(2)
    channels = br.read(8)
    channel_bits = br.read(4)
    positive_bias = bool(br.read(1))
    has_alpha = bool(br.read(1))
    patch_count = br.read(32)
    base_values = [br.read(8) for _ in range(channels)]
    warmup_on = bool(br.read(1))
    warm_w = warm_h = warmup_split = None
    if warmup_on:
        warm_w = br.read(16)
        warm_h = br.read(16)
        warmup_split = br.read(32)
    return (
        downsampled, original_w, original_h, w, h, color_space_names[color_id], channels,
        channel_bits, positive_bias, has_alpha, patch_count, base_values, warmup_on, warm_w,
        warm_h, warmup_split,
    )


def write_patch(bw, patch, channel_bits: int) -> None:
    """Writes one generated-palette patch."""
    bw.write(patch["channel"], channel_bits)
    bw.write(patch["x"], 16)
    bw.write(patch["y"], 16)
    bw.write(patch["w"], 16)
    bw.write(patch["h"], 16)
    bw.write(0, 1)
    mask = patch["mask"]
    bw.write(len(mask), 10)
    for bit in mask:
        bw.write(bit, 1)
    bw.write(patch["neg"], 8)
    bw.write(patch["pos"], 8)
    bw.write(patch["max_bitcount"], 4)
    bw.write(patch["cell_size"], 16)
    write_grid(bw, patch["indices"].ravel().astype(np.int64), patch["bitcount"])


def read_patch(br, channel_bits: int, positive_bias: bool = True):
    """Reads one generated-palette patch and returns its decoded values."""
    channel = br.read(channel_bits)
    x, y = br.read(16), br.read(16)
    w, h = br.read(16), br.read(16)
    if br.read(1) != 0:
        raise ValueError("explicit palette patches were removed in PBC3 3.0 release cleanup")
    mask = [br.read(1) for _ in range(br.read(10))]
    negative_max, positive_max = br.read(8), br.read(8)
    max_bitcount = br.read(4)
    bitcount = ops.resolve_palette_bitcount(mask, max_bitcount, negative_max, positive_max, positive_bias)
    palette = ops.palette_generator(mask, max_bitcount, negative_max, positive_max, positive_bias)
    cell_size = br.read(16)
    gw, gh = ops.ceil_div(w, cell_size), ops.ceil_div(h, cell_size)
    indices = read_grid(br, gh * gw, bitcount).reshape(gh, gw)
    return channel, x, y, w, h, cell_size, palette[indices], bitcount