"""LZO, UE3 packages and soundbanks (ackit/lzo.py, upk.py, bnk.py)."""
import os

import pytest

import synth
from ackit import lzo
from ackit.bnk import Bank, find_banks, fnv1, IN_MEMORY
from ackit.upk import Package


# --- LZO ----------------------------------------------------------------------------------------
@pytest.mark.parametrize("n", [0, 1, 3, 4, 17, 18, 19, 238, 239, 273, 274, 529, 0x20000])
def test_lzo_stored_roundtrip(n):
    data = os.urandom(n)
    assert lzo.decompress(lzo.compress_stored(data), n) == data


def test_lzo_matches():
    # "abc" + M3 match (9 bytes at distance 3) + end
    assert lzo.decompress(bytes([20]) + b"abc" + bytes([39, 0x08, 0x00]) + lzo.EOF) == b"abc" * 4
    # "abcd" + M2 match (4 bytes at distance 4) + end
    assert lzo.decompress(bytes([21]) + b"abcd" + bytes([108, 0]) + lzo.EOF) == b"abcd" * 2


def test_lzo_errors():
    with pytest.raises(lzo.LzoError):
        lzo.decompress(bytes([25]) + b"abc")                       # truncated
    with pytest.raises(lzo.LzoError):
        lzo.decompress(lzo.compress_stored(b"hello") + b"x")        # trailing bytes
    with pytest.raises(lzo.LzoError):
        lzo.decompress(lzo.compress_stored(b"hello"), 6)             # wrong size


# --- packages -----------------------------------------------------------------------------------
@pytest.fixture
def pkg(tmp_path):
    body = bytes(range(256)) * 80                                     # 20480 bytes: 5 blocks, 2 chunks
    p = str(tmp_path / "Synth.upk")
    start = synth.package(p, body)
    return p, body, start


def test_package_read(pkg):
    p, body, start = pkg
    pk = Package(p)
    assert pk.compressed and pk.version == 805 and pk.header_size == start
    assert pk.size == start + len(body) and len(pk.chunks) == 2 and len(pk.blocks) == 5
    assert pk.read(start + 4000, 300) == body[4000:4300]                # crosses a block boundary
    assert pk.read_all() == (start, body)


def test_package_patch(pkg):
    p, body, start = pkg
    raw = open(p, "rb").read()
    pk = Package(p)
    assert pk.patched([]) == raw                                       # nothing changed -> same bytes
    new = pk.patched([(start + 16334, b"X" * 100)])                    # second chunk, across two blocks
    q = p + ".new"
    open(q, "wb").write(new)
    pk2 = Package(q)
    exp = bytearray(body)
    exp[16334:16434] = b"X" * 100
    assert pk2.read_all()[1] == bytes(exp)
    # the first chunk is copied verbatim
    c0 = pk.chunks[0]
    assert new[c0[2]:c0[2] + c0[3]] == raw[c0[2]:c0[2] + c0[3]]
    with pytest.raises(ValueError):
        pk.patched([(start + len(body) - 10, b"too long" * 4)])


# --- banks --------------------------------------------------------------------------------------
def test_bank_parse():
    media = {0x111: b"RIFF" + bytes(41), 0x222: b"RIFF" + bytes(100)}
    b = synth.bank(fnv1("CHR_Test"), media, name="CHR_Test")
    bk = Bank(b)
    assert bk.id == fnv1("CHR_Test") and bk.version == 56 and bk.size == len(b)
    assert bk.names() == {bk.id: "CHR_Test"}
    m = {x.id: x for x in bk.media()}
    assert b[m[0x222].off:m[0x222].off + m[0x222].size] == media[0x222]
    assert m[0x222].off % 16 == (bk.chunks["DATA"][0]) % 16
    src = {s.media: s for s in bk.sources()}
    assert src[0x111].stream == IN_MEMORY and (src[0x111].off, src[0x111].size) == (m[0x111].off, m[0x111].size)


def test_find_banks_in_buffer():
    a = synth.bank(1, {5: b"x" * 20})
    b = synth.bank(2, {6: b"y" * 30})
    buf = b"junkBKHDjunk" + a + bytes(7) + b
    assert [x.id for x in find_banks(buf)] == [1, 2]


def test_fnv1():
    assert fnv1("Init") == 0x50C63A23 and fnv1("Music") == 0xEDF036D6


# --- real game (read-only) ----------------------------------------------------------------------
def test_real_package_blocks(real_game_dir):
    from ackit.game import Game
    p = os.path.join(Game(real_game_dir).cooked, "BmGame.upk")
    pk = Package(p)
    assert pk.compressed and pk.patched([]) == open(p, "rb").read()
    start, data = pk.read_all()
    banks = list(find_banks(data))
    assert len(banks) > 20
    bk = banks[0]
    m = bk.media()[0]
    off = start + bk.at + m.off
    assert pk.read(off, m.size) == data[bk.at + m.off:bk.at + m.off + m.size]
    assert pk.read(off, 4) == b"RIFF"