import argparse
from pathlib import Path
import struct
import sys


def xxtea_mix(previous_word, next_word, total, key_word):
    return (
        (((previous_word << 4) ^ (next_word >> 3))
         + ((next_word << 2) ^ (previous_word >> 5)))
        ^ ((total ^ next_word) + (key_word ^ previous_word))
    ) & 0xFFFFFFFF


def decrypt_lua(data, key_words):
    if len(data) % 4:
        raise ValueError("size is not divisible by four")

    word_count = len(data) // 4
    if word_count < 2:
        return data

    words = list(struct.unpack(f"<{word_count}I", data))
    delta = 0x9E3779B9
    total = (delta * (9 + 56 // word_count)) & 0xFFFFFFFF
    next_word = words[0]

    while total:
        key_selector = (total >> 2) & 3

        for word_index in range(word_count - 1, 0, -1):
            previous_word = words[word_index - 1]
            next_word = (
                words[word_index]
                - xxtea_mix(
                    previous_word,
                    next_word,
                    total,
                    key_words[(key_selector ^ word_index) & 3],
                )
            ) & 0xFFFFFFFF
            words[word_index] = next_word

        previous_word = words[-1]
        next_word = (
            words[0]
            - xxtea_mix(
                previous_word,
                next_word,
                total,
                key_words[key_selector],
            )
        ) & 0xFFFFFFFF
        words[0] = next_word
        total = (total - delta) & 0xFFFFFFFF

    return struct.pack(f"<{word_count}I", *words)


def find_lua_files(source):
    if source.is_file():
        return [source]
    if source.is_dir():
        return sorted(source.rglob("*.lua"))
    raise FileNotFoundError("input was not found")


def main():
    parser = argparse.ArgumentParser(
        description="Decrypt Dynasty Legends 2 Lua files."
    )
    parser.add_argument(
        "input", type=Path, help="encrypted Lua file or folder"
    )
    parser.add_argument("output", type=Path, help="output file or folder")
    args = parser.parse_args()

    source = args.input.resolve()
    destination = args.output.resolve()
    files = find_lua_files(source)

    if not files:
        parser.error("no Lua files found")
    if source.is_dir() and destination.is_relative_to(source):
        parser.error("output folder must be outside the input folder")

    key_bytes = "钛核扫地僧2".encode("utf-8")
    key_words = struct.unpack("<4I", key_bytes)
    failed = 0

    for encrypted_file in files:
        if source.is_file():
            output_file = destination
        else:
            output_file = destination / encrypted_file.relative_to(source)

        if encrypted_file.resolve() == output_file.resolve():
            parser.error("input and output are the same file")

        try:
            plaintext = decrypt_lua(encrypted_file.read_bytes(), key_words)
            plaintext.decode("utf-8")
            output_file.parent.mkdir(parents=True, exist_ok=True)
            output_file.write_bytes(plaintext)
            print(output_file)
        except (OSError, UnicodeDecodeError, ValueError) as error:
            failed += 1
            print(f"Failed: {encrypted_file}: {error}", file=sys.stderr)

    passed = len(files) - failed
    print(f"\n{passed} file(s) decrypted; {failed} failed")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())