import argparse
import os
import struct
import zipfile
from pathlib import Path

QOLPAK_HEADER = struct.Struct("<8sI")
QOLPAK_ENTRY = struct.Struct("<II256sI")

def decompress_quicklz(data: bytes) -> bytes:
    header_size = 9 if data[0] & 2 else 3

    # Both QuickLZ header forms turn up in the PAKs.
    if header_size == 9:
        packed_size, output_size = struct.unpack_from("<II", data, 1)
    else:
        packed_size, output_size = data[1], data[2]

    if packed_size != len(data):
        raise ValueError("bad QuickLZ packet size")

    if not data[0] & 1:
        return data[header_size:header_size + output_size]

    level = (data[0] >> 2) & 3
    if level != 3:
        raise ValueError(f"unexpected QuickLZ level: {level}")
    if output_size == 0:
        return b""

    src = header_size
    dst = 0
    cword = 1
    fetch = 0
    output = bytearray(output_size)
    last_match = output_size - 11

    while True:
        if cword == 1:
            cword = struct.unpack_from("<I", data, src)[0]
            src += 4
            if dst <= last_match:
                fetch = struct.unpack_from("<I", data, src)[0]

        if cword & 1:
            cword >>= 1

            # Level 3 has five match forms. The low bits select one.
            if fetch & 3 == 0:
                offset = (fetch & 0xFF) >> 2
                length = 3
                src += 1
            elif fetch & 2 == 0:
                offset = (fetch & 0xFFFF) >> 2
                length = 3
                src += 2
            elif fetch & 1 == 0:
                offset = (fetch & 0xFFFF) >> 6
                length = ((fetch >> 2) & 15) + 3
                src += 2
            elif fetch & 127 != 3:
                offset = (fetch >> 7) & 0x1FFFF
                length = ((fetch >> 2) & 31) + 2
                src += 3
            else:
                offset = fetch >> 15
                length = ((fetch >> 7) & 255) + 3
                src += 4

            if offset == 0 or offset > dst:
                raise ValueError("bad QuickLZ match")

            for _ in range(length):
                output[dst] = output[dst - offset]
                dst += 1

            if dst <= last_match:
                fetch = struct.unpack_from("<I", data, src)[0]
            continue

        if dst <= last_match:
            output[dst] = data[src]
            dst += 1
            src += 1
            cword >>= 1
            fetch = (
                ((fetch >> 8) & 0xFFFF)
                | (data[src + 2] << 16)
                | (data[src + 3] << 24)
            )
            continue

        while dst < output_size:
            if cword == 1:
                src += 4
                cword = 0x80000000
            output[dst] = data[src]
            dst += 1
            src += 1
            cword >>= 1
        return bytes(output)

def extract_qolpak(archive_path: Path, output_path: Path) -> int:
    entries = []

    with archive_path.open("rb") as archive:
        magic, entry_count = QOLPAK_HEADER.unpack(
            archive.read(QOLPAK_HEADER.size)
        )
        if magic != b"QOLPAK10":
            raise ValueError("not a QOLPAK10 archive")

        # The directory stores two sizes, a 256-byte name and the data offset.
        for _ in range(entry_count):
            record = archive.read(QOLPAK_ENTRY.size)
            packed, unpacked, raw_name, offset = QOLPAK_ENTRY.unpack(record)
            name = raw_name.split(b"\0", 1)[0].decode("ascii")
            entries.append((name, packed, unpacked, offset))

        for name, packed, unpacked, offset in entries:
            archive.seek(offset)
            data = decompress_quicklz(archive.read(packed))
            if len(data) != unpacked:
                raise ValueError(f"size mismatch after extracting {name}")

            relative = os.path.normpath(name.replace("\\", os.sep))
            drive, relative = os.path.splitdrive(relative)
            if (
                drive
                or os.path.isabs(relative)
                or relative == ".."
                or relative.startswith(".." + os.sep)
            ):
                raise ValueError(f"bad filename in archive: {name}")

            destination = output_path / relative
            destination.parent.mkdir(parents=True, exist_ok=True)
            destination.write_bytes(data)
            print(name)

    return len(entries)

def extract_zip(archive_path: Path, output_path: Path) -> int:
    with zipfile.ZipFile(archive_path) as archive:
        entries = archive.infolist()
        archive.extractall(output_path)

    for entry in entries:
        if not entry.is_dir():
            print(entry.filename)

    return sum(not entry.is_dir() for entry in entries)

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract files from an Alganon PAK archive."
    )
    parser.add_argument("archive", type=Path, help="PAK file to extract")
    parser.add_argument("output", type=Path, help="output directory")
    arguments = parser.parse_args()

    try:
        with arguments.archive.open("rb") as archive:
            signature = archive.read(8)

        arguments.output.mkdir(parents=True, exist_ok=True)

        # The .pak suffix covers two different formats in Alganon.
        if signature == b"QOLPAK10":
            extracted = extract_qolpak(arguments.archive, arguments.output)
        elif zipfile.is_zipfile(arguments.archive):
            extracted = extract_zip(arguments.archive, arguments.output)
        else:
            raise ValueError("unknown Alganon PAK format")
    except (OSError, ValueError, IndexError, struct.error) as error:
        parser.exit(1, f"{parser.prog}: {error}\n")

    print(f"Extracted {extracted} files to {arguments.output.resolve()}")

if __name__ == "__main__":
    main()