# 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.

A small fixed header (magic/version/size/checksum) followed by a single
uncompressed (ZIP_STORED) zip archive of every asset. Uncompressed rather
than DEFLATEd because the bulk of what's packed (audio, images) is
already compressed at the file level, so a second, zip-level DEFLATE
pass buys negligible space back for real CPU cost, and it lets every
entry - not just a hand-picked subset - support the reliable repeated
seeking/re-opening (e.g. locale line-by-line reads, MP3 loop restarts)
that libzip only supports for uncompressed entries.

DSK2 format:
  Bytes 0-3:  "DSK2" (magic)
  Bytes 4-7:  uint32_t version = 2 (little-endian)
  Bytes 8-11: uint32_t size (byte length of the zip archive that follows)
  Bytes 12-15: uint32_t checksum (CRC32 of the zip archive that follows)
  Bytes 16+:  the zip archive itself

Usage:
  python3 -m tools.asset.pack --input <assets_dir> --output <dusk.dsk>
    [--duskscriptcc <path to the duskscriptcc executable>]
  Every file under <assets_dir> is added, using its path relative to
  <assets_dir> as the zip entry name - except a .ds script, which is
  compiled with duskscriptcc first and added as the .dsb it compiles to
  (same path, extension swapped), not as its own source. --duskscriptcc is
  only required when a .ds file is actually found; a platform this project
  cross-compiles for without a host-built duskscriptcc (see the CMake side)
  can still pack assets that have none.

Called unconditionally on every CMake build (see the DUSK_ASSETS_BUILT
target) rather than gated behind file-level build-system dependencies -
CMake's CONFIGURE_DEPENDS glob only notices added/removed files on the
NEXT configure, so a plain custom_command DEPENDS list can miss edits to
existing files or lag a build behind on adds/removes. Instead this script
itself decides whether anything actually needs repacking: it writes a
manifest of every input file's (mtime, size) next to the output on each
successful pack, and skips the real work if a fresh manifest compares
equal to it and the output still exists. The manifest also carries the
format version, so a manifest written by an older version of this script
(a stale single/dual-archive dusk.dsk sitting from before a format
change) never compares equal and always triggers a fresh pack.
"""

import argparse
import json
import os
import struct
import subprocess
import sys
import tempfile
import zipfile
import zlib
import io

MAGIC = b'DSK'
VERSION = 2
HEADER_FORMAT = '<4sIII'
HEADER_SIZE = struct.calcsize(HEADER_FORMAT)


def script_arcname(relative_path):
  """.ds packs as the .dsb it compiles to, never as its own source."""
  return relative_path[:-len('.ds')] + '.dsb' if relative_path.endswith('.ds') else relative_path


def compile_script(duskscriptcc, full_path, relative_path, temp_dir):
  # Mirrors relative_path's own directories under temp_dir, so two .ds with
  # the same file name in different folders don't collide here.
  output_path = os.path.join(temp_dir, script_arcname(relative_path))
  os.makedirs(os.path.dirname(output_path), exist_ok=True)
  result = subprocess.run(
    [duskscriptcc, full_path, '-o', output_path, '--no-map'],
    capture_output=True, text=True
  )
  if result.returncode != 0:
    sys.stderr.write(result.stdout)
    sys.stderr.write(result.stderr)
    raise SystemExit(f'duskscriptcc failed to compile {relative_path}')
  with open(output_path, 'rb') as f:
    return f.read()


def build_zip_blob(root, relative_paths, duskscriptcc):
  buf = io.BytesIO()
  with tempfile.TemporaryDirectory() as temp_dir, \
    zipfile.ZipFile(buf, 'w', compression=zipfile.ZIP_STORED) as zf:
    for relative_path in sorted(relative_paths):
      full_path = os.path.join(root, relative_path)
      if relative_path.endswith('.ds'):
        if not duskscriptcc:
          raise SystemExit(
            f'{relative_path} needs duskscriptcc to compile, but --duskscriptcc was not given'
          )
        zf.writestr(script_arcname(relative_path), compile_script(duskscriptcc, full_path, relative_path, temp_dir))
      else:
        zf.write(full_path, arcname=relative_path)
  return buf.getvalue()


def manifest_path(output_path):
  return output_path + '.manifest.json'


def compute_manifest(input_dir, relative_paths, duskscriptcc):
  entries = {}
  for relative_path in relative_paths:
    st = os.stat(os.path.join(input_dir, relative_path))
    entries[relative_path] = [st.st_mtime_ns, st.st_size]
  # A duskscriptcc rebuild (format/bugfix) has to invalidate every .ds in the
  # manifest too, even though none of their own mtimes changed.
  compiler_mtime = os.stat(duskscriptcc).st_mtime_ns if duskscriptcc else None
  return {'version': VERSION, 'compiler': compiler_mtime, 'files': entries}


def load_manifest(path):
  try:
    with open(path, 'r', encoding='utf-8') as f:
      return json.load(f)
  except (OSError, ValueError):
    return None


def pack(input_dir, output_path, duskscriptcc=None):
  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, '/')
      )

  manifest = compute_manifest(input_dir, relative_paths, duskscriptcc)
  manifest_file = manifest_path(output_path)
  if os.path.isfile(output_path) and load_manifest(manifest_file) == manifest:
    print(f'{output_path} is up to date, skipping pack')
    return

  blob = build_zip_blob(input_dir, relative_paths, duskscriptcc)

  header = struct.pack(
    HEADER_FORMAT, MAGIC, VERSION, len(blob), zlib.crc32(blob) & 0xFFFFFFFF
  )

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

  with open(manifest_file, 'w', encoding='utf-8') as f:
    json.dump(manifest, f)

  print(
    f'Wrote {output_path}: {len(relative_paths)} file(s) '
    f'({len(blob)} bytes), {len(header) + len(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(
    '--duskscriptcc', default=None,
    help='Path to the duskscriptcc executable, to compile .ds assets with. '
         'Only required when the assets directory actually has a .ds in it.'
  )
  args = parser.parse_args()

  pack(args.input, args.output, args.duskscriptcc)


if __name__ == '__main__':
  main()