namespace Compiler;

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

    public NubTypeStruct ResolveStruct(string moduleName, string name)
    {
        var module = modules.GetValueOrDefault(moduleName);
        if (module == null)
            throw new CompileException(Diagnostic.Error($"Module '{moduleName}' not found").Build());

        var structType = module.ResolveStruct(name);
        if (structType == null)
            throw new CompileException(Diagnostic.Error($"Struct '{moduleName}::{name}' not found").Build());

        return structType;
    }

    public NubType ResolveIdentifier(string moduleName, string name)
    {
        var module = modules.GetValueOrDefault(moduleName);
        if (module == null)
            throw new CompileException(Diagnostic.Error($"Module '{moduleName}' not found").Build());

        var identType = module.ResolveIdentifier(name);
        if (identType == null)
            throw new CompileException(Diagnostic.Error($"Function '{moduleName}::{name}' not found").Build());

        return identType;
    }

    public List<NubTypeStruct> GetStructTypes()
    {
        return modules.SelectMany(x => x.Value.GetStructTypes()).ToList();
    }

    public sealed class Module
    {
        private readonly Dictionary<string, NubTypeStruct> structTypes = new();
        private readonly Dictionary<string, NubType> identifierTypes = new();

        public List<NubTypeStruct> GetStructTypes()
        {
            return structTypes.Values.ToList();
        }

        public List<NubType> GetIdentifiers()
        {
            return identifierTypes.Values.ToList();
        }

        public NubTypeStruct? ResolveStruct(string name)
        {
            return structTypes.GetValueOrDefault(name);
        }

        public NubType? ResolveIdentifier(string name)
        {
            return identifierTypes.GetValueOrDefault(name);
        }

        public void AddStruct(string name, NubTypeStruct structType)
        {
            structTypes.Add(name, structType);
        }

        public void AddIdentifier(string name, NubType identifier)
        {
            identifierTypes.Add(name, identifier);
        }
    }

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

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

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

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

            // First pass: Register modules
            foreach (var ast in asts)
            {
                var moduleDefinitions = ast.Definitions.OfType<NodeDefinitionModule>().ToList();
                var currentModule = moduleDefinitions[0].Name.Ident;

                if (!modules.ContainsKey(currentModule))
                    modules.Add(currentModule, new Module());
            }

            // Second pass: Register struct types without fields
            foreach (var ast in asts)
            {
                var moduleDefinitions = ast.Definitions.OfType<NodeDefinitionModule>().ToList();
                var currentModule = moduleDefinitions[0].Name.Ident;

                if (!modules.TryGetValue(currentModule, out var module))
                {
                    module = new Module();
                    modules[currentModule] = module;
                }

                foreach (var structDef in ast.Definitions.OfType<NodeDefinitionStruct>())
                {
                    module.AddStruct(structDef.Name.Ident, new NubTypeStruct());
                }
            }

            // Third pass: Resolve struct fields
            foreach (var ast in asts)
            {
                var moduleDefinitions = ast.Definitions.OfType<NodeDefinitionModule>().ToList();
                var module = modules[moduleDefinitions[0].Name.Ident];

                foreach (var structDef in ast.Definitions.OfType<NodeDefinitionStruct>())
                {
                    var structType = module.ResolveStruct(structDef.Name.Ident);
                    var fields = structDef.Fields.Select(f => new NubTypeStruct.Field(f.Name.Ident, Resolve(f.Type))).ToList();
                    structType.ResolveFields(fields);
                }
            }

            // Fourth pass: Register identifiers
            foreach (var ast in asts)
            {
                var moduleDefinitions = ast.Definitions.OfType<NodeDefinitionModule>().ToList();
                var module = modules[moduleDefinitions[0].Name.Ident];

                foreach (var funcDef in ast.Definitions.OfType<NodeDefinitionFunc>())
                {
                    var parameters = funcDef.Parameters.Select(x => Resolve(x.Type)).ToList();
                    var returnType = Resolve(funcDef.ReturnType);
                    var funcType = new NubTypeFunc(parameters, returnType);
                    module.AddIdentifier(funcDef.Name.Ident, funcType);
                }
            }

            return new ModuleGraph(modules);

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

            NubTypeStruct ResolveStruct(NodeTypeCustom type)
            {
                var module = modules.GetValueOrDefault(type.Module.Ident);
                if (module == null)
                    throw new CompileException(Diagnostic.Error($"Unknown module: {type.Module.Ident}").Build());

                var structType = module.ResolveStruct(type.Name.Ident);
                if (structType == null)
                    throw new CompileException(Diagnostic.Error($"Unknown custom type: {type.Module.Ident}::{type.Name.Ident}").Build());

                return structType;
            }
        }
    }
}