import hashlib
import struct
import tempfile
import unittest
from pathlib import Path

from rewrite_baked_lua_string import stage


def fixture(root, contents):
    digest = hashlib.md5(contents).digest()
    (root / (digest.hex() + ".bytes")).write_bytes(contents)
    # One block and one module, with a 3-byte LuaV block trailer.
    index = (struct.pack(">IIII", 255, 4, 1, 1) + bytes(8) + digest +
             struct.pack(">III", 0, len(contents), 1) + bytes(8) +
             struct.pack(">II", len(contents), 0) + b"\x00\x00\x80")
    index_hash = hashlib.md5(index).digest()
    (root / ("LuaV_" + index_hash.hex() + ".bytes")).write_bytes(index)
    selector = bytearray(66)
    selector[:4] = b"SRMI"
    for i in range(4):
        selector[28 + 4 * i:32 + 4 * i] = index_hash[4 * i:4 * i + 4][::-1]
    struct.pack_into("<I", selector, 44, len(index))
    (root / "M_LuaV.bytes").write_bytes(selector)


class RewriteTest(unittest.TestCase):
    def test_changed_string_updates_digest_and_round_trips(self):
        with tempfile.TemporaryDirectory() as folder:
            root = Path(folder)
            src, out = root / "src", root / "stage"
            src.mkdir()
            old, new = b"scene/a.prefab", b"scene/big.prefab"
            fixture(src, b"prefix" + bytes([len(old)]) + old + b"suffix")
            result = stage(src, out, old.decode(), new.decode())
            changed = (out / (result["new_block_md5"] + ".bytes")).read_bytes()
            self.assertIn(bytes([len(new)]) + new, changed)
            self.assertEqual(result["module_new_size"] - result["module_old_size"], len(new) - len(old))
            reverse = root / "reverse"
            result2 = stage(out, reverse, new.decode(), old.decode())
            self.assertEqual(result2["new_block_md5"], result["old_block_md5"])
            self.assertEqual(result2["new_index_md5"], result["old_index_md5"])
            self.assertEqual((reverse / "M_LuaV.bytes").read_bytes(), (src / "M_LuaV.bytes").read_bytes())

    def test_duplicate_string_refuses_write(self):
        with tempfile.TemporaryDirectory() as folder:
            root = Path(folder)
            src, out = root / "src", root / "stage"
            src.mkdir()
            old = b"a.prefab"
            fixture(src, (bytes([len(old)]) + old) * 2)
            with self.assertRaisesRegex(ValueError, "ambiguous"):
                stage(src, out, old.decode(), "longer.prefab")
            self.assertFalse(out.exists())


if __name__ == "__main__":
    unittest.main()