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

"""
Generates chunk JSON + companion DMF mesh files from raw chunk JSON files.

JSON input (assetsraw/<map>/chunks/chunk_X_Y_Z.json, one subdirectory per
map - e.g. assetsraw/overworld/chunks/):
  {
    "tiles": [
      { "pos": [x, y, z], "type": <tile_shape_int>, "tile": <uv_tile_int> }
    ],
    "meshes": [
      { "file": "house_5_3.dmf", "pos": [x, y, z] }
    ]
  }

  Tiles absent from the array default to TILE_SHAPE_NULL. Each (x, y)
  column may only have ONE real tile - chunks store one tile per column,
  not one per (x, y, z), and that tile records its own local z (see the
  chunk JSON format below). If two entries share the same (x, y) with
  different z, the later one in the array wins and a warning is printed.
  Mesh files are located by searching under assets/meshes/ and referenced by
  path from the assets root in the output chunk JSON.

  Output chunk JSON is derived automatically, preserving the map
  subdirectory:
    assetsraw/<map>/chunks/chunk_X_Y_Z.json -> assets/<map>/chunks/X_Y_Z.json

Chunk JSON format (loaded at runtime by assetChunkLoaderSync):
  {
    "tiles": [ [shape, z], ... ]  // exactly CHUNK_WIDTH * CHUNK_HEIGHT
                                  // entries, one per (x, y) column,
                                  // x-major (matches tile_t: { shape, z })
    "meshes": [
      { "model": "<path to .json model, relative to assets/>",
        "offset": [x, y, z] }
    ]
  }
  By convention mesh index 0 (when present) is the chunk's auto-generated
  terrain and the rest are props (see sceneOverworldDrawChunksBase/Props
  on the C side).

DMF format:
  Bytes 0-3:  DMF\\x00
  Bytes 4-7:  uint32_t version = 1 (little-endian)
  Bytes 8-11: uint32_t vertCount  (little-endian)
  Bytes 12+:  meshvertex_t vertices[vertCount]
              Each vertex: uv[2] + pos[3] = 5 x float32 LE = 20 bytes

Usage:
  python3 -m tools.asset.chunk              # process all assetsraw/*/chunks/*.json
  python3 -m tools.asset.chunk <file.json>  # process one JSON file
"""

import json
import math
import os
import struct
import sys

_HERE        = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.normpath(os.path.join(_HERE, '..', '..', '..'))
ASSETSRAW_DIR = os.path.join(PROJECT_ROOT, 'assetsraw')
ASSETS_DIR    = os.path.join(PROJECT_ROOT, 'assets')

CHUNK_WIDTH      = 16
CHUNK_HEIGHT     = 16
CHUNK_DEPTH      = 8
# One tile per (x, y) column - the column's local Z (0..CHUNK_DEPTH-1) is
# carried on the tile itself instead of storing a tile per (x, y, z).
CHUNK_TILE_COUNT = CHUNK_WIDTH * CHUNK_HEIGHT  # 256

# World-space Z distance of one Z-layer/story, in world-float space.
# Must match WORLD_LAYER_HEIGHT in src/dusk/rpg/overworld/worldpos.h.
WORLD_LAYER_HEIGHT = 1.0 / math.sqrt(2)

CHUNK_MESH_COUNT_MAX = 10
CHUNK_MESH_NAME_MAX  = 64

# Internal working format for a column's (shape, z) pair, packed/unpacked
# with struct purely so build_terrain_verts can index into it uniformly -
# unrelated to the on-disk format, which is JSON.
TILE_STRUCT_FORMAT = '<IB3x'
TILE_SIZE   = struct.calcsize(TILE_STRUCT_FORMAT)
VERTEX_SIZE = 20

TILE_SHAPE_NULL      = 0
TILE_SHAPE_GROUND    = 1
TILE_SHAPE_RAMP_NORTH     = 2
TILE_SHAPE_RAMP_EAST      = 3
TILE_SHAPE_RAMP_SOUTH     = 4
TILE_SHAPE_RAMP_WEST      = 5
TILE_SHAPE_RAMP_NORTHEAST = 6
TILE_SHAPE_RAMP_NORTHWEST = 7
TILE_SHAPE_RAMP_SOUTHEAST = 8
TILE_SHAPE_RAMP_SOUTHWEST = 9
TILE_SHAPE_RAMP_NORTHEAST_INNER = 10
TILE_SHAPE_RAMP_NORTHWEST_INNER = 11
TILE_SHAPE_RAMP_SOUTHEAST_INNER = 12
TILE_SHAPE_RAMP_SOUTHWEST_INNER = 13

DMF_MAGIC   = b'DMF\x00'
DMF_VERSION = 1


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def tile_index(x, y):
    return x + y * CHUNK_WIDTH


def find_model(filename):
    """Search ASSETS_DIR/models/ recursively for filename.
    Returns the path relative to ASSETS_DIR (forward slashes), or None."""
    models_root = os.path.join(ASSETS_DIR, 'models')
    for dirpath, _, filenames in os.walk(models_root):
        if filename in filenames:
            rel = os.path.relpath(
                os.path.join(dirpath, filename), ASSETS_DIR
            )
            return rel.replace('\\', '/')
    return None


def write_model_json(path, mesh_rel, texture_rel, color):
    """Write a model JSON descriptor file."""
    obj = {'mesh': mesh_rel, 'color': color}
    if texture_rel:
        obj['texture'] = texture_rel
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, 'w') as f:
        json.dump(obj, f, indent=2)
    print(f'  Wrote model JSON {path}')


def derive_chunk_json_path(json_path):
    """assetsraw/<map>/chunks/chunk_X_Y_Z.json -> assets/<map>/chunks/X_Y_Z.json"""
    base = os.path.splitext(os.path.basename(json_path))[0]
    if base.startswith('chunk_'):
        base = base[len('chunk_'):]
    rel_dir = os.path.relpath(os.path.dirname(os.path.abspath(json_path)), ASSETSRAW_DIR)
    return os.path.join(ASSETS_DIR, rel_dir, base + '.json')


def write_dmf(path, vertex_bytes):
    vert_count = len(vertex_bytes) // VERTEX_SIZE
    buf = bytearray()
    buf += DMF_MAGIC
    buf += struct.pack('<I', DMF_VERSION)
    buf += struct.pack('<I', vert_count)
    buf += vertex_bytes
    with open(path, 'wb') as f:
        f.write(buf)
    print(f'  Wrote DMF {path}: {vert_count} vertices, {len(buf)} bytes')


def write_chunk_json(chunk_json_path, tiles, mesh_names, mesh_offsets=None):
    """Write a chunk JSON file referencing the given model asset paths."""
    mesh_count = len(mesh_names)
    if mesh_offsets is None:
        mesh_offsets = [(0.0, 0.0, 0.0)] * mesh_count

    tile_list = []
    for i in range(CHUNK_TILE_COUNT):
        shape, z = struct.unpack_from(TILE_STRUCT_FORMAT, tiles, i * TILE_SIZE)
        tile_list.append([shape, z])

    meshes = []
    for name, offset in zip(mesh_names, mesh_offsets):
        if len(name) >= CHUNK_MESH_NAME_MAX:
            raise ValueError(
                f"Mesh name too long (>= {CHUNK_MESH_NAME_MAX}): {name}"
            )
        meshes.append({'model': name, 'offset': list(offset)})

    obj = {'tiles': tile_list, 'meshes': meshes}
    os.makedirs(os.path.dirname(chunk_json_path), exist_ok=True)
    with open(chunk_json_path, 'w') as f:
        json.dump(obj, f)
    print(
        f'  Wrote chunk JSON {chunk_json_path}: {mesh_count} mesh(es)'
    )


# ---------------------------------------------------------------------------
# JSON -> chunk JSON + DMF
# ---------------------------------------------------------------------------

def _tile_quad(u0, u1, v0, v1, fx, fy, sw_z, se_z, ne_z, nw_z):
    """Six vertices (2 CCW triangles) for a quad with per-corner Z heights.
    Corner order: SW=(fx,fy), SE=(fx+1,fy), NE=(fx+1,fy+1), NW=(fx,fy+1).
    NORTH=+Y, EAST=+X (matches entityDirGetRelative)."""
    verts = [
        (u0, v0, fx,   fy,   sw_z),
        (u1, v0, fx+1, fy,   se_z),
        (u1, v1, fx+1, fy+1, ne_z),
        (u0, v0, fx,   fy,   sw_z),
        (u1, v1, fx+1, fy+1, ne_z),
        (u0, v1, fx,   fy+1, nw_z),
    ]
    buf = bytearray()
    for v in verts:
        buf += struct.pack('<5f', *v)
    return bytes(buf)


# Per-shape corner heights as (sw, se, ne, nw) offsets from tile base Z.
# NORTH=+Y, EAST=+X.  Cardinal ramps: low face at 0, high face at 1.
# Diagonal outer-corner ramps: only the named corner raised, rest at 0.
# Diagonal inner-corner ramps: only the corner opposite the named one is
# lowered to 0, the rest (including the named corner) stay raised at 1.
_RAMP_CORNERS = {
    TILE_SHAPE_GROUND:         (0.0, 0.0, 0.0, 0.0),
    TILE_SHAPE_RAMP_NORTH:     (0.0, 0.0, 1.0, 1.0),  # south=low, north=high
    TILE_SHAPE_RAMP_SOUTH:     (1.0, 1.0, 0.0, 0.0),  # south=high, north=low
    TILE_SHAPE_RAMP_EAST:      (0.0, 1.0, 1.0, 0.0),  # west=low, east=high
    TILE_SHAPE_RAMP_WEST:      (1.0, 0.0, 0.0, 1.0),  # west=high, east=low
    TILE_SHAPE_RAMP_NORTHEAST: (0.0, 0.0, 1.0, 0.0),  # only NE raised
    TILE_SHAPE_RAMP_NORTHWEST: (0.0, 0.0, 0.0, 1.0),  # only NW raised
    TILE_SHAPE_RAMP_SOUTHEAST: (0.0, 1.0, 0.0, 0.0),  # only SE raised
    TILE_SHAPE_RAMP_SOUTHWEST: (1.0, 0.0, 0.0, 0.0),  # only SW raised
    TILE_SHAPE_RAMP_NORTHEAST_INNER: (0.0, 1.0, 1.0, 1.0),  # only SW low
    TILE_SHAPE_RAMP_NORTHWEST_INNER: (1.0, 0.0, 1.0, 1.0),  # only SE low
    TILE_SHAPE_RAMP_SOUTHEAST_INNER: (1.0, 1.0, 1.0, 0.0),  # only NW low
    TILE_SHAPE_RAMP_SOUTHWEST_INNER: (1.0, 1.0, 0.0, 1.0),  # only NE low
}


def build_terrain_verts(tiles_bytes):
    """Generate quads for every non-null tile, shaped to match tile type."""
    buf = bytearray()
    for y in range(CHUNK_HEIGHT):
        for x in range(CHUNK_WIDTH):
            idx = tile_index(x, y)
            tile_type, z = struct.unpack_from(
                '<IB', tiles_bytes, idx * TILE_SIZE
            )
            corners = _RAMP_CORNERS.get(tile_type)
            if corners is None:
                continue
            fx, fy, fz = float(x), float(y), float(z)
            u0 = x / CHUNK_WIDTH
            u1 = (x + 1) / CHUNK_WIDTH
            v0 = y / CHUNK_HEIGHT
            v1 = (y + 1) / CHUNK_HEIGHT
            sw, se, ne, nw = corners
            buf += _tile_quad(
                u0, u1, v0, v1, fx, fy,
                (fz + sw) * WORLD_LAYER_HEIGHT,
                (fz + se) * WORLD_LAYER_HEIGHT,
                (fz + ne) * WORLD_LAYER_HEIGHT,
                (fz + nw) * WORLD_LAYER_HEIGHT
            )
    return bytes(buf)


def from_json(json_path, chunk_json_path):
    with open(json_path) as f:
        data = json.load(f)

    tiles = bytearray(CHUNK_TILE_COUNT * TILE_SIZE)
    column_z = {}
    for tile in data.get('tiles', []):
        pos = tile['pos']
        x, y, z = int(pos[0]), int(pos[1]), int(pos[2])
        if (x, y) in column_z and column_z[(x, y)] != z:
            print(
                f'  WARNING: {json_path}: column ({x}, {y}) has tiles at '
                f'z={column_z[(x, y)]} and z={z} - only one tile per '
                f'column is kept, z={z} wins (last one in the array)'
            )
        column_z[(x, y)] = z
        struct.pack_into(
            TILE_STRUCT_FORMAT, tiles, tile_index(x, y) * TILE_SIZE,
            int(tile['type']), z
        )

    terrain_verts = build_terrain_verts(tiles)
    model_names = []
    mesh_offsets = []

    chunk_base = os.path.splitext(os.path.basename(chunk_json_path))[0]

    # Terrain mesh + model (only written if non-empty)
    terrain_dir = os.path.join(ASSETS_DIR, 'meshes', 'chunks')
    models_dir = os.path.join(ASSETS_DIR, 'models', 'chunks')
    os.makedirs(terrain_dir, exist_ok=True)
    os.makedirs(models_dir, exist_ok=True)
    if terrain_verts:
        dmf_name = f'chunk_{chunk_base}_0.dmf'
        write_dmf(os.path.join(terrain_dir, dmf_name), terrain_verts)
        mesh_rel = f'meshes/chunks/{dmf_name}'
        json_name = f'chunk_{chunk_base}_0.json'
        write_model_json(
            os.path.join(models_dir, json_name),
            mesh_rel,
            'tiles.png',
            [255, 255, 255, 255]
        )
        model_names.append(f'models/chunks/{json_name}')
        mesh_offsets.append((0.0, 0.0, 0.0))

    # External mesh references -> look up matching model JSON
    for mesh in data.get('meshes', []):
        filename = mesh['file']
        pos = mesh.get('pos', [0, 0, 0])
        model_filename = os.path.splitext(filename)[0] + '.json'
        rel = find_model(model_filename)
        if rel is None:
            raise ValueError(
                f"Model not found under assets/models/: {model_filename}"
            )
        model_names.append(rel)
        mesh_offsets.append((float(pos[0]), float(pos[1]), float(pos[2])))
        print(f'  Resolved {filename} -> {rel}')

    write_chunk_json(chunk_json_path, bytes(tiles), model_names, mesh_offsets)


def process_json(json_path):
    chunk_json_path = derive_chunk_json_path(json_path)
    os.makedirs(os.path.dirname(chunk_json_path), exist_ok=True)
    print(f"{json_path} -> {chunk_json_path}")
    from_json(json_path, chunk_json_path)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    args = sys.argv[1:]

    if not args:
        json_files = sorted(
            os.path.join(dirpath, f)
            for dirpath, _, filenames in os.walk(ASSETSRAW_DIR)
            if os.path.basename(dirpath) == 'chunks'
            for f in filenames
            if f.endswith('.json')
        )
        if not json_files:
            print(f"No chunk JSON files found under {ASSETSRAW_DIR}")
            sys.exit(0)
        for p in json_files:
            process_json(p)
        return

    src = args[0]
    if os.path.splitext(src)[1].lower() != '.json':
        print(f"Expected a .json input file, got: {src}")
        sys.exit(1)
    process_json(src)


if __name__ == '__main__':
    main()