# Copyright (c) 2026 Dominic Masters
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT

"""
Packs the assets directory into a dusk.dsk (DSK2) archive.

Replaces a single plain zip with two independent zip archives back to
back, so specific files can be stored uncompressed instead of DEFLATEd -
useful for small files (locale strings, config) accessed at runtime via
repeated seeks/re-opens, which is only reliably supported by libzip for
uncompressed (STORED) entries. Reading a whole compressed archive into
memory just to get reliable access isn't viable once that archive holds
the bulk of a game's binary assets, so the two archives are kept
separate and only the (small, by convention) stored one is expected to
ever need to be buffered whole.

DSK2 format:
  Bytes 0-3:   "DSK2" (magic)
  Bytes 4-7:   uint32_t version = 1 (little-endian)
  Bytes 8-11:  uint32_t compressedOffset
  Bytes 12-15: uint32_t compressedSize
  Bytes 16-19: uint32_t compressedChecksum (CRC32 of the compressed blob)
  Bytes 20-23: uint32_t storedOffset
  Bytes 24-27: uint32_t storedSize
  Bytes 28-31: uint32_t storedChecksum (CRC32 of the stored blob)
  Bytes 32+:   compressed zip blob, then stored zip blob (each a
               complete, independently valid zip archive)

Usage:
  python3 -m tools.asset.pack --input <assets_dir> --output <dusk.dsk>
    [--stored <glob-pattern>]...
  Every file under <assets_dir> is added, using its path relative to
  <assets_dir> as the zip entry name, to the stored archive if its
  relative path matches any --stored pattern (fnmatch, default:
  "locale/*"), otherwise to the compressed archive.
"""

import argparse
import os
import struct
import zipfile
import zlib
import fnmatch
import io

MAGIC = b'DSK2'
VERSION = 1
HEADER_FORMAT = '<4sIIIIIII'
HEADER_SIZE = struct.calcsize(HEADER_FORMAT)
DEFAULT_STORED_PATTERNS = ['locale/*']


def build_zip_blob(root, relative_paths, compression):
  buf = io.BytesIO()
  with zipfile.ZipFile(buf, 'w', compression=compression) as zf:
    for relative_path in sorted(relative_paths):
      zf.write(os.path.join(root, relative_path), arcname=relative_path)
  return buf.getvalue()


def pack(input_dir, output_path, stored_patterns):
  relative_paths = []
  for dirpath, _dirnames, filenames in os.walk(input_dir):
    for filename in filenames:
      full_path = os.path.join(dirpath, filename)
      relative_paths.append(
        os.path.relpath(full_path, input_dir).replace(os.sep, '/')
      )

  stored_paths = [
    path for path in relative_paths
    if any(fnmatch.fnmatch(path, pattern) for pattern in stored_patterns)
  ]
  compressed_paths = [
    path for path in relative_paths if path not in stored_paths
  ]

  compressed_blob = build_zip_blob(
    input_dir, compressed_paths, zipfile.ZIP_DEFLATED
  )
  stored_blob = build_zip_blob(input_dir, stored_paths, zipfile.ZIP_STORED)

  compressed_offset = HEADER_SIZE
  stored_offset = compressed_offset + len(compressed_blob)

  header = struct.pack(
    HEADER_FORMAT,
    MAGIC,
    VERSION,
    compressed_offset,
    len(compressed_blob),
    zlib.crc32(compressed_blob) & 0xFFFFFFFF,
    stored_offset,
    len(stored_blob),
    zlib.crc32(stored_blob) & 0xFFFFFFFF,
  )

  with open(output_path, 'wb') as f:
    f.write(header)
    f.write(compressed_blob)
    f.write(stored_blob)

  print(
    f'Wrote {output_path}: {len(compressed_paths)} compressed file(s) '
    f'({len(compressed_blob)} bytes), {len(stored_paths)} stored file(s) '
    f'({len(stored_blob)} bytes), {len(header) + len(compressed_blob) + len(stored_blob)} bytes total'
  )


def main():
  parser = argparse.ArgumentParser(
    description='Pack the assets directory into a dusk.dsk (DSK2) archive.'
  )
  parser.add_argument('--input', required=True, help='Assets directory to pack')
  parser.add_argument('--output', required=True, help='Path to write dusk.dsk to')
  parser.add_argument(
    '--stored',
    action='append',
    dest='stored_patterns',
    help='fnmatch pattern (relative to --input) of files to store '
         'uncompressed instead of compressing. May be given multiple '
         'times. Defaults to "locale/*" if never given.'
  )
  args = parser.parse_args()

  pack(args.input, args.output, args.stored_patterns or DEFAULT_STORED_PATTERNS)


if __name__ == '__main__':
  main()