using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Drawing;
using System.Drawing.Imaging;
using System.IO;
using System.Runtime.InteropServices;

namespace RMAPack
{
    public struct ImageDesc
    {
        public string Roughness;
        public string Metallicty;
        public string BaseColor;
        public string AmbientOcclusion;
        public string Normal;
    }

    class Program
    {
        private static readonly string[] BASE_COLOR = new string[] { "_base_color", "_basecolor" };
        private static readonly string[] NORMAL = new string[] { "_normal" };
        private static readonly string[] ROUGHNESS = new string[] { "_roughness", "_rough", "_r"};
        private static readonly string[] METALLIC = new string[] { "_metallic", "_metallicity", "_m"};
        private static readonly string[] AMBIENT_OCCLUSION = new string[] { "_ambientocclusion" };

        static void Main(string[] args)
        {
            if (args.Length < 2)
            {
                Usage();
                return;
            }
            var sourceDir = args[0];
            var targetDir = args[1];
            ProcessFolder(sourceDir, targetDir);
        }

        static void ProcessFolder(string sourceDir, string targetDir)
        {
            var sourceFiles = Directory.EnumerateFiles(sourceDir);
            var desces = CollectDescriptions(sourceFiles);
            foreach (var desc in desces)
            {
                if (string.IsNullOrEmpty(desc.Value.Roughness) || string.IsNullOrEmpty(desc.Value.Metallicty)
                    || string.IsNullOrEmpty(desc.Value.Normal) || string.IsNullOrEmpty(desc.Value.BaseColor))
                {
                    Console.WriteLine($@"Skipping {desc.Key} - not enough maps");
                    continue;
                }
                using (var normal = Bitmap.FromFile(desc.Value.Normal))
                using (var baseColor = Bitmap.FromFile(desc.Value.BaseColor))
                using (var mask = ProcessImage(desc.Key, desc.Value))
                {
                    var maskOutPath = Path.Combine(targetDir, $"T_{desc.Key}_M.png");
                    var baseOutPath = Path.Combine(targetDir, $"T_{desc.Key}_BC.png");
                    var normalOutPath = Path.Combine(targetDir, $"T_{desc.Key}_N.png");
                    mask.Save(maskOutPath, ImageFormat.Png);
                    normal.Save(normalOutPath, ImageFormat.Png);
                    baseColor.Save(baseOutPath, ImageFormat.Png);
                }
            }
        }

        static Bitmap ProcessImage(string imageName, ImageDesc desc)
        {
            Console.WriteLine($@"Processing {imageName}");
            using (var roughness = Bitmap.FromFile(desc.Roughness))
            using (var metallicity = Bitmap.FromFile(desc.Metallicty))
            {
                Console.WriteLine(" Loading source images");
                var ao = string.IsNullOrEmpty(desc.AmbientOcclusion) ? null : Bitmap.FromFile(desc.AmbientOcclusion);

                var roughnessData = ExtractRedChannel(roughness as Bitmap);
                var metallicityData = ExtractRedChannel(metallicity as Bitmap);
                var aoData = ao != null ? ExtractRedChannel(ao as Bitmap) : null;
                var resultImage = new Bitmap(roughness);
                BitmapData data = null;
                try
                {
                    data = resultImage.LockBits(new Rectangle(0, 0, resultImage.Width, resultImage.Height),
                        ImageLockMode.WriteOnly, PixelFormat.Format24bppRgb);
                    var resultData = new byte[data.Stride * resultImage.Width];
                    Console.WriteLine(" Composing final image");
                    for (var row = 0; row < resultImage.Height; row++)
                    {
                        for (var offset = 0; offset < resultImage.Width; offset++)
                        {
                            var inShift = row * resultImage.Width + offset;
                            var outShift = (row * data.Stride + offset * 3);
                            resultData[outShift + 2] = roughnessData[inShift];
                            resultData[outShift + 1] = metallicityData[inShift];
                            resultData[outShift] = aoData != null ? aoData[inShift] : (byte)255;

                        }
                    }
                    Marshal.Copy(resultData, 0, data.Scan0, resultData.Length);
                    return resultImage;
                }
                finally
                {
                    if (data != null)
                    {
                        resultImage.UnlockBits(data);
                    }
                    if (ao != null)
                    {
                        ao.Dispose();
                    }
                }
            }
        }

        static byte[] ExtractRedChannel(Bitmap bitmap)
        {
            BitmapData data = null;
            try
            {
                data = bitmap.LockBits(new Rectangle(0, 0, bitmap.Width, bitmap.Height), ImageLockMode.ReadOnly,
                    PixelFormat.Format24bppRgb);
                var result = new byte[bitmap.Width * bitmap.Height];
                var sourceData = new byte[data.Stride * bitmap.Height];
                Marshal.Copy(data.Scan0, sourceData, 0, sourceData.Length);
                for (var row = 0; row < bitmap.Height; row++)
                {
                    for (var shift = 0; shift < bitmap.Width; shift++)
                    {
                        var red = sourceData[(row * data.Stride + shift * 3)];
                        result[row * bitmap.Width + shift] = red;
                    }
                }
                return result;
            }
            finally
            {
                if (data != null)
                {
                    bitmap.UnlockBits(data);
                }

            }

        }

        private static Dictionary<string, ImageDesc> CollectDescriptions(IEnumerable<string> sourceFiles)
        {
            var allImageDesc = new Dictionary<string, ImageDesc>();
            foreach (var sourceFile in sourceFiles)
            {
                var fullName = Path.GetFileNameWithoutExtension(sourceFile.ToLower());
                var lastUnderscorePosition = fullName.LastIndexOf("_");
                if (fullName.EndsWith(BASE_COLOR[0]))
                {
                    lastUnderscorePosition = fullName.LastIndexOf(BASE_COLOR[0]); // HACK!
                }
                if (lastUnderscorePosition < 0)
                {
                    continue;
                }
                var imageName = fullName.Substring(0, lastUnderscorePosition);
                ImageDesc desc;
                if (!allImageDesc.TryGetValue(imageName, out desc))
                {
                    desc = new ImageDesc();
                }
                if (IsNameEndsWith(fullName, ROUGHNESS))
                {
                    desc.Roughness = sourceFile;
                }
                else if (IsNameEndsWith(fullName, METALLIC))
                {
                    desc.Metallicty = sourceFile;
                }
                else if (IsNameEndsWith(fullName, BASE_COLOR))
                {
                    desc.BaseColor = sourceFile;
                }
                else if (IsNameEndsWith(fullName, NORMAL))
                {
                    desc.Normal = sourceFile;
                }
                else if (IsNameEndsWith(fullName, AMBIENT_OCCLUSION))
                {
                    desc.AmbientOcclusion = sourceFile;
                }
                allImageDesc[imageName] = desc;
            }
            return allImageDesc;
        }

        static bool IsNameEndsWith(string name, string[] ends)
        {
            foreach (var end in ends)
            {
                if (name.EndsWith(end))
                {
                    return true;
                }
            }
            return false;
        }

        static void Usage()
        {
            Console.WriteLine("Usage: rmapack sourcedir targetdir");
        }
    }
}