using System.Collections.Generic;
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[]>(
					"paletteSpriteNames",
					"A series of palette sprites 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[] paletteSpriteNames, 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");

            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");

                        for (var i = 0; i < sprite.Frames.Length; i += 1)
						{
                            var spriteFrame = sprite.Frames[i];
                            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);
							}

							// Now write out the layer
							// There's just 1 always because nobody uses GM's shit ass sprite editor
							var layerSpriteWriteDir = Path.Combine(spriteWriteDir, "layers", sprite.Frames[i].CompositeImage.FrameID.Name);
							Directory.CreateDirectory(layerSpriteWriteDir);
							var layerSpriteWritePath = Path.Combine(layerSpriteWriteDir, sprite.Layers[0].Name + ".png");
                            File.Copy(spriteWritePath, layerSpriteWritePath);
						}
                    }
				}
			}

			foreach (var paletteSpriteName in paletteSpriteNames)
            {
                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);

                    var indexToGrayscaleMap = new Dictionary<int, byte>();

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

                        if (palette.ColorToGrayscaleMap.ContainsKey(sourceColor))
                        {
                            indexToGrayscaleMap.Add(i, palette.ColorToGrayscaleMap[sourceColor]);
                        }
                    }

					// build the alternate rows
                    for (var j = 1; j < paletteImage.Width; j += 1)
                    {
                        var colorPalette = new Color[256];

                        for (var i = 0; i < paletteImage.Height; i += 1)
                        {
                            if (!indexToGrayscaleMap.ContainsKey(i)) { continue; }
                            var grayscale = indexToGrayscaleMap[i];

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

                        palette.AddAlternateColorRow(colorPalette);
                    }
                }

				// Write final palette sprite to PNG
				var paletteOutputPath = Path.Combine(outputDir.FullName, textureGroup + "_" + paletteSpriteName + "_Palette.png");
				using (var stream = File.OpenWrite(paletteOutputPath))
				{
					writer.WritePng(palette.CreateIndexedPaletteBitmap(), palette.Width, palette.Height, StbImageWriteSharp.ColorComponents.RedGreenBlueAlpha, stream);
				}

                palette.ClearAlternateColorRows();
            }
        }

		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.GrayscaleCount == 256)
					{
						console.Out.Write("Too many colors! Bailing!\n");
						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;
		}
    }
}