#!/usr/bin/env python3
"""Extract a single compressed RED4 archive segment by resource hash."""

from __future__ import annotations

import argparse
import struct
from pathlib import Path


FILE_ENTRY_SIZE = 56
FILE_SEGMENT_SIZE = 16


def parse_hash(value: str) -> int:
    return int(value, 0)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("archive", type=Path)
    parser.add_argument("resource_hash", type=parse_hash)
    parser.add_argument("outpath", type=Path)
    args = parser.parse_args()

    with args.archive.open("rb") as archive:
        header = archive.read(40)
        if len(header) != 40:
            raise SystemExit("archive header is too short")
        magic, _version, index_pos, _index_size, *_rest = struct.unpack("<IIQIQIQ", header)
        if magic != 1380009042:
            raise SystemExit(f"unexpected archive magic: {magic}")

        archive.seek(index_pos + 16)
        file_count, segment_count, _dependency_count = struct.unpack("<III", archive.read(12))

        found: tuple[int, int] | None = None
        for _ in range(file_count):
            entry = archive.read(FILE_ENTRY_SIZE)
            (
                key,
                _timestamp,
                _flags,
                segments_start,
                segments_end,
                _dependencies_start,
                _dependencies_end,
                _sha1,
            ) = struct.unpack("<QqIIIII20s", entry)
            if key == args.resource_hash:
                if segments_end - segments_start != 1:
                    raise SystemExit(
                        f"resource has {segments_end - segments_start} segments; expected 1"
                    )
                found = (segments_start, segments_end)

        if found is None:
            raise SystemExit(f"resource hash {args.resource_hash} not found")

        segments = []
        for _ in range(segment_count):
            segments.append(struct.unpack("<QII", archive.read(FILE_SEGMENT_SIZE)))

        offset, z_size, size = segments[found[0]]
        archive.seek(offset)
        data = archive.read(z_size)
        if len(data) != z_size:
            raise SystemExit("archive ended before segment data was fully read")

    args.outpath.parent.mkdir(parents=True, exist_ok=True)
    args.outpath.write_bytes(data)
    print(f"wrote {args.outpath} z_size={z_size} size={size}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())