import struct
import sys
import os
import json
from args import args
from assetcache import assetCache, assetGetCache
from assethelpers import getAssetRelativePath

CHUNK_WIDTH = 16
CHUNK_HEIGHT = 16
CHUNK_DEPTH = 8
CHUNK_TILE_COUNT = CHUNK_WIDTH * CHUNK_HEIGHT * CHUNK_DEPTH
TILE_WIDTH = 16.0
TILE_HEIGHT = 16.0
TILE_DEPTH = 11.36

def processTile(tileIndex, x=0, y=0, z=0):
  vertices = []
  indices = []

  # Placement X, Y, Z
  px = x * TILE_WIDTH
  py = y * TILE_HEIGHT
  pz = z * TILE_DEPTH

  if tileIndex == 0:
    # Tile 0, nothing
    return None
  
  elif tileIndex == 2:
    # Tile 2, ramp up 
    color = (255,0,0)
    vertices = [
      {'position': (px, py, pz + TILE_DEPTH), 'color': color, 'uv': (0,0)},      # 0,0
      {'position': (px + TILE_WIDTH, py, pz + TILE_DEPTH), 'color': color, 'uv': (1,0)}, # 1,0
      {'position': (px + TILE_WIDTH, py + TILE_HEIGHT, pz), 'color': color, 'uv': (1,1)}, # 1,1
      {'position': (px, py, pz + TILE_DEPTH), 'color': color, 'uv': (0,0)},      # 0,0 (repeat)
      {'position': (px + TILE_WIDTH, py + TILE_HEIGHT, pz), 'color': color, 'uv': (1,1)}, # 1,1 (repeat)
      {'position': (px, py + TILE_HEIGHT, pz), 'color': color, 'uv': (0,1)} # 0,1
    ]
    indices = [0, 1, 2, 3, 4, 5]

  else:
    # Determine color for checkerboard pattern
    if tileIndex == 1:
      color = (0, 255, 0)
    else:
      color = (0, 0, 255)

    vertices = [
      {'position': (px, py, pz), 'color': color, 'uv': (0,0)},      # 0,0
      {'position': (px + TILE_WIDTH, py, pz), 'color': color, 'uv': (1,0)}, # 1,0
      {'position': (px + TILE_WIDTH, py + TILE_HEIGHT, pz), 'color': color, 'uv': (1,1)}, # 1,1
      {'position': (px, py, pz), 'color': color, 'uv': (0,0)},      # 0,0 (repeat)
      {'position': (px + TILE_WIDTH, py + TILE_HEIGHT, pz), 'color': color, 'uv': (1,1)}, # 1,1 (repeat)
      {'position': (px, py + TILE_HEIGHT, pz), 'color': color, 'uv': (0,1)} # 0,1
    ]
    indices = [0, 1, 2, 3, 4, 5]

  return {
    'vertices': vertices,
    'indices': indices
  }

def processChunk(path):
  cache = assetGetCache(path)
  if cache:
    return cache

  # Read input file as JSON
  with open(path, 'r') as f:
    inData = json.load(f)

  chunk = {
    'tiles': [0] * CHUNK_TILE_COUNT,
    'models': []
  }

  baseModel = {
    'vertices': [],
    'indices': [],
    'vertexCount': 0,
    'indexCount': 0
  }

  # Append the model to chunk.models
  chunk['models'].append(baseModel)

  for i, tile in enumerate(inData['tiles']):
    # Set to chunk
    chunk['tiles'][i] = tile

    # Calculate x, y, z from i
    x = i % CHUNK_WIDTH
    y = (i // CHUNK_WIDTH) % CHUNK_HEIGHT
    z = i // (CHUNK_WIDTH * CHUNK_HEIGHT)

    # Add tile 3D model
    result = processTile(tile, x, y, z)
    if result is not None and len(result['vertices']) > 0:
      base = len(baseModel['vertices'])
      quad_indices = [base + idx for idx in result['indices']]
      baseModel['vertices'].extend(result['vertices'])
      baseModel['indices'].extend(quad_indices)
      baseModel['vertexCount'] = len(baseModel['vertices'])
      baseModel['indexCount'] = len(baseModel['indices'])

  # Generate binary buffer for efficient output
  buffer = bytearray()
  buffer.extend(b'DCF')# Header
  buffer.extend(len(chunk['tiles']).to_bytes(4, 'little'))  # Number of tiles
  buffer.extend(len(chunk['models']).to_bytes(1, 'little'))  # Number of models

  # Buffer tile data as array of uint8_t
  for tileIndex in chunk['tiles']:
    buffer.append(tileIndex.to_bytes(1, 'little')[0])

  # For each model
  for model in chunk['models']:
    # Write vertex count and index count
    buffer.extend(model['vertexCount'].to_bytes(4, 'little'))
    
    # For each vertex
    for vertex in model['vertices']:
      # This is not tightly packed in memory.
      # R G B A U V X Y Z
      # Color is 4 bytes (RGBA)
      # Rest is floats
      r, g, b = vertex['color']
      a = 255
      buffer.extend(r.to_bytes(1, 'little'))
      buffer.extend(g.to_bytes(1, 'little'))
      buffer.extend(b.to_bytes(1, 'little'))
      buffer.extend(a.to_bytes(1, 'little'))
      u, v = vertex['uv']
      buffer.extend(bytearray(struct.pack('<f', u)))
      buffer.extend(bytearray(struct.pack('<f', v)))
      x, y, z = vertex['position']
      buffer.extend(bytearray(struct.pack('<f', x)))
      buffer.extend(bytearray(struct.pack('<f', y)))
      buffer.extend(bytearray(struct.pack('<f', z)))

  # Write out map file
  relative = getAssetRelativePath(path)
  fileNameWithoutExt = os.path.splitext(os.path.basename(path))[0]
  outputFileRelative = os.path.join(os.path.dirname(relative), f"{fileNameWithoutExt}.dcf")
  outputFilePath = os.path.join(args.output_assets, outputFileRelative)
  os.makedirs(os.path.dirname(outputFilePath), exist_ok=True)
  with open(outputFilePath, "wb") as f:
    f.write(buffer)
  
  outChunk = {
    'files': [ outputFilePath ],
    'chunk': chunk
  }

  return assetCache(path, outChunk)


def processMap(asset):
  cache = assetGetCache(asset['path'])
  if cache is not None:
    return cache
  
  # Path provided should be a directory.
  if not os.path.isdir(asset['path']):
    print(f"Error: Asset path {asset['path']} is not a directory.")
    sys.exit(1)

  # List files
  chunkFiles = []
  for fileName in os.listdir(asset['path']):
    if not fileName.endswith('.json'):
      continue
    result = processChunk(os.path.join(asset['path'], fileName))
    chunkFiles.extend(result['files'])

  outMap = {
    'files': chunkFiles
  }

  return assetCache(asset['path'], outMap)