"""Slideshow tests: storage budget math, viewer assembly fit, and a full disk
round-trip (viewer + N data files written and read back).

Run with `pytest` or directly: `python tests/test_slideshow.py`.
"""

import os
import subprocess
import sys
import tempfile

import numpy as np
from PIL import Image

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

from lenser import slideshow as ss                       # noqa: E402
from lenser import imageprep                             # noqa: E402
from lenser.diskimage import DiskError, have_c1541       # noqa: E402
from lenser.viewer.assemble import (build_slideshow_prg,  # noqa: E402
                                     have_xa)


def _gradient_png(path, w=160, h=200):
    yy, xx = np.mgrid[0:h, 0:w]
    rgb = np.stack([(xx * 255 // w), (yy * 255 // h),
                    ((xx + yy) * 255 // (w + h))], axis=-1).astype(np.uint8)
    Image.fromarray(rgb, "RGB").save(path)


def test_slideshow_budget():
    # CBM block rounding: a PRG of L bytes -> ceil((L+2)/254) blocks.
    assert ss.item_blocks("c64", "d64", 10001) == 40    # multicolor
    assert ss.item_blocks("c64", "d64", 9000) == 36     # hires/mono
    assert ss.item_blocks("c64", "d64", 1) == 1

    b = ss.budget("c64", "d64", [10001, 9000, 9000], viewer_len=283)
    assert b.fits and b.files == 4
    assert b.used_blocks == ss._cbm_blocks(283) + 40 + 36 + 36

    # too many multicolor images for a d64, but fine on a d81
    many = [10001] * 20
    assert not ss.budget("c64", "d64", many, 283).fits
    assert ss.budget("c64", "d81", many, 283).fits

    # directory-entry cap (d81 = 296) is enforced even when blocks would fit:
    # 300 one-block files use ~301 blocks (well under 3160) but 301 dir entries.
    tiny = [10] * 300
    over = ss.budget("c64", "d81", tiny, 283)
    assert over.used_blocks < over.total_blocks      # blocks would fit
    assert not over.fits and "directory" in over.reason


def test_slideshow_stub_assembles_and_fits():
    if not have_xa():
        return
    for advance in ("key", "seconds", "both"):
        for loop in (True, False):
            prg = build_slideshow_prg([0, 1, 0], advance=advance, seconds=5,
                                      loop=loop)
            assert prg[:2] == bytes([0x01, 0x08])          # load address $0801
            end = 0x0801 + (len(prg) - 2)
            assert end < 0x2000                            # code clears picture RAM


def test_slideshow_disk_roundtrip():
    if not (have_xa() and have_c1541()):
        return
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src)
        items = [
            ss.SlideItem(src, mode="multicolor"),
            ss.SlideItem(src, mode="hires",
                         prep=imageprep.PrepOptions(brightness=1.2)),
            ss.SlideItem(src, mode="mono", mono_base="green"),
        ]
        show = ss.Slideshow(platform="c64", disk_format="d64", advance="both",
                            seconds=5, loop=True, items=items)
        convs = [ss.convert_item(show, it) for it in items]
        out = os.path.join(td, "show.d64")
        ss.build_disk(show, out, convs=convs)

        # directory lists the viewer first, then 00, 01, 02 in order
        dirout = subprocess.run(["c1541", "-attach", out, "-dir"],
                                capture_output=True, text=True).stdout
        assert '"show"' in dirout
        for i in range(3):
            assert f'"{i:02d}"' in dirout

        # each data file reads back byte-identical to the conversion payload
        for i, c in enumerate(convs):
            host = os.path.join(td, f"r{i}.prg")
            subprocess.run(["c1541", "-attach", out, "-read", f"{i:02d}", host],
                           capture_output=True)
            raw = open(host, "rb").read()
            assert raw[0] | (raw[1] << 8) == c.data_addr
            assert raw[2:] == c.data


def test_slideshow_flavor_detection():
    class C:
        def __init__(self, v):
            self.viewer = v
    assert ss.slideshow_flavor([C("hires"), C("multicolor"), C("hires")]) == "simple"
    assert ss.slideshow_flavor([C("fli"), C("fli")]) == "fli"
    assert ss.slideshow_flavor([C("interlace")]) == "interlace"
    for bad in ([C("hires"), C("fli")], [C("fli"), C("interlace")]):
        try:
            ss.slideshow_flavor(bad)
            assert False, "mixed flavors should be rejected"
        except ValueError:
            pass


def test_fli_interlace_stubs_fit():
    if not have_xa():
        return
    # fli code must clear $4000, interlace must clear $2000 (their data regions)
    for flavor, limit in (("fli", 0x4000), ("interlace", 0x2000)):
        for video in ("pal", "ntsc"):
            prg = build_slideshow_prg([0, 0, 0], advance="both", seconds=5,
                                      loop=True, video=video, flavor=flavor)
            assert prg[:2] == bytes([0x01, 0x08])
            assert 0x0801 + (len(prg) - 2) < limit


def test_fli_interlace_disk_roundtrip():
    if not (have_xa() and have_c1541()):
        return
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src)
        for flavor, mode in (("fli", "fli"), ("interlace", "interlace")):
            items = [ss.SlideItem(src, mode=mode), ss.SlideItem(src, mode=mode)]
            show = ss.Slideshow(platform="c64", disk_format="d64", advance="both",
                                seconds=5, loop=True, items=items)
            convs = [ss.convert_item(show, it) for it in items]
            assert ss.slideshow_flavor(convs) == flavor
            out = os.path.join(td, f"{flavor}.d64")
            ss.build_disk(show, out, convs=convs)
            for i, c in enumerate(convs):
                host = os.path.join(td, f"{flavor}{i}.prg")
                subprocess.run(["c1541", "-attach", out, "-read", f"{i:02d}", host],
                               capture_output=True)
                raw = open(host, "rb").read()
                assert raw[0] | (raw[1] << 8) == c.data_addr   # $4000 fli / $2000 ifli
                assert raw[2:] == c.data


def test_c128_slideshow():
    assert ss.supports_slideshow("c128")
    # CBM-DOS budget shared with the C64
    assert ss.item_blocks("c128", "d64", 16384) == ss.item_blocks("c64", "d64", 16384)
    if not (have_xa() and have_c1541()):
        return
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src)
        items = [ss.SlideItem(src, mode="hicolor"), ss.SlideItem(src, mode="mono")]
        show = ss.Slideshow(platform="c128", disk_format="d64", advance="both",
                            seconds=5, loop=True, items=items)
        convs = [ss.convert_item(show, it) for it in items]
        out = os.path.join(td, "c128.d64")
        ss.build_disk(show, out, convs=convs)
        dirout = subprocess.run(["c1541", "-attach", out, "-dir"],
                                capture_output=True, text=True).stdout
        assert '"pic"' in dirout                 # boots via RUN"PIC"
        for i in range(2):
            assert f'"{i:02d}"' in dirout
            host = os.path.join(td, f"r{i}.prg")
            subprocess.run(["c1541", "-attach", out, "-read", f"{i:02d}", host],
                           capture_output=True)
            raw = open(host, "rb").read()
            assert raw[0] | (raw[1] << 8) == 0x4000   # VDC images load to $4000
            assert raw[2:] == convs[i].data
        # the 80x100 'color' VDC mode is not yet supported in slideshows
        cshow = ss.Slideshow(platform="c128", items=[ss.SlideItem(src, mode="color")])
        try:
            ss.check_modes("c128", [ss.convert_item(cshow, cshow.items[0])])
            assert False, "color mode should be rejected"
        except ValueError:
            pass


def test_atari_slideshow():
    assert ss.supports_slideshow("atari")
    assert ss.disk_formats("atari") == ["atr"]
    # ATR sectors are 128 bytes, no 2-byte load-address overhead
    assert ss.item_blocks("atari", "atr", 8196) == 65
    assert ss.item_blocks("atari", "atr", 256) == 2
    if not have_xa():
        return
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src, 160, 192)
        items = [ss.SlideItem(src, mode="gr15") for _ in range(3)]
        show = ss.Slideshow(platform="atari", disk_format="atr", advance="both",
                            seconds=5, loop=True, items=items)
        convs = [ss.convert_item(show, it) for it in items]
        out = os.path.join(td, "show.atr")
        ss.build_disk(show, out, convs=convs)
        raw = open(out, "rb").read()
        assert raw[:2] == bytes([0x96, 0x02])         # ATR magic
        assert raw[16 + 1] == 8                        # boot loads 8 sectors
        spi = ss.item_blocks("atari", "atr", len(convs[0].data))
        for i, c in enumerate(convs):
            sector = 9 + i * spi                       # boot_sectors(8)+1 + i*spi
            off = 16 + (sector - 1) * 128
            assert raw[off:off + len(c.data)] == c.data
        # gr15 / gr9 / gr8 / gr15dli all build a disk
        for mode in ("gr9", "gr8", "gr15dli"):
            sh = ss.Slideshow(platform="atari", disk_format="atr",
                              items=[ss.SlideItem(src, mode=mode)] * 2)
            ss.build_disk(sh, os.path.join(td, f"{mode}.atr"))
        # but modes may not be MIXED in one slideshow
        mix = [ss.convert_item(ss.Slideshow(platform="atari"),
                               ss.SlideItem(src, mode=m)) for m in ("gr15", "gr9")]
        try:
            ss.check_modes("atari", mix)
            assert False, "mixed Atari modes should be rejected"
        except ValueError:
            pass


def test_bbc_slideshow():
    assert ss.supports_slideshow("bbc")
    assert ss.disk_formats("bbc") == ["ssd"]
    assert ss.item_blocks("bbc", "ssd", 20480) == 80      # 256-byte DFS sectors
    if not have_xa():
        return
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src, 320, 256)
        items = [ss.SlideItem(src, mode="mode1") for _ in range(3)]
        show = ss.Slideshow(platform="bbc", disk_format="ssd", advance="both",
                            seconds=5, loop=True, items=items)
        convs = [ss.convert_item(show, it) for it in items]
        out = os.path.join(td, "show.ssd")
        ss.build_disk(show, out, convs=convs)
        d = open(out, "rb").read()
        nfiles = d[0x105] // 8
        names = [d[8 + i * 8:8 + i * 8 + 7].decode("ascii").rstrip()
                 for i in range(nfiles)]
        assert "!BOOT" in names and "PIC" in names
        for i in range(3):
            assert f"{i:02d}" in names
            j = names.index(f"{i:02d}")
            e = 0x100 + 8 + j * 8
            start = ((d[e + 6] & 3) << 8) | d[e + 7]
            length = (((d[e + 6] >> 4) & 3) << 16) | (d[e + 5] << 8) | d[e + 4]
            off = start * 256
            assert d[off:off + length] == convs[i].data
        # a single BBC screen mode per slideshow
        m2 = ss.convert_item(ss.Slideshow(platform="bbc"),
                             ss.SlideItem(src, mode="mode2"))
        try:
            ss.check_modes("bbc", [convs[0], m2])
            assert False, "mixed BBC modes should be rejected"
        except ValueError:
            pass


def test_apple_slideshow():
    assert ss.supports_slideshow("apple")
    assert ss.disk_formats("apple") == ["dsk"]
    assert ss.item_blocks("apple", "dsk", 8192) == 32
    assert ss.budget("apple", "dsk", [8192] * 4, 0).fits        # RAM holds 4 HGR
    assert not ss.budget("apple", "dsk", [8192] * 5, 0).fits
    if not have_xa():
        return
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src, 280, 192)
        items = [ss.SlideItem(src, mode="hgr_color") for _ in range(3)]
        show = ss.Slideshow(platform="apple", disk_format="dsk", advance="both",
                            seconds=5, loop=True, items=items)
        convs = [ss.convert_item(show, it) for it in items]
        out = os.path.join(td, "show.dsk")
        ss.build_disk(show, out, convs=convs)
        d = open(out, "rb").read()
        assert len(d) == 143360                  # DOS 3.3 .dsk
        assert d[0] == 0x01                       # boot sector byte 0
        # too many images for RAM
        try:
            ss.check_modes("apple", convs * 2)    # 6 > 4
            assert False, "over-RAM Apple slideshow should be rejected"
        except ValueError:
            pass
        # DHGR is not supported for slideshows yet
        dh = ss.convert_item(ss.Slideshow(platform="apple"),
                             ss.SlideItem(src, mode="dhgr"))
        try:
            ss.check_modes("apple", [dh])
            assert False, "DHGR should be rejected"
        except ValueError:
            pass


def test_iigs_slideshow():
    assert ss.supports_slideshow("iigs")
    assert ss.disk_formats("iigs") == ["dsk"]
    assert ss.item_blocks("iigs", "dsk", 32768) == 128
    assert ss.budget("iigs", "dsk", [32768] * 4, 0).fits
    assert not ss.budget("iigs", "dsk", [32768] * 5, 0).fits
    if not have_xa():
        return
    # the two-stage loader: boot must fit one 256-byte sector
    from lenser.iigs.viewer.assemble import build_slideshow
    boot, stage2, pages = build_slideshow(3, advance="both", seconds=5, loop=True)
    assert len(boot) <= 256 and len(stage2) > 0 and pages >= 1
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src, 320, 200)
        items = [ss.SlideItem(src, mode="shr") for _ in range(3)]
        show = ss.Slideshow(platform="iigs", disk_format="dsk", advance="both",
                            seconds=5, loop=True, items=items)
        convs = [ss.convert_item(show, it) for it in items]
        out = os.path.join(td, "show.dsk")
        ss.build_disk(show, out, convs=convs)
        d = open(out, "rb").read()
        assert len(d) == 143360 and d[0] == 0x01      # bootable .dsk
        try:
            ss.check_modes("iigs", convs * 2)          # 6 > 4
            assert False, "over-capacity IIgs slideshow should be rejected"
        except ValueError:
            pass


def test_amiga_slideshow():
    assert ss.supports_slideshow("amiga")
    assert ss.disk_formats("amiga") == ["adf"]
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src, 320, 256)
        items = [ss.SlideItem(src, mode="lowres") for _ in range(3)]
        show = ss.Slideshow(platform="amiga", disk_format="adf", advance="both",
                            seconds=4, loop=True, items=items)
        convs = [ss.convert_item(show, it) for it in items]
        # Amiga conv.data is a dict; its on-disk size comes from the blob
        nb = ss.image_nbytes(convs[0])
        assert nb > 8000 and ss.item_blocks("amiga", "adf", nb) == -(-nb // 512)
        out = os.path.join(td, "show.adf")
        ss.build_disk(show, out, convs=convs)
        d = open(out, "rb").read()
        assert len(d) == 901120 and d[:4] == b"DOS\x00"   # bootable 880K .adf
        # chip-RAM bound rejects an oversized show
        assert not ss.budget("amiga", "adf", [40000] * 20, 0).fits


def test_slideshow_overfill_raises():
    if not have_xa():
        return
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "g.png")
        _gradient_png(src)
        # 25 multicolor images (~40 blocks each) cannot fit a 664-block d64
        items = [ss.SlideItem(src, mode="multicolor") for _ in range(25)]
        show = ss.Slideshow(disk_format="d64", items=items)
        # reuse one conversion for all to keep the test fast
        c0 = ss.convert_item(show, items[0])
        convs = [c0] * 25
        out = os.path.join(td, "show.d64")
        try:
            ss.build_disk(show, out, convs=convs)
            assert False, "expected DiskError for an over-budget slideshow"
        except DiskError as e:
            assert "does not fit" in str(e)


def _imgobj(w, h):
    arr = np.zeros((h, w, 3), np.uint8)
    return Image.fromarray(arr, "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)} slideshow tests passed.")