using System.Text;
using System.Text.RegularExpressions;

namespace AcDream.Tools.ShaderCompiler;

/// <summary>
/// Campaign V slice V6c: give every stage-to-stage varying an explicit
/// <c>layout(location = N)</c>, which SPIR-V requires and desktop GL does not.
///
/// <para>This is the one place the compiler edits a shader body rather than
/// prepending to it, and it is limited to inserting a qualifier in front of a
/// declaration it already matched exactly. Locations are keyed BY NAME, taken
/// from the vertex stage's outputs and looked up by the fragment stage's inputs,
/// so the two stages cannot drift apart if someone reorders a declaration in one
/// of them. Assigning by ordinal would look identical today and produce silently
/// swapped varyings the first time an author moved a line.</para>
///
/// <para>Anything already carrying a <c>layout(</c> qualifier is left alone, as
/// are interface blocks (which open a brace) and the redeclared
/// <c>gl_PerVertex</c> block that lets a core-profile shader write
/// <c>gl_ClipDistance</c>.</para>
/// </summary>
internal static partial class GlslVaryingLocations
{
    // GLSL allows the interpolation qualifier on either side of the direction —
    // both `out flat uvec2 v;` and `flat out uvec2 v;` appear in acdream's
    // shaders — so the pattern accepts either and neither.
    [GeneratedRegex(
        @"^(?<indent>\s*)(?<leading>(flat|noperspective|smooth|centroid)\s+)*(?<direction>in|out)\s+(?<interp>(flat|noperspective|smooth|centroid)\s+)*(?<type>[A-Za-z_][A-Za-z0-9_]*)\s+(?<name>[A-Za-z_][A-Za-z0-9_]*)\s*(?<array>\[[^\]]*\])?\s*;\s*(?<trailing>//.*)?$",
        RegexOptions.ExplicitCapture)]
    private static partial Regex VaryingDeclaration();

    /// <summary>
    /// Rewrites <paramref name="source"/> so every user varying carries a
    /// location. <paramref name="locationsByName"/> is shared across the two
    /// stages of a pair and is populated by whichever stage declares a name
    /// first.
    /// </summary>
    internal static string Apply(string source, string stage, Dictionary<string, int> locationsByName)
    {
        ArgumentNullException.ThrowIfNull(source);
        ArgumentNullException.ThrowIfNull(locationsByName);

        string[] lines = source.Replace("\r\n", "\n").Split('\n');
        var output = new StringBuilder();
        int nextFragmentOutput = 0;

        foreach (string line in lines)
        {
            if (line.Contains("layout(", StringComparison.Ordinal) || line.Contains('{'))
            {
                output.AppendLine(line);
                continue;
            }

            Match match = VaryingDeclaration().Match(line);
            if (!match.Success)
            {
                output.AppendLine(line);
                continue;
            }

            string direction = match.Groups["direction"].Value;
            string name = match.Groups["name"].Value;

            // Vertex inputs are attributes, not varyings: their locations are
            // the vertex layout the pipeline declares, and every acdream vertex
            // shader already states them. Leaving an unqualified one alone makes
            // the compiler say so rather than inventing a binding.
            if (stage == "vert" && direction == "in")
            {
                output.AppendLine(line);
                continue;
            }

            int location;
            if (stage == "frag" && direction == "out")
            {
                // Fragment outputs live in their own location space; acdream has
                // exactly one colour attachment everywhere.
                location = nextFragmentOutput++;
            }
            else if (locationsByName.TryGetValue(name, out int existing))
            {
                location = existing;
            }
            else
            {
                location = locationsByName.Count == 0 ? 0 : locationsByName.Values.Max() + 1;
                locationsByName[name] = location;
            }

            string indent = match.Groups["indent"].Value;
            string rest = line[indent.Length..];
            output.AppendLine($"{indent}layout(location = {location}) {rest}");
        }

        return output.ToString();
    }
}