#!/usr/bin/env python3
"""Create a staged Lua bundle rewrite; never modify the source installation.

The caller supplies the Client's Lua/Windows directory, output directory and
two exact asset paths. This handles the observed LuaV big-endian index and
SRMI digest format only; it refuses ambiguous input instead of guessing.
"""
import argparse
import hashlib
import json
import struct
from pathlib import Path


def md5(data):
    return hashlib.md5(data).digest()


def replace_once(buf, old, new):
    if len(old) >= 128 or len(new) >= 128:
        raise ValueError("only single-byte Lua string lengths are supported")
    original = bytes([len(old)]) + old
    replacement = bytes([len(new)]) + new
    if buf.count(original) != 1:
        raise ValueError("the exact length-prefixed string must occur once")
    index = buf.index(original)
    return buf[:index] + replacement + buf[index + len(original):], index


def stage(source, output, old, new):
    source, output = source.resolve(strict=True), output.resolve()
    if output == source or source in output.parents:
        raise ValueError("output must be outside the source installation")
    if output.exists() and any(output.iterdir()):
        raise ValueError("output must be empty")
    selector = (source / "M_LuaV.bytes").read_bytes()
    if len(selector) != 66 or selector[:4] != b"SRMI":
        raise ValueError("unsupported LuaV selector")
    digest = b"".join(selector[i:i + 4][::-1] for i in range(28, 44, 4))
    index_path = source / ("LuaV_" + digest.hex() + ".bytes")
    index = bytearray(index_path.read_bytes())
    if md5(index) != digest or struct.unpack_from("<I", selector, 44)[0] != len(index):
        raise ValueError("LuaV selector digest or size does not match")
    if len(index) < 16 or struct.unpack_from(">II", index, 0) != (255, 4):
        raise ValueError("unsupported LuaV index header")
    blocks, total_modules = struct.unpack_from(">II", index, 8)
    pos, entries_seen, matches = 16, 0, []
    old, new = old.encode("utf-8"), new.encode("utf-8")
    for _ in range(blocks):
        if pos + 36 > len(index):
            raise ValueError("truncated LuaV block header")
        block_digest = bytes(index[pos + 8:pos + 24])
        reserved, size, count = struct.unpack_from(">III", index, pos + 24)
        table = pos + 36
        end = table + count * 16
        if reserved != 0 or end > len(index):
            raise ValueError("unsupported LuaV block table")
        block_path = source / (block_digest.hex() + ".bytes")
        data = block_path.read_bytes()
        if len(data) != size or md5(data) != block_digest:
            raise ValueError("Lua block digest or size differs")
        entries_seen += count
        needle = bytes([len(old)]) + old
        if data.count(needle):
            if data.count(needle) != 1:
                raise ValueError("ambiguous path within Lua block")
            location = data.index(needle)
            entries = [struct.unpack_from(">QII", index, table + 16 * i) for i in range(count)]
            containing = [(i, off, length) for i, (_, length, off) in enumerate(entries)
                          if off <= location and location + len(needle) <= off + length]
            if len(containing) != 1:
                raise ValueError("Lua module index does not uniquely cover the string")
            matches.append((pos, table, data, entries, containing[0], location))
        if bytes(index[end:end + 3]) != b"\x00\x00\x80":
            raise ValueError("unsupported LuaV block trailer")
        pos = end + 3
    if entries_seen != total_modules or pos != len(index) or len(matches) != 1:
        raise ValueError("LuaV index size/count or matching block differs")
    pos, table, data, entries, (module, module_offset, module_len), location = matches[0]
    rewritten, _ = replace_once(data, old, new)
    delta = len(rewritten) - len(data)
    for i, (_, length, offset) in enumerate(entries):
        if i == module:
            length += delta
        elif offset >= module_offset + module_len:
            offset += delta
        elif offset < module_offset + module_len and offset + length > module_offset:
            raise ValueError("overlapping Lua module entries")
        if length < 0 or offset < 0 or offset + length > len(rewritten):
            raise ValueError("rewritten Lua module bounds invalid")
        struct.pack_into(">II", index, table + 16 * i + 8, length, offset)
    index[pos + 8:pos + 24] = md5(rewritten)
    struct.pack_into(">I", index, pos + 28, len(rewritten))
    new_index_digest = md5(index)
    updated_selector = bytearray(selector)
    for i in range(4):
        updated_selector[28 + i * 4:32 + i * 4] = new_index_digest[i * 4:i * 4 + 4][::-1]
    output.mkdir(parents=True, exist_ok=True)
    (output / (md5(rewritten).hex() + ".bytes")).write_bytes(rewritten)
    (output / ("LuaV_" + new_index_digest.hex() + ".bytes")).write_bytes(index)
    (output / "M_LuaV.bytes").write_bytes(updated_selector)
    return {"old_block_md5": md5(data).hex(), "new_block_md5": md5(rewritten).hex(),
            "old_index_md5": digest.hex(), "new_index_md5": new_index_digest.hex(),
            "module_offset": module_offset, "module_old_size": module_len,
            "module_new_size": module_len + delta, "string_offset": location}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("old_path")
    parser.add_argument("new_path")
    args = parser.parse_args()
    print(json.dumps(stage(args.source, args.output, args.old_path, args.new_path), indent=2))


if __name__ == "__main__":
    main()