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

namespace Palettizer
{
    class Program
    {
        static int Main(string[] args)
        {
            var project = new Command("project")
            {
                new Argument<FileInfo>(
                    "project",
                    "Path to a GMS2.3 yyp file."
                ),
				new Argument<string>(
					"textureGroup",
					"The name of the texture group to palettize."
				),
				new Argument<string>(
					"paletteSpriteName",
					"The palette sprite to use in indexing."
				)
            };

            var root = new RootCommand{
                project
            };

            project.Handler = CommandHandler.Create<FileInfo, string, string, IConsole>(HandleProject);
            return root.Invoke(args);
        }

        static void HandleProject(FileInfo project, string textureGroup, string paletteSpriteName, IConsole console)
        {
			var jsonReaderOptions = new JsonDocumentOptions();
            jsonReaderOptions.AllowTrailingCommas = true;

            var gmProject = Importer.ImportProject(project);
            var projectDir = project.Directory;

            console.Out.Write("Parsing: " + gmProject.Name + "\n");

            if (Directory.Exists("output"))
            {
                Directory.Delete("output", true);
            }
            var outputDir = Directory.CreateDirectory("output");

			var palette = new Palette();
			var writer = new StbImageWriteSharp.ImageWriter();

			// Add all sprite colors to palette
			foreach (var resource in gmProject.Resources)
			{
				var path = Path.Combine(projectDir.FullName, resource.ID.Path);
                var spriteDir = Path.GetDirectoryName(path);
                var file = new FileInfo(path);

				if (file.Name.StartsWith("pal_")) { continue; }

                var resourceJsonString = File.ReadAllText(path);

				var document = JsonDocument.Parse(resourceJsonString, jsonReaderOptions);
				if (document.RootElement.GetProperty("resourceType").GetString() == "GMSprite")
				{
					var sprite = Importer.ImportSprite(file);

					if (sprite.TextureGroupID.Name == textureGroup)
					{
						console.Out.Write(sprite.Name + "\n");

						foreach (var spriteFrame in sprite.Frames)
						{
							ImageResult image;
							byte[] grayscaleImage;

							var spriteImageFilePath = Path.Combine(spriteDir, spriteFrame.CompositeImage.FrameID.Name + ".png");

							// Add sprite to palette
							using (var stream = File.OpenRead(spriteImageFilePath))
							{
								image = ImageResult.FromStream(stream, ColorComponents.RedGreenBlueAlpha);
								grayscaleImage = AddSpriteToPalette(palette, image, console);
							}

                            var spriteWriteDir = Path.Combine(outputDir.FullName, Path.GetRelativePath(projectDir.FullName, spriteDir));
                            Directory.CreateDirectory(spriteWriteDir);

                            var spriteWritePath = Path.Combine(spriteWriteDir, spriteFrame.CompositeImage.FrameID.Name + ".png");
                            // Now write the image to output directory as grayscale
                            using (var stream = File.OpenWrite(spriteWritePath))
							{
								writer.WritePng(grayscaleImage, image.Width, image.Height, StbImageWriteSharp.ColorComponents.RedGreenBlueAlpha, stream);
							}
						}
					}
				}
			}

			var paletteSpriteJSONDir = Path.Combine(projectDir.FullName, "sprites", paletteSpriteName);
			var paletteSpriteJSONPath = Path.Combine(paletteSpriteJSONDir, paletteSpriteName + ".yy");

			var paletteSpriteData = Importer.ImportSprite(new FileInfo(paletteSpriteJSONPath));
			var paletteSpriteImagePath = Path.Combine(paletteSpriteJSONDir, paletteSpriteData.Frames[0].CompositeImage.FrameID.Name + ".png");

			// Add alternate color palettes to palette
			// NOTE: our original palettes were column-indexed, we want to convert to row-index
			using (var stream = File.OpenRead(paletteSpriteImagePath))
			{
				ImageResult paletteImage = ImageResult.FromStream(stream, ColorComponents.RedGreenBlueAlpha);

				for (var j = 0; j < paletteImage.Width; j += 1)
				{
					var colorPalette = new Color[paletteImage.Height];

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

					palette.AddAlternateColorRow(colorPalette);
				}
			}

			// Write final palette sprite to PNG
			var paletteOutputPath = Path.Combine(outputDir.FullName, textureGroup + "_Palette.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, IConsole console)
		{
			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.GrayscaleIndex == 255)
					{
						console.Out.Write("Too many colors! Bailing!\n");
						return grayscaleImage;
					}
					else
					{
						palette.ColorToGrayscaleMap.Add(color, palette.GrayscaleIndex);
						palette.GrayscaleIndex += 1;
					}
				}
			}

			return grayscaleImage;
		}
    }
}