using System;
using System.IO;
using System.Reflection;
using System.Text.RegularExpressions;
using SharpDX.D3DCompiler;

namespace Blacklight.SharpDX9.PerformanceTests
{
    /// <summary>
    /// Compiles the main pixel shader the way the backend does and reports its instruction and register counts, the
    /// numbers a GPU's occupancy and the ps_3_0 slot limit depend on.
    /// </summary>
    public static class ShaderStats
    {
        public const string Argument = "--shader-stats";

        /// <summary>
        /// Gets the text of the main shader with its includes expanded.
        /// </summary>
        public static string PreprocessedMainShader()
        {
            Assembly assembly = typeof(Module).Assembly;
            string name = Module.ShaderNamespace + ".Main.fx";
            using (Stream stream = assembly.GetManifestResourceStream(name))
            using (var reader = new StreamReader(stream))
            {
                return ShaderBytecode.Preprocess(reader.ReadToEnd(), null, new Module.IncludeHandler(), out string _);
            }
        }

        /// <summary>
        /// Gets the disassembly of one entry point of the main shader.
        /// </summary>
        public static string Disassemble(string source, string entryPoint, string profile)
        {
            // D3DX9's compiler, the one production D3D9 shaders go through.
            using (SharpDX.Direct3D9.ShaderBytecode compiled = SharpDX.Direct3D9.ShaderBytecode.Compile(
                source, entryPoint, profile, SharpDX.Direct3D9.ShaderFlags.OptimizationLevel3))
            {
                byte[] bytes = compiled.Data.ReadRange<byte>((int)compiled.Data.Length);
                return new ShaderBytecode(bytes).Disassemble(DisassemblyFlags.None);
            }
        }

        /// <summary>
        /// Gets the highest temporary register a disassembly uses, plus one.
        /// </summary>
        public static int TemporaryRegisterCount(string disassembly)
        {
            int highest = -1;
            foreach (Match match in Regex.Matches(disassembly, @"\br(\d+)\b"))
            {
                highest = Math.Max(highest, int.Parse(match.Groups[1].Value));
            }

            return highest + 1;
        }

        /// <summary>
        /// Writes the full disassembly of each entry point to a file in `directory`.
        /// </summary>
        public static int Dump(string directory)
        {
            string source = PreprocessedMainShader();
            foreach ((string entryPoint, string profile) in new[] { ("MainPS", "ps_3_0"), ("MainVS", "vs_3_0") })
            {
                string path = Path.Combine(directory, entryPoint + ".asm");
                File.WriteAllText(path, Disassemble(source, entryPoint, profile));
            }

            return 0;
        }

        public static int Run()
        {
            string source = PreprocessedMainShader();
            foreach ((string entryPoint, string profile) in new[] { ("MainPS", "ps_3_0"), ("MainVS", "vs_3_0") })
            {
                string disassembly = Disassemble(source, entryPoint, profile);
                Match slots = Regex.Match(disassembly, @"approximately (\d+) instruction slots used[^\r\n]*");
                Console.WriteLine($"{entryPoint}: {slots.Value}");
                Console.WriteLine($"{entryPoint}: {TemporaryRegisterCount(disassembly)} temporary registers");
            }

            return 0;
        }
    }
}