import os
from PIL import Image
from args import args
import sys
import datetime

palettes = []

def extractPaletteFromImage(image):
  # goes through and finds all unique colors in the image
  if image.mode != 'RGBA':
    image = image.convert('RGBA')
  pixels = list(image.getdata())
  uniqueColors = []
  for color in pixels:
    # We treat all alpha 0 as rgba(0,0,0,0) for palette purposes
    if color[3] == 0:
      color = (0, 0, 0, 0)
    if color not in uniqueColors:
      uniqueColors.append(color)
  return uniqueColors

def processPalette(asset):
  print(f"Processing palette: {asset['path']}")

  paletteIndex = len(palettes)
  image = Image.open(asset['path'])
  pixels = extractPaletteFromImage(image)

  fileNameWithoutExt = os.path.splitext(os.path.basename(asset['path']))[0]
  fileNameWithoutPalette = os.path.splitext(fileNameWithoutExt)[0]

  # PSP requires that the palette size be a power of two, so we will pad the
  # palette with transparent colors if needed.
  # if args.platform == "psp":
  def mathNextPowTwo(x):
    return 1 << (x - 1).bit_length()
  
  nextPowTwo = mathNextPowTwo(len(pixels))
  while len(pixels) < nextPowTwo:
    pixels.append((0, 0, 0, 0))

  # Header
  now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
  data = f"// Palette Generated for {asset['path']} at {now}\n"
  data += f"#include \"display/palette/palette.h\"\n\n"
  data += f"#define PALETTE_{paletteIndex}_COLOR_COUNT {len(pixels)}\n\n"
  data += f"#pragma pack(push, 1)\n"
  data += f"static const color_t PALETTE_{paletteIndex}_COLORS[PALETTE_{paletteIndex}_COLOR_COUNT] = {{\n"
  for pixel in pixels:
    data += f"  {{ 0x{pixel[0]:02X}, 0x{pixel[1]:02X}, 0x{pixel[2]:02X}, 0x{pixel[3]:02X} }},\n"
  data += f"}};\n\n"
  data += f"#pragma pack(pop)\n\n"
  data += f"static const palette_t PALETTE_{paletteIndex} = {{\n"
  data += f"  .colorCount = PALETTE_{paletteIndex}_COLOR_COUNT,\n"
  data += f"  .colors = PALETTE_{paletteIndex}_COLORS,\n"
  data += f"}};\n"

  # Write Header
  outputFile = os.path.join(args.headers_dir, "display", "palette", f"palette_{paletteIndex}.h")
  os.makedirs(os.path.dirname(outputFile), exist_ok=True)
  with open(outputFile, "w") as f:
    f.write(data)

  palette = {
    "paletteIndex": paletteIndex,
    "paletteName": fileNameWithoutPalette,
    "pixels": pixels,
    "headerFile": os.path.relpath(outputFile, args.headers_dir),
    "asset": asset,
    "files": [ ],# No zippable files.
  }

  palettes.append(palette)
  return palette

def processPaletteList():
  data = f"// Auto-generated palette list\n"
  for palette in palettes:
    data += f"#include \"{palette['headerFile']}\"\n"
  data += f"\n"
  data += f"#define PALETTE_LIST_COUNT {len(palettes)}\n\n"
  data += f"static const palette_t* PALETTE_LIST[PALETTE_LIST_COUNT] = {{\n"
  for palette in palettes:
    data += f"  &PALETTE_{palette['paletteIndex']},\n"
  data += f"}};\n"

  # Write the palette list to a header file
  outputFile = os.path.join(args.headers_dir, "display", "palette", "palettelist.h")
  os.makedirs(os.path.dirname(outputFile), exist_ok=True)
  with open(outputFile, "w") as f:
    f.write(data)