using System.CommandLine;
using System.CommandLine.Invocation;
using System.IO;
using StbImageSharp;

namespace Palettizer2
{
	class Program
	{
		static int Main(string[] args)
		{
			var palettize = new Command("palettize")
			{
				new Argument<DirectoryInfo>(
					"directory",
					"Path to a directory containing PNG files."
				),
				new Argument<string>(
					"paletteName",
					"The name of the palette sprite to output."
				)
			};

			var root = new RootCommand
			{
				palettize
			};

			palettize.Handler = CommandHandler.Create<DirectoryInfo, string>(HandlePalettize);
			return root.Invoke(args);
		}

		static void HandlePalettize(DirectoryInfo directory, string paletteName)
		{
			var outputDir = Directory.CreateDirectory("output");
			var palette = new Palette();
			var writer = new StbImageWriteSharp.ImageWriter();

			// add all sprite colors to palette
			foreach (var imageFile in directory.EnumerateFiles("*.png", SearchOption.AllDirectories))
			{
				System.Console.WriteLine(imageFile.FullName);

				var spriteDir = Path.GetDirectoryName(imageFile.FullName);

				ImageResult image;
				byte[] grayscaleImageBytes;
				using (var stream = File.OpenRead(imageFile.FullName))
				{
					image = ImageResult.FromStream(stream, ColorComponents.RedGreenBlueAlpha);
					grayscaleImageBytes = AddSpriteToPalette(palette, image);
				}

				var spriteWriteDir = Path.Combine(
					outputDir.FullName,
					Path.GetRelativePath(directory.FullName, spriteDir));

				Directory.CreateDirectory(spriteWriteDir);
				var spriteWritePath = Path.Combine(spriteWriteDir, imageFile.Name);

				using (var stream = File.OpenWrite(spriteWritePath))
				{
					writer.WritePng(
						grayscaleImageBytes,
						image.Width,
						image.Height,
						StbImageWriteSharp.ColorComponents.RedGreenBlueAlpha,
						stream);
				}
			}

			// write palette sprite to PNG
			var paletteOutputPath = Path.Combine(outputDir.FullName, paletteName + ".png");
			using (var stream = File.OpenWrite(paletteOutputPath))
			{
				writer.WritePng(
					palette.CreateIndexedPaletteBitmap(),
					Palette.Width,
					palette.Height,
					StbImageWriteSharp.ColorComponents.RedGreenBlueAlpha,
					stream);
			}
		}

		static byte[] AddSpriteToPalette(Palette palette, ImageResult image)
		{
			var grayscaleImage = new byte[image.Width * image.Height * 4];

			for (var i = 0; i < image.Width * image.Height * 4; i += 4)
			{
				var color = new Color
				{
					R = image.Data[i],
					G = image.Data[i + 1],
					B = image.Data[i + 2],
					A = image.Data[i + 3]
				};

				if (color.A == 0)
				{
					continue;
				}

				if (palette.ColorToGrayscaleMap.ContainsKey(color))
				{
					var grayscaleColor = palette.ColorToGrayscaleMap[color];

					grayscaleImage[i] = grayscaleColor;
					grayscaleImage[i + 1] = grayscaleColor;
					grayscaleImage[i + 2] = grayscaleColor;
					grayscaleImage[i + 3] = color.A;
				}
				else
				{
					if (palette.GrayscaleCount == 256)
					{
						System.Console.WriteLine("Too many colors! Bailing!");
						return grayscaleImage;
					}
					else
					{
						palette.AddColor(color);

						var grayscaleColor = palette.ColorToGrayscaleMap[color];

						grayscaleImage[i] = grayscaleColor;
						grayscaleImage[i + 1] = grayscaleColor;
						grayscaleImage[i + 2] = grayscaleColor;
						grayscaleImage[i + 3] = color.A;
					}
				}
			}

			return grayscaleImage;
		}
	}
}