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

tilesets = []

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/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)

  print(f"Write header for tileset: {outputFile}")

  tileset = {
    "files": [],
    "image": tilesetData['image'],
    "headerFile": os.path.relpath(outputFile, args.headers_dir),
    "tilesetName": tilesetName,
    "tilesetNameUpper": tilesetNameUpper,
    "tilesetIndex": len(tilesets),
    "tilesetData": tilesetData,
    "files": tilesetData['image']['files'],
  }

  tilesets.append(tileset)
  return assetCache(asset['path'], tileset)

def processTilesetList():
  data = f"// Tileset List Generated at {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
  data += f"#pragma once\n"
  for tileset in tilesets:
    data += f"#include \"{tileset['headerFile']}\"\n"
  data += f"\n"
  data += f"#define TILESET_LIST_COUNT {len(tilesets)}\n\n"
  data += f"static const tileset_t* TILESET_LIST[TILESET_LIST_COUNT] = {{\n"
  for tileset in tilesets:
    data += f"  &TILESET_{tileset['tilesetNameUpper']},\n"
  data += f"}};\n"

  # Write header.
  outputFile = os.path.join(args.headers_dir, "display", "tileset", f"tilesetlist.h")
  os.makedirs(os.path.dirname(outputFile), exist_ok=True)
  with open(outputFile, 'w') as f:
    f.write(data)