using System;
using System.Collections.Generic;
using System.Linq;
using System.Text.RegularExpressions;
using System.Xml.Linq;

public sealed class VkSpec {
    // ---------- Data model ----------

    public sealed record Handle(string Name, bool Dispatchable);
    public sealed record Param(string Raw, string Type, string Name, bool IsPointer, bool IsConst, string? LenAttr, bool IsHandle);
    public sealed record Command(string Name, string ReturnType, List<Param> Params);

    public sealed record StructMember(string Type, string Name, bool IsPointer, string? LenAttr, bool IsHandle);
    public sealed record StructDef(string Name, List<StructMember> Members);

    // Owned pair (created/destroyed or allocated/freed)
    public sealed record Pair(
        string HandleType,
        string CreateOrAllocate, string DestroyOrFree,
        string ParentType,
        Param OutHandleParam,
        List<string> ExtraOwnersForDestroy
    );

    // Simple producer (Get*/Enumerate* with single parent handle and a single out handle param)
    public sealed record SimpleProducer(
        string HandleType,
        string CommandName,
        string ParentType,
        Param OutHandleParam,
        bool IsMany
    );

    public sealed record OwnerUse(string Type, bool IsMany, string SourceStruct, string SourceMember);

    public Dictionary<string, Handle> Handles { get; private init; } = new();
    public Dictionary<string, Command> Commands { get; private init; } = new();
    public Dictionary<string, StructDef> Structs { get; private init; } = new();
    public HashSet<string> Extensions { get; private init; } = new();
    public List<Pair> Pairs { get; private init; } = new();
    public List<SimpleProducer> Producers { get; private init; } = new();

    // ---------- Load & parse ----------

    public static VkSpec Load(string xmlPath) {
        var doc = XDocument.Load(xmlPath);

        var handles = ParseHandles(doc);
        var structs = ParseStructs(doc, handles);
        var (real, aliasElems) = ParseCommandsPre(doc);
        var cmds = BuildCommands(real, handles);
        ApplyAliases(cmds, aliasElems, handles);
        var exts = ParseExtensions(doc);
        var pairs = BuildPairs(cmds, handles);
        var producers = BuildSimpleProducers(cmds, handles);

        return new VkSpec {
            Handles = handles,
            Structs = structs,
            Commands = cmds,
            Extensions = exts,
            Pairs = pairs,
            Producers = producers
        };
    }

    static Dictionary<string, Handle> ParseHandles(XDocument doc) {
        var dict = new Dictionary<string, Handle>();
        var types = doc.Root!.Element("types")!.Elements("type")
            .Where(t => (string?)t.Attribute("category") == "handle");

        foreach (var t in types) {
            var name = t.Element("name")?.Value;
            if (string.IsNullOrEmpty(name)) {
                var txt = (t.Value ?? "");
                var m = Regex.Match(txt, @"VK_DEFINE_(?:NON_)?DISPATCHABLE_HANDLE\((Vk\w+)\)");
                if (m.Success) name = m.Groups[1].Value;
            }
            if (string.IsNullOrEmpty(name)) continue;

            var body = t.Value ?? string.Empty;
            bool dispatchable =
                body.Contains("VK_DEFINE_HANDLE(") ||
                body.Contains("VK_DEFINE_DISPATCHABLE_HANDLE(");
            dict[name!] = new Handle(name!, dispatchable);
        }
        return dict;
    }

    static Dictionary<string, StructDef> ParseStructs(XDocument doc, Dictionary<string, Handle> handles) {
        var dict = new Dictionary<string, StructDef>();
        var xs = doc.Root!.Element("types")!.Elements("type")
            .Where(t => (string?)t.Attribute("category") == "struct");

        foreach (var t in xs) {
            var name = t.Element("name")?.Value ?? (string?)t.Attribute("name");
            if (string.IsNullOrEmpty(name)) continue;

            var members = new List<StructMember>();
            foreach (var m in t.Elements("member")) {
                var type = m.Element("type")?.Value ?? "void";
                var mname = m.Element("name")?.Value ?? "";
                var raw = Regex.Replace(m.Value ?? "", @"\s+", " ");
                bool isPtr = raw.Contains("*");
                string? len = (string?)m.Attribute("len");
                bool isHandle = handles.ContainsKey(type);
                members.Add(new StructMember(type, mname, isPtr, len, isHandle));
            }
            dict[name!] = new StructDef(name!, members);
        }

        return dict;
    }

    static (List<XElement> real, List<XElement> aliases) ParseCommandsPre(XDocument doc) {
        var all = doc.Root!.Element("commands")!.Elements("command").ToList();
        var real = new List<XElement>();
        var aliases = new List<XElement>();
        foreach (var c in all) {
            if (c.Attribute("alias") != null) aliases.Add(c);
            else if (c.Element("proto") != null) real.Add(c);
        }
        return (real, aliases);
    }

    static string Cleanup(string s) => Regex.Replace(s.Trim(), @"\s+", " ");

    static Dictionary<string, Command> BuildCommands(List<XElement> realCmds, Dictionary<string, Handle> handles) {
        bool IsPointer(string raw) => raw.Contains("*");
        bool IsConst(string raw) => raw.Contains("const ");

        var dict = new Dictionary<string, Command>();

        foreach (var c in realCmds) {
            var proto = c.Element("proto")!;
            var name = proto.Element("name")!.Value;

            string retType = string.Concat(proto.Nodes().Select(n => {
                if (n is XElement e) return e.Name == "name" ? "" : e.Value;
                return ((XText)n).Value;
            }));
            retType = Cleanup(retType.Replace(name, ""));

            var prms = new List<Param>();
            foreach (var p in c.Elements("param")) {
                var type = p.Element("type")?.Value ?? "void";
                var pname = p.Element("name")?.Value ?? "";
                var raw = Cleanup(p.Value);
                var len = (string?)p.Attribute("len");
                bool isPtr = IsPointer(raw);
                bool isConst = IsConst(raw);
                bool isHandle = handles.ContainsKey(type);
                prms.Add(new Param(raw, type, pname, isPtr, isConst, len, isHandle));
            }

            dict[name] = new Command(name, retType, prms);
        }

        return dict;
    }

    static void ApplyAliases(Dictionary<string, Command> dict, List<XElement> aliases, Dictionary<string, Handle> handles) {
        foreach (var c in aliases) {
            var name = (string?)c.Attribute("name");
            var aliasOf = (string?)c.Attribute("alias");
            if (string.IsNullOrWhiteSpace(name) || string.IsNullOrWhiteSpace(aliasOf)) continue;

            if (dict.TryGetValue(aliasOf!, out var baseCmd)) {
                var clone = new Command(name!, baseCmd.ReturnType,
                    baseCmd.Params.Select(p => new Param(p.Raw, p.Type, p.Name, p.IsPointer, p.IsConst, p.LenAttr, p.IsHandle)).ToList());
                dict[name!] = clone;
            }
        }
    }

    static HashSet<string> ParseExtensions(XDocument doc) {
        var set = new HashSet<string>(StringComparer.Ordinal);
        var exts = doc.Root!.Element("extensions")!.Elements("extension");
        foreach (var e in exts) {
            var name = (string?)e.Attribute("name");
            if (!string.IsNullOrWhiteSpace(name)) set.Add(name!);
        }
        return set;
    }

    static List<Pair> BuildPairs(Dictionary<string, Command> cmds, Dictionary<string, Handle> handles) {
        var list = new List<Pair>();

        // vkCreateX / vkDestroyX
        foreach (var create in cmds.Values.Where(x => x.Name.StartsWith("vkCreate"))) {
            var obj = create.Name.Substring("vkCreate".Length);
            var destroyName = "vkDestroy" + obj;
            if (!cmds.TryGetValue(destroyName, out var destroy)) continue;

            var parent = create.Params.First().Type;
            var outParam = create.Params.LastOrDefault(p => p.IsPointer && handles.ContainsKey(p.Type));
            if (outParam is null) continue;

            var handleType = outParam.Type;
            var extraOwners = destroy.Params
                .Where(p => p.IsHandle && p.Type != handleType)
                .Select(p => p.Type).Distinct().ToList();

            list.Add(new Pair(handleType, create.Name, destroyName, parent, outParam, extraOwners));
        }

        // vkAllocateX / vkFreeX
        foreach (var alloc in cmds.Values.Where(x => x.Name.StartsWith("vkAllocate"))) {
            var obj = alloc.Name.Substring("vkAllocate".Length);
            var freeName = "vkFree" + obj;
            if (!cmds.TryGetValue(freeName, out var free)) continue;

            var parent = alloc.Params.First().Type;
            var outParam = alloc.Params.LastOrDefault(p => p.IsPointer && handles.ContainsKey(p.Type));
            if (outParam is null) continue;

            var handleType = outParam.Type;
            var extraOwners = free.Params
                .Where(p => p.IsHandle && p.Type != handleType)
                .Select(p => p.Type).Distinct().ToList();

            list.Add(new Pair(handleType, alloc.Name, freeName, parent, outParam, extraOwners));
        }

        return list
            .GroupBy(p => p.HandleType)
            .Select(g => g.First())
            .ToList();
    }

    // Build "simple producers": vkGet*/vkEnumerate* that output a handle with exactly one parent handle argument
    static List<SimpleProducer> BuildSimpleProducers(Dictionary<string, Command> cmds, Dictionary<string, Handle> handles) {
        var list = new List<SimpleProducer>();

        foreach (var c in cmds.Values) {
            if (!(c.Name.StartsWith("vkGet") || c.Name.StartsWith("vkEnumerate")))
                continue;

            // pick first handle output param (pointer to handle)
            int outIdx = c.Params.FindIndex(p => p.IsPointer && handles.ContainsKey(p.Type));
            if (outIdx < 0) continue;
            var outParam = c.Params[outIdx];

            // consider handle params BEFORE out param as potential parents
            var parents = c.Params.Take(outIdx).Where(p => p.IsHandle).ToList();
            if (parents.Count != 1) continue; // only "simple" case supported: a single parent

            var parent = parents[0].Type;
            bool isMany = outParam.LenAttr != null;

            list.Add(new SimpleProducer(outParam.Type, c.Name, parent, outParam, isMany));
        }

        return list;
    }

    public enum LoadLevel { Global, Instance, Device }
    public LoadLevel ClassifyLoadLevel(Command c) {
        if (c.Params.Count == 0) return LoadLevel.Global;
        var first = c.Params[0].Type;
        if (first == "VkInstance" || first == "VkPhysicalDevice") return LoadLevel.Instance;
        if (first == "VkDevice" || first == "VkQueue" || first == "VkCommandBuffer") return LoadLevel.Device;
        return LoadLevel.Global;
    }

    // ---------- Deterministic keep filter ----------
    public static class LifetimeFilter {
        // Explicit "do not keep" cases (semantically not required after creation)
        private static readonly HashSet<(string st, string mem, string ht)> Exact = new()
        {
            ("VkPipelineShaderStageCreateInfo", "module", "VkShaderModule"),
            ("VkFramebufferCreateInfo",         "renderPass", "VkRenderPass"),
            ("VkGraphicsPipelineCreateInfo",    "renderPass", "VkRenderPass"),
            ("VkSwapchainCreateInfoKHR",        "oldSwapchain", "VkSwapchainKHR"),
        };

        private static readonly string[] NameHints = { "old", "scratch", "staging" };

        public static bool ShouldKeep(string structName, string memberName, string handleType) {
            if (Exact.Contains((structName, memberName, handleType)))
                return false;

            var mn = memberName.ToLowerInvariant();
            if (NameHints.Any(h => mn.Contains(h)))
                return false;

            // Immutable samplers ARE kept
            return true;
        }
    }

    // Traverse CreateInfo/AllocateInfo trees and collect keep-alive owners
    public List<OwnerUse> AnalyzeCreateInfoOwners(Command create, string handleType) {
        var owners = new Dictionary<(string Type, string Struct, string Member), OwnerUse>();

        void Mark(string t, bool many, string srcStruct, string srcMember) {
            if (t == handleType) return; // do not keep self
            if (!LifetimeFilter.ShouldKeep(srcStruct, srcMember, t)) return;

            var key = (t, srcStruct, srcMember);
            if (owners.TryGetValue(key, out var ex))
                owners[key] = new OwnerUse(t, ex.IsMany || many, srcStruct, srcMember);
            else
                owners[key] = new OwnerUse(t, many, srcStruct, srcMember);
        }

        var stack = new Stack<(string type, bool many)>();

        foreach (var p in create.Params.Where(p => p.Name.Contains("CreateInfo") || p.Name.Contains("AllocateInfo"))) {
            if (Structs.ContainsKey(p.Type))
                stack.Push((p.Type, p.IsPointer && p.LenAttr != null));
        }

        var visited = new HashSet<string>();
        while (stack.Count > 0) {
            var (stype, parentMany) = stack.Pop();
            if (!visited.Add(stype)) { /* ok */ }

            if (!Structs.TryGetValue(stype, out var s)) continue;
            foreach (var m in s.Members) {
                bool manyHere = parentMany || (m.IsPointer && m.LenAttr != null);

                if (Handles.ContainsKey(m.Type)) {
                    Mark(m.Type, manyHere, stype, m.Name);
                } else if (Structs.ContainsKey(m.Type)) {
                    stack.Push((m.Type, manyHere));
                }
            }
        }

        return owners.Values
            .GroupBy(o => o.Type)
            .Select(g => new OwnerUse(
                g.Key,
                g.Any(x => x.IsMany),
                g.First().SourceStruct,
                g.First().SourceMember
            ))
            .ToList();
    }
}