import os
from args import args
from PIL import Image
import struct
import sys
from assethelpers import getAssetRelativePath

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 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 savePalette(pixels, outputFilePath):
  colorCount = len(pixels)
  buf = bytearray(b"DPF")
  buf += struct.pack("<i", colorCount)           # little-endian int32_t
  for r, g, b, a in pixels:                      # each channel 0..255
    buf += struct.pack("BBBB", r, g, b, a)     # raw bytes
  os.makedirs(os.path.dirname(outputFilePath), exist_ok=True)
  with open(outputFilePath, "wb") as f:
    f.write(buf)

def processPalette(asset):
  assetPath = asset['path']

  # Process the image file
  print(f"Processing palette: {assetPath}")

  # Load the image
  image = Image.open(assetPath)
  pixels = extractPaletteFromImage(image)

  # Save the processed image to the output directory
  fileNameWithoutExt = os.path.splitext(os.path.basename(assetPath))[0]
  outputFilePath = os.path.join(args.output_assets, f"{fileNameWithoutExt}.dpf")
  os.makedirs(os.path.dirname(outputFilePath), exist_ok=True)
  savePalette(pixels, outputFilePath)

  outputRelative = os.path.relpath(outputFilePath, args.output_assets)

  palette = {
    "outputFile": outputRelative,
    "paletteColors": len(pixels),
    "files": [ outputFilePath ],
    "pixels": pixels
  }

  PALETTES.append(palette)

  return palette

def processImage(asset):
  assetPath = asset['path']

  print(f"Processing image: {assetPath}")
  
  # Load the image
  image = Image.open(assetPath)

  # Get the image's palette because we are going to try and find a matching
  # palette from the already processed palettes...
  imagePalette = extractPaletteFromImage(image)
  
  # Now find which palette has every single color in this image's palette
  paletteIndex = -1
  for i, palette in enumerate(PALETTES):
    paletteColors = palette["pixels"]
    if all(color in paletteColors for color in imagePalette):
      paletteIndex = i
      break

  # Did we manage to find a matching palette?
  if paletteIndex == -1:
    print(f"Error: No matching palette found for image {assetPath}. Please process a suitable palette first.")
    sys.exit(1)
    
  # We found the palette, so now we can convert the image to use that palette
  palette = PALETTES[paletteIndex]
  indexes = []
  for color in imagePalette:
    if color in palette["pixels"]:
      index = palette["pixels"].index(color)
      indexes.append(index)
    else:
      print(f"Error: Color {color} in image {assetPath} not found in palette.")
      sys.exit(1)

  relative = getAssetRelativePath(assetPath)
  fileNameWithoutExt = os.path.splitext(os.path.basename(assetPath))[0]
  outputFileRelative = os.path.join(os.path.dirname(relative), f"{fileNameWithoutExt}.dpi")
  outputFilePath = os.path.join(args.output_assets, outputFileRelative)
  os.makedirs(os.path.dirname(outputFilePath), exist_ok=True)
  with open(outputFilePath, "wb") as f:
    # Write header
    f.write(b"DPI")  # Dusk Palettized Image
    f.write(struct.pack("<i", image.width))   # Width
    f.write(struct.pack("<i", image.height))  # Height
    f.write(struct.pack("B", paletteIndex))   # Palette index
    
    # Write uint8_t pixel index.
    for index in indexes:
      f.write(struct.pack("B", index))

  return {
    "paletteIndex": paletteIndex,
    "palette": PALETTES[paletteIndex],
    "outputFile": outputFileRelative,
    "files": [ assetPath, outputFilePath ]
  }