#Compressed Chunk Export

#!/usr/bin/env python3
import csv
import json
import math
import struct
import sys
import zipfile
import zlib
from collections.abc import Iterator
from pathlib import Path
from typing import Any

Frame = dict[str, Any]
GridFormat = dict[str, int]

def grid_format(bits_per_cell: int) -> GridFormat:
  """Return packing constants for a bits-per-cell value.

  :param bits_per_cell: Number of bits used by one packed cell.
  :returns: A dictionary with the cell mask and cells-per-word count.
  :rtype: dict[str, int]
  """
  cells_per_word = 32 // bits_per_cell
  return {
    "bits_per_cell": bits_per_cell,
    "cells_per_word": cells_per_word,
    "cell_mask": (1 << bits_per_cell) - 1,
  }

def frame_byte_size(cols: int, rows: int, bits_per_cell: int) -> int:
  """Return the byte length of one packed frame.

  :param cols: Number of grid columns.
  :param rows: Number of grid rows.
  :param bits_per_cell: Number of bits used by one packed cell.
  :returns: The packed frame size in bytes.
  :rtype: int
  """
  fmt = grid_format(bits_per_cell)
  packed_cols = math.ceil(cols / fmt["cells_per_word"])
  return packed_cols * rows * 4

def inflate_chunk(data: bytes, codec: str) -> bytes:
  """Decode a chunk payload according to its manifest codec.

  :param data: Compressed or raw bytes read from the ZIP entry.
  :param codec: Codec name from the chunk manifest entry.
  :returns: Raw packed frame bytes.
  :rtype: bytes
  :raises ValueError: If the codec is not supported by this script.
  """
  if codec == "deflate-raw":
    return zlib.decompress(data, wbits=-zlib.MAX_WBITS)
  if codec == "raw-packed" or codec == "stored-packed":
    return data
  raise ValueError(f"Unsupported chunk codec: {codec}")

def unpack_frame(frame_bytes: bytes, cols: int, rows: int, bits_per_cell: int) -> list[list[int]]:
  """Unpack one frame into a two-dimensional list of tribe indexes.

  :param frame_bytes: Raw packed frame bytes.
  :param cols: Number of grid columns.
  :param rows: Number of grid rows.
  :param bits_per_cell: Number of bits used by one packed cell.
  :returns: A list of rows, where each value is a tribe index.
  :rtype: list[list[int]]
  """
  fmt = grid_format(bits_per_cell)
  packed_cols = math.ceil(cols / fmt["cells_per_word"])
  values: list[list[int]] = []

  for y in range(rows):
    row: list[int] = []
    for x in range(cols):
      word_index = y * packed_cols + (x // fmt["cells_per_word"])
      word = struct.unpack_from("<I", frame_bytes, word_index * 4)[0]
      shift = (x % fmt["cells_per_word"]) * bits_per_cell
      row.append((word >> shift) & fmt["cell_mask"])
    values.append(row)

  return values

def iter_frames(zip_path: str | Path) -> Iterator[Frame]:
  """Yield decoded frames from a compressed chunk export ZIP.

  :param zip_path: Path to the exported ZIP file.
  :returns: An iterator of frame dictionaries with generation, values, and metadata.
  :rtype: Iterator[dict[str, object]]
  :raises ValueError: If a decoded chunk has an unexpected byte length.
  """
  with zipfile.ZipFile(zip_path) as zf:
    manifest = json.loads(zf.read("manifest.json"))
    metadata = json.loads(zf.read("metadata.json"))
    cols = metadata["cols"]
    rows = metadata["rows"]
    tribes = metadata["tribes"]

    for chunk in manifest["chunks"]:
      bits = chunk.get("gridFormat", metadata["gridFormat"])["bitsPerCell"]
      frame_size = frame_byte_size(cols, rows, bits)
      raw_entry = zf.read(f"chunks/{chunk['filename']}")
      payload = inflate_chunk(raw_entry, chunk["codec"])
      expected_size = frame_size * chunk["blockCount"]

      if len(payload) != expected_size:
        raise ValueError(
          f"{chunk['filename']} decoded to {len(payload)} bytes, expected {expected_size}"
        )

      for local_index, generation in enumerate(chunk["generations"]):
        start = local_index * frame_size
        end = start + frame_size
        values = unpack_frame(payload[start:end], cols, rows, bits)
        yield {
          "generation": generation,
          "values": values,
          "tribes": tribes,
          "cols": cols,
          "rows": rows,
        }

def write_first_frame_csv(zip_path: str | Path, output_path: str | Path) -> None:
  """Write the first exported frame to a CSV file using tribe IDs.

  :param zip_path: Path to the exported ZIP file.
  :param output_path: Path where the CSV file should be written.
  :returns: ``None``. Writes the CSV file to disk.
  :rtype: None
  :raises ValueError: If the export does not contain any frames.
  """
  first = next(iter_frames(zip_path), None)
  if first is None:
    raise ValueError("No frames found in export")

  tribes = first["tribes"]
  with open(output_path, "w", newline="", encoding="utf-8") as handle:
    writer = csv.writer(handle)
    for row in first["values"]:
      writer.writerow([
        tribes[value]["id"] if value < len(tribes) else f"unknown:{value}"
        for value in row
      ])

  print(f"Wrote generation {first['generation']} to {output_path}")

def main() -> None:
  """Read an export ZIP and write its first frame as CSV.

  :returns: ``None``. Writes the first exported frame to disk.
  :rtype: None
  :raises SystemExit: If the command line arguments are invalid.
  """
  if len(sys.argv) != 2:
    raise SystemExit("Usage: python read_chunk_export.py golt-export.zip")

  zip_path = Path(sys.argv[1])
  with zipfile.ZipFile(zip_path) as zf:
    manifest = json.loads(zf.read("manifest.json"))
    metadata = json.loads(zf.read("metadata.json"))
    print(f"Grid: {metadata['cols']} x {metadata['rows']}")
    print(f"Topology: {metadata.get('topology', 'toroidal')}")
    print(f"Random seed: {metadata.get('randomSeed', 42)}")
    print(f"Frames: {manifest['selectedRange']['framesTotal']}")
    print(f"Chunks: {len(manifest['chunks'])}")
    print("Tribes:", ", ".join(tribe["id"] for tribe in metadata["tribes"]))

  write_first_frame_csv(zip_path, "first-frame.csv")

if __name__ == "__main__":
  main()