using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;

namespace Compiler;

public class ModuleGraph(Dictionary<string, ModuleGraph.Module> modules)
{
    public static Builder Create() => new();

    public List<Module> GetModules()
    {
        return modules.Values.ToList();
    }

    public bool TryResolveModule(string moduleName, [NotNullWhen(true)] out Module? module)
    {
        module = modules.GetValueOrDefault(moduleName);
        return module != null;
    }

    public Manifest CreateManifest()
    {
        return new Manifest(1, GetModules().Select(x => x.CreateManifestModule()).ToList());
    }

    public record Manifest(int Version, IReadOnlyList<Module.ManifestModule> Modules);

    public sealed class Module(string name)
    {
        public string Name { get; } = name;
        private readonly Dictionary<string, CustomTypeInfo> customTypes = new();
        private readonly Dictionary<string, IdentifierInfo> identifierTypes = new();

        public IReadOnlyList<NubType> GetCustomTypes()
        {
            return customTypes.Values.Select(x => x.Type).ToList();
        }

        public IReadOnlyDictionary<string, NubType> GetIdentifierTypes()
        {
            return identifierTypes.ToDictionary(x => x.Key, x => x.Value.Type);
        }

        public bool TryResolveCustomType(string name, bool searchPrivate, [NotNullWhen(true)] out NubType? customType)
        {
            var info = customTypes.GetValueOrDefault(name);
            if (info == null)
            {
                customType = null;
                return false;
            }

            if (info.Exported || searchPrivate)
            {
                customType = info.Type;
                return true;
            }

            customType = null;
            return false;
        }

        public bool TryResolveIdentifierType(string name, bool searchPrivate, [NotNullWhen(true)] out NubType? identifierType)
        {
            var info = identifierTypes.GetValueOrDefault(name);
            if (info == null)
            {
                identifierType = null;
                return false;
            }

            if (info.Exported || searchPrivate)
            {
                identifierType = info.Type;
                return true;
            }

            identifierType = null;
            return false;
        }

        public void AddCustomType(string name, CustomTypeInfo info)
        {
            customTypes.Add(name, info);
        }

        public void AddIdentifier(string name, IdentifierInfo info)
        {
            identifierTypes.Add(name, info);
        }

        public ManifestModule CreateManifestModule()
        {
            var manifestCustomTypes = customTypes.ToDictionary(x => x.Key, x => x.Value.CreateManifestCustomTypeInfo());
            var manifestIdentifiers = identifierTypes.ToDictionary(x => x.Key, x => x.Value.CreateManifestIdentifierInfo());
            return new ManifestModule(Name, manifestCustomTypes, manifestIdentifiers);
        }

        public enum Source
        {
            Ast,
            Lib,
        }

        public class CustomTypeInfo(NubType type, bool exported, Source source)
        {
            public NubType Type { get; } = type;
            public bool Exported { get; } = exported;
            public Source Source { get; } = source;

            public ManifestCustomTypeInfo CreateManifestCustomTypeInfo()
            {
                return new ManifestCustomTypeInfo(Type, Exported);
            }
        }

        public class IdentifierInfo(NubType type, bool exported, Source source)
        {
            public NubType Type { get; } = type;
            public bool Exported { get; } = exported;
            public Source Source { get; } = source;

            public ManifestIdentifierInfo CreateManifestIdentifierInfo()
            {
                return new ManifestIdentifierInfo(Type, Exported);
            }
        }

        public record ManifestModule(string Name, Dictionary<string, ManifestCustomTypeInfo> CustomTypes, Dictionary<string, ManifestIdentifierInfo> Identifiers);
        public record ManifestCustomTypeInfo(NubType Type, bool Exported);
        public record ManifestIdentifierInfo(NubType Type, bool Exported);
    }

    public class Builder
    {
        private readonly List<Ast> asts = [];
        private readonly List<Manifest> manifests = [];

        public void AddAst(Ast ast)
        {
            asts.Add(ast);
        }

        public void AddManifest(Manifest manifest)
        {
            manifests.Add(manifest);
        }

        public ModuleGraph? Build(out List<Diagnostic> diagnostics)
        {
            diagnostics = [];

            var modules = new Dictionary<string, Module>();

            // First pass: Register libraries
            foreach (var manifest in manifests)
            {
                foreach (var manifestModule in manifest.Modules)
                {
                    if (!modules.TryGetValue(manifestModule.Name, out var module))
                    {
                        module = new Module(manifestModule.Name);
                        modules.Add(manifestModule.Name, module);
                    }

                    foreach (var customType in manifestModule.CustomTypes)
                    {
                        module.AddCustomType(customType.Key, new Module.CustomTypeInfo(customType.Value.Type, customType.Value.Exported, Module.Source.Lib));
                    }

                    foreach (var customType in manifestModule.Identifiers)
                    {
                        module.AddIdentifier(customType.Key, new Module.IdentifierInfo(customType.Value.Type, customType.Value.Exported, Module.Source.Lib));
                    }
                }
            }

            var astModuleCache = new Dictionary<Ast, Module>();

            // First pass: Register modules from ast
            foreach (var ast in asts)
            {
                if (!modules.ContainsKey(ast.ModuleName.Ident))
                {
                    var module = new Module(ast.ModuleName.Ident);
                    modules.Add(ast.ModuleName.Ident, module);
                    astModuleCache[ast] = module;
                }
            }

            // Second pass: Register struct types without fields
            foreach (var ast in asts)
            {
                var module = astModuleCache[ast];

                foreach (var structDef in ast.Definitions.OfType<NodeDefinitionStruct>())
                {
                    var type = new NubTypeStruct(module.Name, structDef.Name.Ident, structDef.Packed);
                    var info = new Module.CustomTypeInfo(type, structDef.Exported, Module.Source.Ast);
                    module.AddCustomType(structDef.Name.Ident, info);
                }
            }

            // Third pass: Resolve struct fields
            foreach (var ast in asts)
            {
                var module = astModuleCache[ast];

                foreach (var structDef in ast.Definitions.OfType<NodeDefinitionStruct>())
                {
                    if (!module.TryResolveCustomType(structDef.Name.Ident, true, out var customType))
                        throw new UnreachableException($"{nameof(customType)} should always be registered");

                    var fields = structDef.Fields.Select(f => new NubTypeStruct.Field(f.Name.Ident, ResolveType(f.Type, module.Name))).ToList();
                    ((NubTypeStruct)customType).ResolveFields(fields);
                }
            }

            // Fourth pass: Register identifiers
            foreach (var ast in asts)
            {
                var module = astModuleCache[ast];

                foreach (var funcDef in ast.Definitions.OfType<NodeDefinitionFunc>())
                {
                    var parameters = funcDef.Parameters.Select(x => ResolveType(x.Type, module.Name)).ToList();
                    var returnType = ResolveType(funcDef.ReturnType, module.Name);
                    var funcType = NubTypeFunc.Get(parameters, returnType);
                    var info = new Module.IdentifierInfo(funcType, funcDef.Exported, Module.Source.Ast);
                    module.AddIdentifier(funcDef.Name.Ident, info);
                }

                foreach (var globalVariable in ast.Definitions.OfType<NodeDefinitionGlobalVariable>())
                {
                    var type = ResolveType(globalVariable.Type, module.Name);
                    var info = new Module.IdentifierInfo(type, globalVariable.Exported, Module.Source.Ast);
                    module.AddIdentifier(globalVariable.Name.Ident, info);
                }
            }

            if (diagnostics.Any(x => x.Severity == DiagnosticSeverity.Error))
                return null;

            return new ModuleGraph(modules);

            NubType ResolveType(NodeType node, string currentModule)
            {
                return node switch
                {
                    NodeTypeBool => NubTypeBool.Instance,
                    NodeTypeCustom type => ResolveCustomType(type, currentModule),
                    NodeTypeFunc type => NubTypeFunc.Get(type.Parameters.Select(x => ResolveType(x, currentModule)).ToList(), ResolveType(type.ReturnType, currentModule)),
                    NodeTypePointer type => NubTypePointer.Get(ResolveType(type.To, currentModule)),
                    NodeTypeSInt type => NubTypeSInt.Get(type.Width),
                    NodeTypeUInt type => NubTypeUInt.Get(type.Width),
                    NodeTypeString => NubTypeString.Instance,
                    NodeTypeVoid => NubTypeVoid.Instance,
                    _ => throw new ArgumentOutOfRangeException(nameof(node))
                };
            }

            NubType ResolveCustomType(NodeTypeCustom type, string currentModule)
            {
                var module = modules.GetValueOrDefault(type.Module.Ident);
                if (module == null)
                    throw new CompileException(Diagnostic.Error($"Unknown module: {type.Module.Ident}").Build());

                var includePrivate = currentModule == type.Module.Ident;

                if (!module.TryResolveCustomType(type.Name.Ident, includePrivate, out var customType))
                    throw new CompileException(Diagnostic.Error($"Unknown custom type: {type.Module.Ident}::{type.Name.Ident}").Build());

                return customType;
            }
        }
    }
}