"""Regression tests: decode each mode's emitted VIC-II bytes and check they
reproduce the converter's own index image, and that every viewer assembles and
fits.  Run with `pytest` or directly: `python tests/test_roundtrip.py`.

These tests exercise the byte-packing that the GUI preview deliberately does *not*
touch (the preview renders from the index image), so they are the safety net that
catches an encoding bug before it reaches a real C64.
"""

import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from c64view import imageprep, palette as pal              # noqa: E402
from c64view.convert import fli, hires, ifli, multicolor   # noqa: E402
from c64view.viewer.assemble import SOURCES, build_viewer_prg, have_xa  # noqa: E402


def _gradient(w, h):
    yy, xx = np.mgrid[0:h, 0:w]
    rgb = np.stack([(xx * 255 // w), (yy * 255 // h),
                    ((xx + yy) * 255 // (w + h))], axis=-1)
    return rgb.astype(np.uint8)


def _decode_mc(bitmap, screen, colram, bg):
    dec = np.zeros((200, 160), np.uint8)
    for cr in range(25):
        for cc in range(40):
            ci = cr * 40 + cc
            lut = [bg, screen[ci] >> 4, screen[ci] & 0xF, colram[ci] & 0xF]
            for r in range(8):
                byte = bitmap[cr * 320 + cc * 8 + r]
                for x in range(4):
                    dec[cr * 8 + r, cc * 4 + x] = lut[(byte >> (6 - 2 * x)) & 3]
    return dec


def test_multicolor_roundtrip():
    img = imageprep.prepare(_imgobj(160, 200), 160, 200, 2.0, imageprep.PrepOptions())
    c = multicolor.convert(img)
    d = c.data
    dec = _decode_mc(np.frombuffer(d[:8000], np.uint8),
                     np.frombuffer(d[8000:9000], np.uint8),
                     np.frombuffer(d[9000:10000], np.uint8), d[10000])
    assert np.array_equal(dec, c.index_image)


def test_hires_roundtrip():
    img = imageprep.prepare(_imgobj(320, 200), 320, 200, 1.0, imageprep.PrepOptions())
    c = hires.convert(img)
    bitmap = np.frombuffer(c.data[:8000], np.uint8)
    screen = np.frombuffer(c.data[8000:9000], np.uint8)
    dec = np.zeros((200, 320), np.uint8)
    for cr in range(25):
        for cc in range(40):
            ci = cr * 40 + cc
            fg, bgc = screen[ci] >> 4, screen[ci] & 0xF
            for r in range(8):
                byte = bitmap[cr * 320 + cc * 8 + r]
                for x in range(8):
                    dec[cr * 8 + r, cc * 8 + x] = fg if (byte >> (7 - x)) & 1 else bgc
    assert np.array_equal(dec, c.index_image)


def test_fli_roundtrip():
    img = imageprep.prepare(_imgobj(160, 200), 160, 200, 2.0, imageprep.PrepOptions())
    c = fli.convert(img)
    d = c.data
    screens = [np.frombuffer(d[L * 1024:L * 1024 + 1000], np.uint8) for L in range(8)]
    bitmap = np.frombuffer(d[8192:8192 + 8000], np.uint8)
    colram = np.frombuffer(d[16384:16384 + 1000], np.uint8)
    bg = d[17384]
    dec = np.zeros((200, 160), np.uint8)
    for cr in range(25):
        for cc in range(40):
            ci = cr * 40 + cc
            for r in range(8):
                sb = screens[r][ci]
                lut = [bg, sb >> 4, sb & 0xF, colram[ci] & 0xF]
                byte = bitmap[cr * 320 + cc * 8 + r]
                for x in range(4):
                    dec[cr * 8 + r, cc * 4 + x] = lut[(byte >> (6 - 2 * x)) & 3]
    assert np.array_equal(dec, c.index_image)


def test_interlace_blend_better():
    """Interlace blend error should beat plain multicolor on a gradient."""
    img = imageprep.prepare(_imgobj(160, 200), 160, 200, 2.0, imageprep.PrepOptions())
    assert ifli.convert(img).error < multicolor.convert(img).error + 1e-6
    assert len(ifli.convert(img).data) == 25577


def test_viewers_assemble_and_fit():
    if not have_xa():
        return  # xa not installed; skip
    sizes = {"hires": 9000, "multicolor": 10001, "fli": 17385,
             "fli_ntsc": 17385, "interlace": 25577}
    for key in SOURCES:
        prg = build_viewer_prg(key, bytes(sizes[key]),
                               0x4000 if key.startswith("fli") else 0x2000)
        assert prg[:2] == bytes([0x01, 0x08])     # PRG load address $0801


def _imgobj(w, h):
    from PIL import Image
    return Image.fromarray(_gradient(w, h), "RGB")


if __name__ == "__main__":
    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    for fn in fns:
        fn()
        print(f"PASS {fn.__name__}")
    print(f"\nAll {len(fns)} tests passed.")