"""ZX Spectrum image encoder.

256x192, two colours per 8x8 cell (ink + paper), like C64 hires -- but the two
colours of a cell must share the BRIGHT bit, so each cell's pair is chosen from
one brightness group (normal 0-7 or bright 8-15).  For error-diffusion dithers
the pair is picked dither-aware (segment metric) so the two colours bracket the
cell and dithering blends to the true shade; ordered/none use nearest-colour.

Produces the classic 6912-byte screen: 6144-byte interleaved bitmap + 768
attribute bytes (FLASH BRIGHT PAPER INK).
"""

from __future__ import annotations

import numpy as np

from ... import dither, palette as c64pal
from ...convert import base
from .. import palette as spal

WIDTH, HEIGHT = 256, 192
CELL_W, CELL_H = 8, 8
PIXEL_ASPECT = 1.0
N_COLS, N_ROWS = 32, 24
BITMAP_BYTES = 6144
ATTR_BYTES = 768


def _select(cells, plab, dither_mode):
    """Per cell, pick the best two-colour pair from whichever brightness group
    (normal/bright) fits better -- both colours share the BRIGHT bit."""
    if dither_mode in base.DIFFUSION_DITHERS:
        seg = base.segment_distances(cells, plab)
        sets_n, err_n = base.select_cell_sets_dither(
            cells, plab, spal.NORMAL_GROUP, n_free=2, seg=seg)
        sets_b, err_b = base.select_cell_sets_dither(
            cells, plab, spal.BRIGHT_GROUP, n_free=2, seg=seg)
    else:
        dist = base.cell_distance(cells, plab)
        sets_n, err_n = base.select_cell_sets(dist, spal.NORMAL_GROUP, n_free=2)
        sets_b, err_b = base.select_cell_sets(dist, spal.BRIGHT_GROUP, n_free=2)
    use_b = err_b < err_n
    return np.where(use_b[:, None], sets_b, sets_n)


def _bitmap_offset(y: int, cx: int) -> int:
    """Offset within the 6144-byte bitmap for pixel row y, char column cx."""
    return ((y & 7) << 8) | ((y & 0x38) << 2) | ((y & 0xC0) << 5) | cx


def convert(img_rgb, palette_name="spectrum", dither_mode="floyd",
            intensive=False, base_color=None):
    plab = spal.palette_lab()
    prgb = spal.get_palette().astype(np.uint8)
    img_lab = c64pal.srgb_to_lab(img_rgb)

    cells, rows, cols = base.cells_lab(img_lab, CELL_W, CELL_H)
    sets = _select(cells, plab, dither_mode)

    allowed = base.per_pixel_allowed(sets, rows, cols, CELL_W, CELL_H, HEIGHT, WIDTH)
    idx = dither.quantize(img_lab, allowed, plab, dither_mode).astype(np.uint8)

    scr = _encode(idx, sets, rows, cols)

    return base.Conversion(
        mode="hires", width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
        index_image=idx.astype(np.uint16), data=bytes(scr), data_addr=0x4000,
        viewer="spectrum", preview_rgb=prgb[idx],
        error=base.perceptual_error(idx, img_lab, plab),
        meta={"palette": "spectrum", "dither": dither_mode},
    )


def _encode(idx, sets, rows, cols):
    """Build the 6912-byte screen (interleaved bitmap + attributes)."""
    scr = bytearray(BITMAP_BYTES + ATTR_BYTES)
    for cy in range(rows):
        for cx in range(cols):
            ci = cy * cols + cx
            paper, ink = int(sets[ci, 0]), int(sets[ci, 1])
            bright = ink >> 3            # both colours share the BRIGHT bit
            scr[BITMAP_BYTES + ci] = (bright << 6) | ((paper & 7) << 3) | (ink & 7)
            for r in range(8):
                y = cy * 8 + r
                row = idx[y, cx * 8:cx * 8 + 8]
                byte = 0
                for px in range(8):
                    byte = (byte << 1) | (1 if row[px] == ink else 0)
                scr[_bitmap_offset(y, cx)] = byte
    return scr