import json
from processimage import processImage
import sys
from assethelpers import getAssetRelativePath
import os
import datetime
from args import args
from xml.etree import ElementTree
from assetcache import assetGetCache, assetCache

def loadTilesetFromTSX(asset):
  # Load the TSX file
  tree = ElementTree.parse(asset['path'])
  root = tree.getroot()

  # Expect tileheight, tilewidth, columns and tilecount attributes
  if 'tilewidth' not in root.attrib or 'tileheight' not in root.attrib or 'columns' not in root.attrib or 'tilecount' not in root.attrib:
    print(f"Error: TSX file {asset['path']} is missing required attributes (tilewidth, tileheight, columns, tilecount)")
    sys.exit(1)

  tileWidth = int(root.attrib['tilewidth'])
  tileHeight = int(root.attrib['tileheight'])
  columns = int(root.attrib['columns'])
  tileCount = int(root.attrib['tilecount'])
  rows = (tileCount + columns - 1) // columns  # Calculate rows based on tileCount and columns

  # Find the image element
  imageElement = root.find('image')
  if imageElement is None or 'source' not in imageElement.attrib:
    print(f"Error: TSX file {asset['path']} is missing an image element with a source attribute")
    sys.exit(1)

  imagePath = imageElement.attrib['source']

  # Image is relative to the TSX file
  imageAssetPath = os.path.join(os.path.dirname(asset['path']), imagePath)

  image = processImage({
    'path': imageAssetPath,
    'options': asset['options'],
  })

  return {
    "image": image,
    "tileWidth": tileWidth,
    "tileHeight": tileHeight,
    "columns": columns,
    "rows": rows,
    "originalWidth": tileWidth * columns,
    "originalHeight": tileHeight * rows,
  }

def loadTilesetFromArgs(asset):
  # We need to determine how big each tile is. This can either be provided as
  # an arg of tileWidth/tileHeight or as a count of rows/columns.
  # Additionally, if the image has been factored, then the user can provide both
  # tile sizes AND cols/rows to indicate the original size of the image.
  image = processImage(asset)

  tileWidth, tileHeight = None, None
  columns, rows = None, None
  originalWidth, originalHeight = image['width'], image['height']

  if 'tileWidth' in asset['options'] and 'columns' in asset['options']:
    tileWidth = int(asset['options']['tileWidth'])
    columns = int(asset['options']['columns'])
    originalWidth = tileWidth * columns
  elif 'tileWidth' in asset['options']:
    tileWidth = int(asset['options']['tileWidth'])
    columns = image['width'] // tileWidth
  elif 'columns' in asset['options']:
    columns = int(asset['options']['columns'])
    tileWidth = image['width'] // columns
  else:
    print(f"Error: Tileset {asset['path']} must specify either tileWidth or columns")
    sys.exit(1)

  if 'tileHeight' in asset['options'] and 'rows' in asset['options']:
    tileHeight = int(asset['options']['tileHeight'])
    rows = int(asset['options']['rows'])
    originalHeight = tileHeight * rows
  elif 'tileHeight' in asset['options']:
    tileHeight = int(asset['options']['tileHeight'])
    rows = image['height'] // tileHeight
  elif 'rows' in asset['options']:
    rows = int(asset['options']['rows'])
    tileHeight = image['height'] // rows
  else:
    print(f"Error: Tileset {asset['path']} must specify either tileHeight or rows")
    sys.exit(1)
    
  return {
    "image": image,
    "tileWidth": tileWidth,
    "tileHeight": tileHeight,
    "columns": columns,
    "rows": rows,
    "originalWidth": originalWidth,
    "originalHeight": originalHeight,
  }

def processTileset(asset):
  cache = assetGetCache(asset['path'])
  if cache is not None:
    return cache

  print(f"Processing tileset: {asset['path']}")
  tilesetData = None
  if asset['path'].endswith('.tsx'):
    tilesetData = loadTilesetFromTSX(asset)
  else:
    tilesetData = loadTilesetFromArgs(asset)
  
  fileNameWithoutExtension = os.path.splitext(os.path.basename(asset['path']))[0]
  now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
  tilesetName = fileNameWithoutExtension
  tilesetNameUpper = tilesetName.upper()

  widthScale = tilesetData['originalWidth'] / tilesetData['image']['width']
  heightScale = tilesetData['originalHeight'] / tilesetData['image']['height']

  # Create header
  data = f"// Tileset Generated for {asset['path']} at {now}\n"
  data += f"#pragma once\n"
  data += f"#include \"display/tileset.h\"\n\n"
  data += f"static const tileset_t TILESET_{tilesetNameUpper} = {{\n"
  data += f"  .tileWidth = {tilesetData['tileWidth']},\n"
  data += f"  .tileHeight = {tilesetData['tileHeight']},\n"
  data += f"  .tileCount = {tilesetData['columns'] * tilesetData['rows']},\n"
  data += f"  .columns = {tilesetData['columns']},\n"
  data += f"  .rows = {tilesetData['rows']},\n"
  data += f"  .uv = {{ {widthScale / tilesetData['columns']}f, {heightScale / tilesetData['rows']}f }},\n"
  data += f"  .image = {json.dumps(tilesetData['image']['imagePath'])},\n"
  data += f"}};\n"

  # Write Header
  outputFile = os.path.join(args.headers_dir, "display", "tileset", f"tileset_{tilesetName}.h")
  os.makedirs(os.path.dirname(outputFile), exist_ok=True)
  with open(outputFile, 'w') as f:
    f.write(data)
  
  return assetCache(asset['path'], {
    "files": [],
    "image": tilesetData['image'],
    "headerFile": outputFile,
    "files": tilesetData['image']['files'],
  })