import sys, os
import argparse
from datetime import datetime
import xml.etree.ElementTree as ET
from PIL import Image

# Input font information.
CHARACTER_MAP = [
  'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
  'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
  'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
  'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
  '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '+', '-', '=',
  '(', ')', '[', ']', '{', '}', '<', '>', '/', '*', ':', '#', '%',
  '!', '?', '.', ',', "'", '"', '@', '&', '$', ' '
]
CHAR_START = 0x20  # ASCII space character
CHAR_END = 0x7E    # ASCII tilde character (exclusive)
CHARS_TOTAL = CHAR_END - CHAR_START

# 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='File to write header')
parser.add_argument('--input', required=True, help='Input XML from tiled')
args = parser.parse_args()

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

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

# Find root element
tree = ET.parse(inputFile)
root = tree.getroot()
# Check if the root element is 'tileset'
if root.tag != 'tileset':
  print(f"Error: Expected root element 'tileset', found '{root.tag}'")
  sys.exit(1)

# Shoul have tilewidth and tileheight attributes
if 'tilewidth' not in root.attrib or 'tileheight' not in root.attrib:
  print("Error: 'tileset' element must have 'tilewidth' and 'tileheight' attributes")
  sys.exit(1)

if 'tilecount' not in root.attrib:
  print("Error: 'tileset' element must have 'tilecount' attribute")
  sys.exit(1)

# Find image element
image = root.find('image')
if image is None:
  print("Error: 'tileset' element must contain an 'image' element")
  sys.exit(1)

# Ensure image has 'source' attribute
if 'source' not in image.attrib:
  print("Error: 'image' element must have a 'source' attribute")
  sys.exit(1)

# Ensure image source exists
inputDir = os.path.dirname(inputFile)

imageSource = os.path.join(inputDir, image.attrib['source'])
if not os.path.exists(imageSource):
  print(f"Error: Image source '{imageSource}' does not exist.")
  sys.exit(1)

# Ensure image has 'width' and 'height' attributes
if 'width' not in image.attrib or 'height' not in image.attrib:
  print("Error: 'image' element must have 'width' and 'height' attributes")
  sys.exit(1)

# Ensure image is readable
try:
  img = Image.open(imageSource)
except Exception as e:
  print(f"Error: Unable to open image '{imageSource}': {e}")
  sys.exit(1)

# Ensure image dimensions match the attributes
if img.width != int(image.attrib['width']) or img.height != int(image.attrib['height']):
  print(f"Error: Image dimensions ({img.width}x{img.height}) do not match attributes ({image.attrib['width']}x{image.attrib['height']})")
  sys.exit(1)

# Prepare header content
now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
tileCount = int(root.attrib['tilecount'])

outputTileIndexes = []
for i in range(CHARS_TOTAL):
  # For some reason the input tilemap is not in ASCII, so we are going to map
  # the tiles to the ASCII characters we expect, we start at 0x20 (space)

  # Find the index of the character in the CHARACTER_MAP
  inputIndex = -1
  for j, char in enumerate(CHARACTER_MAP):
    if ord(char) == (CHAR_START + i):
      inputIndex = j
      break
  
  if inputIndex == -1:
    print(f"Warning: Character '{chr(CHAR_START + i)}' not found in CHARACTER_MAP")
    outputTileIndexes.append(0)  # Use 0 for missing characters (space)
    continue

  outputTileIndexes.append(inputIndex)

with open(outputFile, 'w') as f:
  f.write(f"// Generated at {now}\n")
  f.write("#pragma once\n")
  f.write("#include \"dusk.h\"\n\n")
  f.write(f"#define FONT_TILE_WIDTH {root.attrib['tilewidth']}\n")
  f.write(f"#define FONT_TILE_HEIGHT {root.attrib['tileheight']}\n")
  f.write(f"#define FONT_TILE_COUNT {len(outputTileIndexes)}\n")
  f.write(f"#define FONT_CHAR_START {CHAR_START}\n")
  f.write(f"#define FONT_CHAR_END {CHAR_END}\n")
  f.write(f"#define FONT_CHAR_COUNT {CHARS_TOTAL}\n")
  f.write(f"#define FONT_COLUMN_COUNT {img.width // int(root.attrib['tilewidth'])}\n")
  f.write(f"#define FONT_ROW_COUNT {img.height // int(root.attrib['tileheight'])}\n\n")

  f.write("static const uint8_t TILE_INDEXES[FONT_CHAR_COUNT] = {\n")
  f.write("  ")
  for i in range(len(outputTileIndexes)):
    tileIndex = outputTileIndexes[i]
    f.write(f"{tileIndex}, ")
  f.write("\n};\n\n")
  
  f.write("static const uint8_t TILE_PIXEL_DATA[FONT_TILE_COUNT][FONT_TILE_WIDTH * FONT_TILE_HEIGHT] = {\n")
  for i in range(len(outputTileIndexes)):
    tileIndex = outputTileIndexes[i]
    f.write(f"  // Character {i} ('{chr(CHAR_START + i)}'). Read from {tileIndex} tile.\n")
    f.write(f"  {{")
    
    # Read the tile from the image
    tileX = (tileIndex % (img.width // int(root.attrib['tilewidth']))) * int(root.attrib['tilewidth'])
    tileY = (tileIndex // (img.width // int(root.attrib['tilewidth']))) * int(root.attrib['tileheight'])
    tile = img.crop((tileX, tileY, tileX + int(root.attrib['tilewidth']), tileY + int(root.attrib['tileheight'])))

    # Pixel is either 0 (transparent) or 1 (opaque)
    for y in range(int(root.attrib['tileheight'])):
      for x in range(int(root.attrib['tilewidth'])):
        pixel = tile.getpixel((x, y))
        f.write(f"0x{1 if pixel[3] > 0 else 0:02X}, ")
    f.write("},\n\n")


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