// VkSpec.cs
// Robust vk.xml parser + semantic model for the generator.

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Xml.Linq;

public sealed class VkSpec {
    // ==========================
    // Public data model
    // ==========================

    public sealed class Handle {
        public string Name = "";              // e.g., "VkBuffer"
        public bool IsDispatchable;           // VK_DEFINE_HANDLE vs NON_DISPATCHABLE
        public string? Parent;                // optional vk.xml "parent" attr
    }

    public sealed class Param {
        public string Name = "";              // e.g., "pCreateInfo"
        public string Type = "";              // e.g., "VkBufferCreateInfo"
        public bool IsPointer;                // '*' in declaration
        public bool IsConst;                  // "const "
        public string? LenAttr;               // len="..."
        public bool IsHandle;                 // Type is a known handle
    }

    public sealed class Command {
        public string Name = "";              // e.g., "vkCreateBuffer"
        public string ReturnType = "void";    // e.g., "VkResult"
        public List<Param> Params = new();    // parameters in order
        public HashSet<string> ProtectMacros { get; } = new(StringComparer.Ordinal);  // VK_USE_PLATFORM_*
        public string? AliasOf;               // if this command is alias of another
    }

    public sealed class StructMember {
        public string Name = "";
        public string Type = "";
        public bool IsPointer;
        public bool IsConst;
        public string? LenAttr;
        public bool IsHandle;
        public bool IsStruct;
    }

    public sealed class VkStruct {
        public string Name = "";
        public List<StructMember> Members = new();
    }

    public sealed class Extension {
        public string Name = "";
        public string? Protect;                   // VK_USE_PLATFORM_* if present
        public string? Platform;                  // platform="android"/"win32"/...
        public string? Supported;                 // "vulkan", "vulkansc", "disabled", ...
        public List<string> RequireCommands = new();
    }

    public sealed class Pair {
        public string HandleType = "";            // e.g., "VkBuffer"
        public string ParentType = "";            // creation/alloc parent (may be empty for Instance)
        public string CreateOrAllocate = "";      // e.g., "vkCreateBuffer" or "vkAllocateCommandBuffers"
        public string DestroyOrFree = "";         // e.g., "vkDestroyBuffer" or "vkFreeCommandBuffers"
        public Param OutHandleParam = new();      // pBuffer / pCommandBuffers
        public List<string> ExtraOwnersForDestroy = new();
    }

    public sealed class Producer {
        public string CommandName = "";           // e.g., "vkGetDeviceQueue"
        public string ParentType = "";            // first param type
        public string HandleType = "";            // produced handle type
        public Param OutHandleParam = new();      // pQueue / pPhysicalDevices
        public bool IsMany;                       // has len=""
    }

    public sealed class OwnerUse {
        public string Type = "";                  // handle type to keep
        public bool IsMany;                       // array vs single
        public string SourceStruct = "";
        public string SourceMember = "";
    }

    public enum LoadLevel { Global, Instance, Device }

    public Dictionary<string, Handle> Handles { get; } = new(StringComparer.Ordinal);
    public Dictionary<string, Command> Commands { get; } = new(StringComparer.Ordinal);
    public Dictionary<string, VkStruct> Structs { get; } = new(StringComparer.Ordinal);
    public Dictionary<string, Extension> Extensions { get; } = new(StringComparer.Ordinal);

    public List<Pair> Pairs { get; } = new();
    public List<Producer> Producers { get; } = new();

    public Dictionary<string, string> PlatformGuards { get; } = new(StringComparer.Ordinal);

    public HashSet<string> ReachableCommands { get; } = new(StringComparer.Ordinal);

    // ==========================
    // Load
    // ==========================
    public static VkSpec Load(string vkXmlPath) {
        if (!File.Exists(vkXmlPath))
            throw new FileNotFoundException("vk.xml not found", vkXmlPath);

        var xdoc = XDocument.Load(vkXmlPath,
            LoadOptions.PreserveWhitespace | LoadOptions.SetBaseUri | LoadOptions.SetLineInfo);

        var spec = new VkSpec();
        var registry = xdoc.Root ?? throw new InvalidDataException("Malformed vk.xml: missing <registry>");

        spec.ParseTypes(registry);
        spec.ParseCommands(registry);
        spec.ParsePlatforms(registry);
        spec.ParseExtensions(registry);
        spec.ParseFeatures(registry, "vulkan");
        spec.ComputeReachableCommands("vulkan");
        spec.BuildPairsAndProducers();

        return spec;
    }

    // --------------------------
    // <types>
    // --------------------------
    void ParseTypes(XElement registry) {
        var types = registry.Element("types") ?? throw new InvalidDataException("<types> missing");

        foreach (var t in types.Elements("type")) {
            var category = (string?)t.Attribute("category") ?? "";

            if (category == "handle") {
                var name = (string?)t.Element("name") ?? (string?)t.Attribute("name") ?? "";
                if (string.IsNullOrWhiteSpace(name)) continue;

                string inner = string.Concat(t.Nodes().OfType<XText>().Select(n => n.Value));
                bool isDispatchable = inner.Contains("VK_DEFINE_HANDLE", StringComparison.Ordinal);

                Handles[name] = new Handle {
                    Name = name,
                    IsDispatchable = isDispatchable,
                    Parent = (string?)t.Attribute("parent")
                };
            } else if (category == "struct") {
                var name = (string?)t.Attribute("name") ?? "";
                if (string.IsNullOrWhiteSpace(name)) continue;

                var s = new VkStruct { Name = name };
                foreach (var mem in t.Elements("member")) {
                    var m = new StructMember {
                        Name = (string?)mem.Element("name") ?? "",
                        Type = (string?)mem.Element("type") ?? ""
                    };
                    m.IsHandle = Handles.ContainsKey(m.Type);

                    string decl = GetNodeText(mem);
                    m.IsPointer = decl.Contains("*");
                    m.IsConst = decl.Contains("const ");
                    m.LenAttr = (string?)mem.Attribute("len");
                    m.IsStruct = (!m.IsHandle && Structs.ContainsKey(m.Type));
                    s.Members.Add(m);
                }
                Structs[name] = s;
            }
        }

        // second pass mark struct-refs
        foreach (var s in Structs.Values)
            foreach (var m in s.Members)
                if (!m.IsHandle && Structs.ContainsKey(m.Type))
                    m.IsStruct = true;
    }

    // --------------------------
    // <commands>
    // --------------------------
    void ParseCommands(XElement registry) {
        var commands = registry.Element("commands") ?? throw new InvalidDataException("<commands> missing");
        var aliasToTarget = new List<(string alias, string target)>();

        foreach (var ce in commands.Elements("command")) {
            var aliasAttr = (string?)ce.Attribute("alias");
            var nameAttr = (string?)ce.Attribute("name");
            if (!string.IsNullOrEmpty(aliasAttr) && !string.IsNullOrEmpty(nameAttr)) {
                Commands[nameAttr] = new Command { Name = nameAttr, AliasOf = aliasAttr };
                aliasToTarget.Add((nameAttr, aliasAttr));
                continue;
            }

            var proto = ce.Element("proto");
            if (proto == null) continue;

            var cmdName = (string?)proto.Element("name") ?? "";
            if (string.IsNullOrWhiteSpace(cmdName)) continue;

            var retType = (string?)proto.Element("type") ?? "void";
            var c = new Command { Name = cmdName, ReturnType = retType };

            foreach (var pe in ce.Elements("param")) {
                var p = new Param {
                    Name = (string?)pe.Element("name") ?? "",
                    Type = (string?)pe.Element("type") ?? ""
                };
                p.IsHandle = Handles.ContainsKey(p.Type);

                string decl = GetNodeText(pe);
                p.IsPointer = decl.Contains("*");
                p.IsConst = decl.Contains("const ");
                p.LenAttr = (string?)pe.Attribute("len");

                c.Params.Add(p);
            }

            Commands[cmdName] = c;
        }

        // resolve aliases
        foreach (var (alias, target) in aliasToTarget) {
            if (!Commands.TryGetValue(alias, out var a) || !Commands.TryGetValue(target, out var t)) continue;
            a.ReturnType = t.ReturnType;
            a.Params = t.Params.Select(pp => new Param {
                Name = pp.Name,
                Type = pp.Type,
                IsConst = pp.IsConst,
                IsPointer = pp.IsPointer,
                LenAttr = pp.LenAttr,
                IsHandle = pp.IsHandle
            }).ToList();
            a.ProtectMacros.UnionWith(t.ProtectMacros);
        }
    }

    // --------------------------
    // <platforms>
    // --------------------------
    void ParsePlatforms(XElement registry) {
        var plats = registry.Element("platforms");
        if (plats == null) return;

        foreach (var p in plats.Elements("platform")) {
            var name = (string?)p.Attribute("name");
            var protect = (string?)p.Attribute("protect");
            if (string.IsNullOrWhiteSpace(name) || string.IsNullOrWhiteSpace(protect)) continue;
            PlatformGuards[name] = protect; // android -> VK_USE_PLATFORM_ANDROID_KHR
        }
    }

    // --------------------------
    // <extensions>
    // --------------------------
    void ParseExtensions(XElement registry) {
        var exts = registry.Element("extensions");
        if (exts == null) return;

        foreach (var xe in exts.Elements("extension")) {
            var ext = new Extension {
                Name = (string?)xe.Attribute("name") ?? "",
                Protect = (string?)xe.Attribute("protect"),
                Platform = (string?)xe.Attribute("platform"),
                Supported = (string?)xe.Attribute("supported")
            };
            if (string.IsNullOrWhiteSpace(ext.Name)) continue;

            // If protect isn't set but platform is known, resolve guard
            if (ext.Protect == null && ext.Platform != null &&
                PlatformGuards.TryGetValue(ext.Platform, out var guardMacro))
                ext.Protect = guardMacro;

            foreach (var req in xe.Elements("require")) {
                foreach (var cmd in req.Elements("command")) {
                    var cmdName = (string?)cmd.Attribute("name") ?? "";
                    if (string.IsNullOrWhiteSpace(cmdName)) continue;
                    ext.RequireCommands.Add(cmdName);

                    if (ext.Protect != null && Commands.TryGetValue(cmdName, out var c))
                        c.ProtectMacros.Add(ext.Protect);
                }
            }

            Extensions[ext.Name] = ext;
        }
    }

    // --------------------------
    // <feature api="...">
    // --------------------------
    void ParseFeatures(XElement registry, string apiFilter) {
        foreach (var feat in registry.Elements("feature")) {
            var api = (string?)feat.Attribute("api") ?? "vulkan";
            if (!api.Split(',', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries)
                    .Contains(apiFilter, StringComparer.Ordinal))
                continue;

            foreach (var req in feat.Elements("require"))
                foreach (var cmd in req.Elements("command")) {
                    var name = (string?)cmd.Attribute("name");
                    if (!string.IsNullOrWhiteSpace(name)) ReachableCommands.Add(name!);
                }
        }
    }

    // --------------------------
    // Compute enabled commands (core + enabled extensions)
    // --------------------------
    void ComputeReachableCommands(string apiFilter) {
        foreach (var ext in Extensions.Values) {
            var supported = (ext.Supported ?? "vulkan")
                .Split(',', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries);

            bool ok = supported.Contains(apiFilter, StringComparer.Ordinal)
                      && !supported.Contains("disabled", StringComparer.Ordinal);
            if (!ok) continue;

            foreach (var cmd in ext.RequireCommands)
                ReachableCommands.Add(cmd);
        }

        // Alias coherence
        foreach (var c in Commands.Values)
            if (c.AliasOf != null && ReachableCommands.Contains(c.AliasOf))
                ReachableCommands.Add(c.Name);
    }

    // --------------------------
    // Pairs + Producers
    // --------------------------
    void BuildPairsAndProducers() {
        // Create/Allocate
        foreach (var c in Commands.Values) {
            if (!ReachableCommands.Contains(c.Name)) continue;

            bool isCreate = c.Name.StartsWith("vkCreate", StringComparison.Ordinal);
            bool isAlloc = c.Name.StartsWith("vkAllocate", StringComparison.Ordinal);
            if (!isCreate && !isAlloc) continue;

            var outHandle = c.Params.FirstOrDefault(p => p.IsPointer && p.IsHandle);
            if (outHandle == null) continue;

            var handleType = outHandle.Type;
            if (!Handles.ContainsKey(handleType)) continue;

            string parentType = "";
            if (c.Params.Count > 0 && Handles.ContainsKey(c.Params[0].Type))
                parentType = c.Params[0].Type;

            var destroy = FindDestroyOrFreeFor(handleType);

            var extraOwners = new List<string>();
            if (destroy != null) {
                foreach (var p in destroy.Params) {
                    if (!p.IsHandle) continue;
                    if (p.Type == handleType) continue;
                    if (parentType != "" && p.Type == parentType) continue;
                    if (!extraOwners.Contains(p.Type)) extraOwners.Add(p.Type);
                }
            }

            Pairs.Add(new Pair {
                HandleType = handleType,
                ParentType = parentType,
                CreateOrAllocate = c.Name,
                DestroyOrFree = destroy?.Name ?? "",
                OutHandleParam = outHandle,
                ExtraOwnersForDestroy = extraOwners
            });
        }

        // Producers
        foreach (var c in Commands.Values) {
            if (!ReachableCommands.Contains(c.Name)) continue;
            bool isProducerName = c.Name.StartsWith("vkGet", StringComparison.Ordinal) ||
                                  c.Name.StartsWith("vkEnumerate", StringComparison.Ordinal);
            if (!isProducerName) continue;

            if (c.Params.Count == 0) continue;
            var parent = c.Params[0];
            if (!Handles.ContainsKey(parent.Type)) continue;

            var outHandle = c.Params.FirstOrDefault(p => p.IsPointer && Handles.ContainsKey(p.Type));
            if (outHandle == null) continue;

            Producers.Add(new Producer {
                CommandName = c.Name,
                ParentType = parent.Type,
                HandleType = outHandle.Type,
                OutHandleParam = outHandle,
                IsMany = !string.IsNullOrEmpty(outHandle.LenAttr)
            });
        }
    }

    Command? FindDestroyOrFreeFor(string handleType) {
        foreach (var c in Commands.Values) {
            if (!(c.Name.Contains("vkDestroy", StringComparison.Ordinal) ||
                  c.Name.Contains("vkFree", StringComparison.Ordinal)))
                continue;

            if (!c.Params.Any(p => p.Type == handleType)) continue;
            if (c.Params.Count == 0 || !Handles.ContainsKey(c.Params[0].Type)) continue;
            if (!ReachableCommands.Contains(c.Name)) continue;

            return c;
        }
        return null;
    }

    // --------------------------
    // Semantics for generator
    // --------------------------
    public LoadLevel ClassifyLoadLevel(Command c) {
        if (c.Params.Count == 0) return LoadLevel.Global;

        var firstType = c.Params[0].Type;
        if (!Handles.TryGetValue(firstType, out var hFirst) || !hFirst.IsDispatchable)
            return LoadLevel.Global;

        if (firstType == "VkInstance" || firstType == "VkPhysicalDevice")
            return LoadLevel.Instance;

        return LoadLevel.Device;
    }

    public IEnumerable<OwnerUse> AnalyzeCreateInfoOwners(Command createOrAlloc, string handleType) {
        var infoParam = createOrAlloc.Params.FirstOrDefault(p =>
            p.Type.EndsWith("CreateInfo", StringComparison.Ordinal) ||
            p.Type.EndsWith("AllocateInfo", StringComparison.Ordinal));
        if (infoParam == null) return Enumerable.Empty<OwnerUse>();

        if (!Structs.TryGetValue(infoParam.Type, out var infoStruct))
            return Enumerable.Empty<OwnerUse>();

        var results = new List<OwnerUse>();
        var visited = new HashSet<string>(StringComparer.Ordinal);

        void Recurse(string structName, bool arrayCtx) {
            if (!Structs.TryGetValue(structName, out var s)) return;
            if (!visited.Add(structName)) return;

            foreach (var m in s.Members) {
                if (m.IsHandle) {
                    if (m.Type != handleType) {
                        bool isMany = arrayCtx || m.IsPointer || !string.IsNullOrEmpty(m.LenAttr);
                        results.Add(new OwnerUse {
                            Type = m.Type,
                            IsMany = isMany,
                            SourceStruct = structName,
                            SourceMember = isMany ? m.Name + "[]" : m.Name
                        });
                    }
                } else if (m.IsStruct) {
                    bool childArr = arrayCtx || m.IsPointer || !string.IsNullOrEmpty(m.LenAttr);
                    Recurse(m.Type, childArr);
                }
            }

            visited.Remove(structName);
        }

        Recurse(infoStruct.Name, false);

        return results
            .GroupBy(r => (r.Type, r.IsMany, r.SourceStruct, r.SourceMember))
            .Select(g => g.First());
    }

    static string GetNodeText(XElement e) {
        return string.Concat(e.Nodes().Select(n => n switch {
            XText t => t.Value,
            XElement x => x.Value,
            _ => ""
        }));
    }
}