import sys, os
import argparse
from datetime import datetime
import json
import math

# Values defined within C
CHUNK_WIDTH = 8
CHUNK_HEIGHT = 8
CHUNK_TILE_COUNT = CHUNK_WIDTH * CHUNK_HEIGHT
CHUNK_ENTITY_COUNT_MAX = 8
TILE_WIDTH = 16
TILE_HEIGHT = 16

ENTITY_TYPE_MAP = {
  "npc": "ENTITY_TYPE_NPC",
}

# Check if the script is run with the correct arguments
parser = argparse.ArgumentParser(description="Generate chunk header files")
parser.add_argument('--output', required=True, help='Dir to output headers')
parser.add_argument('--input', required=True, help='Input JSON file from tiled')
args = parser.parse_args()

# Ensure outdir exists
outputDir = args.output
os.makedirs(outputDir, exist_ok=True)

# Create world directory if it does not exist
worldDir = os.path.join(outputDir, "world")
os.makedirs(worldDir, exist_ok=True)

# Create chunks directory if it does not exist
chunksDir = os.path.join(worldDir, "chunk")
os.makedirs(chunksDir, exist_ok=True)

# Some vars used during printing
now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

# Read the input JSON file
inputFile = args.input
if not os.path.isfile(inputFile):
  print(f"Error: Input file '{inputFile}' does not exist.")
  sys.exit(1)

with open(inputFile, 'r') as f:
  data = json.load(f)
  
# Data should have height key
if 'height' not in data or 'width' not in data:
  print(f"Error: Input file '{inputFile}' does not contain 'height' or 'width' key.")
  sys.exit(1)

if 'tilewidth' not in data or 'tileheight' not in data:
  print(f"Error: Input file '{inputFile}' does not contain 'tilewidth' or 'tileheight' key.")
  sys.exit(1)

if 'infinite' not in data or not isinstance(data['infinite'], bool):
  print(f"Error: Input file '{inputFile}' does not contain 'infinite' key.")
  sys.exit(1)

# Need layers
if 'layers' not in data or not isinstance(data['layers'], list):
  print(f"Error: Input file '{inputFile}' does not contain 'layers' key.")
  sys.exit(1)

layers = data['layers']
if len(layers) == 0:
  print(f"Error: Input file '{inputFile}' does not contain any layers.")
  sys.exit(1)

# Object layer
objectLayer = None
for layer in layers:
  if layer.get('type') == 'objectgroup':
    objectLayer = layer
    break

if objectLayer is None:
  print(f"Error: Input file '{inputFile}' does not contain an object layer.")
  sys.exit(1)

if 'objects' not in objectLayer or not isinstance(objectLayer['objects'], list):
  print(f"Error: Object layer in '{inputFile}' does not contain 'objects' key or it is not a list.")
  sys.exit(1)

# Tile Layers
tileLayers = []
for layer in layers:
  if layer.get('type') == 'tilelayer':
    tileLayers.append(layer)

if len(tileLayers) == 0:
  print(f"Error: Input file '{inputFile}' does not contain any tile layers.")
  sys.exit(1)

# First layer
firstLayer = tileLayers[0]
if 'width' not in firstLayer or 'height' not in firstLayer:
  print(f"Error: First layer in '{inputFile}' does not contain 'width' or 'height' key.")
  sys.exit(1)

if 'chunks' not in firstLayer or not isinstance(firstLayer['chunks'], list):
  print(f"Error: First layer in '{inputFile}' does not contain 'chunks' key.")
  sys.exit(1)

if len(firstLayer['chunks']) == 0:
  print(f"Error: First layer in '{inputFile}' does not contain any chunks.")
  sys.exit(1)

firstLayerFirstChunk = firstLayer['chunks'][0]

# Now determine the input map bounds.
isMinXFound = False
isMaxXFound = False
isMinYFound = False
isMaxYFound = False
inputMapLowestX = 0
inputMapHighestX = 0
inputMapLowestY = 0
inputMapHighestY = 0
inputLayerWidthInTiles = firstLayerFirstChunk['width']
inputLayerHeightInTiles = firstLayerFirstChunk['height']

for chunk in firstLayer['chunks']:
  if 'x' not in chunk or 'y' not in chunk:
    print(f"Error: Chunk in first layer does not contain 'x' or 'y' key.")
    sys.exit(1)

  # Check chunk is not empty
  if 'data' not in chunk or not isinstance(chunk['data'], list):
    print(f"Error: Chunk in first layer does not contain 'data' key or it is not a list.")
    sys.exit(1)

  if len(chunk['data']) != inputLayerWidthInTiles * inputLayerHeightInTiles:
    print(f"Error: Chunk in first layer does not contain the expected number of tiles ({inputLayerWidthInTiles * inputLayerHeightInTiles}).")
    sys.exit(1)

  chunkEmpty = True
  for tile in chunk['data']:
    if tile == 0:
      continue
    chunkEmpty = False
    break

  if chunkEmpty:
    print(f"Warning: Chunk at ({chunk['x']}, {chunk['y']}) is empty, skipping.")
    continue
    
  chunkX = chunk['x']
  chunkY = chunk['y']

  if inputMapLowestX > chunkX or not isMinXFound:
    inputMapLowestX = chunkX
    isMinXFound = True
  if inputMapHighestX < chunkX or not isMaxXFound:
    inputMapHighestX = chunkX
    isMaxXFound = True

  if inputMapLowestY > chunkY or not isMinYFound:
    inputMapLowestY = chunkY
    isMinYFound = True
  if inputMapHighestY < chunkY or not isMaxYFound:
    inputMapHighestY = chunkY
    isMaxYFound = True

inputMapHighestX += inputLayerWidthInTiles
inputMapHighestY += inputLayerHeightInTiles

print(f"Input map lowest X: {inputMapLowestX}, highest X: {inputMapHighestX}")
print(f"Input map lowest Y: {inputMapLowestY}, highest Y: {inputMapHighestY}")

# We now offset all chunks by the lowest X/Y values to make them start at (0, 0).
for layerIndex, layer in enumerate(tileLayers):
  if layer['startx'] != inputMapLowestX or layer['starty'] != inputMapLowestY:
    continue
  
  for chunkIndex, chunk in enumerate(layer['chunks']):
    chunk['x'] -= inputMapLowestX
    chunk['y'] -= inputMapLowestY
    layer['chunks'][chunkIndex] = chunk

  layers[layerIndex] = layer


# Pre generate entity data
nextEntityId = 1
for obIndex, ob in enumerate(objectLayer['objects']):
  if 'x' not in ob or 'y' not in ob:
    print(f"Error: Object in object layer does not contain 'x' or 'y' key.")
    sys.exit(1)
  
  ob['x'] -= inputMapLowestX * TILE_WIDTH
  ob['y'] -= inputMapLowestY * TILE_HEIGHT

  # Objects are bottom aligned in tiled, so we need to adjust the Y coordinate.
  ob['y'] -= TILE_HEIGHT

  # Round off the coordinates
  ob['x'] = round(ob['x'])
  ob['y'] = round(ob['y'])
  ob['id'] = nextEntityId
  nextEntityId += 1

  objectLayer['objects'][obIndex] = ob

mapWidthInTiles = inputMapHighestX - inputMapLowestX
mapHeightInTiles = inputMapHighestY - inputMapLowestY
mapWidthInRealChunks =  math.ceil(float(mapWidthInTiles) / float(CHUNK_WIDTH))
mapHeightInRealChunks = math.ceil(float(mapHeightInTiles) / float(CHUNK_HEIGHT))
print(f"Map width in chunks: {mapWidthInRealChunks}, height in chunks: {mapHeightInRealChunks}")

if inputLayerWidthInTiles < CHUNK_WIDTH or inputLayerHeightInTiles < CHUNK_HEIGHT:
  print(f"Error: Input layer size {inputLayerWidthInTiles}x{inputLayerHeightInTiles} is smaller than chunk size {CHUNK_WIDTH}x{CHUNK_HEIGHT}.")
  sys.exit(1)

# For each output chunk.
worldWidth = 0
worldHeight = 0
chunksDone = set()

for chunkY in range(mapHeightInRealChunks):
  for chunkX in range(mapWidthInRealChunks):
    # Top left X/Y based on real chunk size
    topLeftTileX = chunkX * CHUNK_WIDTH
    topLeftTileY = chunkY * CHUNK_HEIGHT

    # Top left coordinates based on input layer size
    inputTopLeftTileX = math.floor(float(topLeftTileX) / float(inputLayerWidthInTiles)) * inputLayerWidthInTiles
    inputTopLeftTileY = math.floor(float(topLeftTileY) / float(inputLayerHeightInTiles)) * inputLayerHeightInTiles

    # Get the layers for this chunk.
    chunkLayers = []
    for layerIndex, layer in enumerate(tileLayers):
      foundChunk = None
      if 'chunks' not in layer or not isinstance(layer['chunks'], list):
        print(f"Error: Layer {layerIndex} in '{inputFile}' does not contain 'chunks' key.")
        sys.exit(1)

      # Find the chunk in this layer that matches the output chunk coordinates.
      chunks = layer['chunks']
      for chunk in chunks:
        if 'x' not in chunk or 'y' not in chunk:
          print(f"Error: Chunk in layer {layerIndex} does not contain 'x' or 'y' key.")
          sys.exit(1)

        if chunk['x'] == inputTopLeftTileX and chunk['y'] == inputTopLeftTileY:
          foundChunk = chunk
          break

      # If we did not find a chunk for this layer, append None.
      if foundChunk is None:
        chunkLayers.append(None)
        continue

      # Is this chunk layer just empty?
      layerEmpty = True
      for tile in foundChunk.get('data', []):
        if tile == 0:
          continue
        layerEmpty = False
        break
      
      if layerEmpty:
        chunkLayers.append(None)
        continue
        
      # Append the found chunk to the chunkLayers list.
      chunkLayers.append(foundChunk)
      
    # Now we have a chunkLayers list with the found chunks for each layer.
    if all(chunk is None for chunk in chunkLayers) or len(chunkLayers) == 0:
      continue

    # If we have more than 2 layers, we cannot handle this (yet).
    if len(chunkLayers) > 2:
      print(f"Error: Expected 2 layers for chunk at ({chunkX}, {chunkY}), found {len(chunkLayers)}.")
      sys.exit(1)

    entities = []
    for ob in objectLayer['objects']:
      if 'x' not in ob or 'y' not in ob:
        print(f"Error: Object in object layer does not contain 'x' or 'y' key.")
        sys.exit(1)
        
      # Is this object within the chunk?
      if ob['x'] < topLeftTileX * TILE_WIDTH:
        continue
      if ob['x'] >= (topLeftTileX + CHUNK_WIDTH) * TILE_WIDTH:
        continue
      if ob['y'] < topLeftTileY * TILE_HEIGHT:
        continue
      if ob['y'] >= (topLeftTileY + CHUNK_HEIGHT) * TILE_HEIGHT:
        continue
      
      entities.append(ob)
    
    # Shorthand functions
    def getInputLocalTileX(absoluteTileX):
      return absoluteTileX % inputLayerWidthInTiles
    
    def getInputLocalTileY(absoluteTileY):
      return absoluteTileY % inputLayerHeightInTiles
    
    # Determine base layer data.
    layerBase = chunkLayers[0]
    layerBaseData = []
    for y in range(CHUNK_HEIGHT):
      for x in range(CHUNK_WIDTH):
        absoluteTileX = topLeftTileX + x
        absoluteTileY = topLeftTileY + y
        inputLocalTileX = getInputLocalTileX(absoluteTileX)
        inputLocalTileY = getInputLocalTileY(absoluteTileY)
        inputTileIndex = inputLocalTileY * inputLayerWidthInTiles + inputLocalTileX
        outputTileIndex = y * CHUNK_WIDTH + x
        layerBaseData.append(layerBase['data'][inputTileIndex])

    if len(layerBaseData) != CHUNK_TILE_COUNT:
      print(f"Error: Layer base data length {len(layerBaseData)} does not match expected chunk tile count {CHUNK_TILE_COUNT}.")
      sys.exit(1)

    # This is a valid chunk.
    worldWidth = max(worldWidth, chunkX + 1)
    worldHeight = max(worldHeight, chunkY + 1)
    chunksDone.add((chunkX, chunkY))

    chunkHeaderPath = os.path.join(chunksDir, f"chunk_{chunkX}_{chunkY}.h")
    with open(chunkHeaderPath, 'w') as f:
      f.write(f"// Generated chunk header for chunk at position ({chunkX}, {chunkY})\n")
      f.write(f"// Generated at {now}\n")
      f.write("#pragma once\n")
      f.write("#include \"world/chunkdata.h\"\n\n")
      f.write(f"static const chunkdata_t CHUNK_{chunkX}_{chunkY} = {{\n")
      f.write(f"  .layerBase = {{\n")
      for y in range(CHUNK_HEIGHT):
        f.write(f"    ")
        for x in range(CHUNK_WIDTH):
          i = y * CHUNK_WIDTH + x
          byte = layerBaseData[i]
          f.write(f"0x{byte:02x}, ")
        f.write(f"\n")
      f.write("  },\n\n")

      f.write(f"  .layerOverlay = {{}},\n")

      f.write(f"  .entities = {{\n")
      for entity in entities:
        # Entities are center aligned in tiled.
        localX = entity['x'] - (topLeftTileX * TILE_WIDTH)
        localY = entity['y'] - (topLeftTileY * TILE_HEIGHT)

        print(f"Entity at ({entity['x']}, {entity['y']}) in chunk ({chunkX}, {chunkY}) is at local position ({localX}, {localY})")

        f.write("    {\n")
        f.write(f"      .id = {entity['id']},\n")
        f.write(f"      .type = ENTITY_TYPE_NPC,\n")
        f.write(f"      .x = {localX},\n")
        f.write(f"      .y = {localY},\n")
        f.write(f"      .dir = ENTITY_DIR_SOUTH,\n")
        f.write("    },\n")
      f.write(f"  }},\n")

      f.write("};\n\n")

# Output header file.
headerPath = os.path.join(worldDir, "world.h")
with open(headerPath, 'w') as f:
  f.write(f"// Generated chunks file. Generated at {now}\n\n")
  f.write("#pragma once\n")
  f.write("#include \"dusk.h\"\n")

  # Now, for each chunk, include its header file
  for (x, y) in chunksDone:
    chunk_header = f"world/chunk/chunk_{x}_{y}.h"
    f.write(f"#include \"{chunk_header}\"\n")
  
  f.write("\n")
  f.write(f"#define WORLD_WIDTH {worldWidth}\n")
  f.write(f"#define WORLD_HEIGHT {worldHeight}\n\n")
  f.write(f"static const chunkdata_t* WORLD_CHUNKS[] = {{\n")
  for i in range(worldHeight):
    f.write("  ")
    for j in range(worldWidth):
      if (j, i) in chunksDone:
        f.write(f"&CHUNK_{j}_{i}, ")
      else:
        f.write("NULL, ")
    f.write("\n")
  f.write("};\n\n")

print(f"chunks.h generated at: {headerPath}")