using System.Data.Common;
using System.Diagnostics.CodeAnalysis;

namespace Compiler;

public class TypeChecker
{
    public static TypedNodeDefinitionFunc? CheckFunction(string fileName, string currentModule, NodeDefinitionFunc function, ModuleGraph moduleGraph, out List<Diagnostic> diagnostics)
    {
        return new TypeChecker(fileName, currentModule, function, moduleGraph).CheckFunction(out diagnostics);
    }

    private TypeChecker(string fileName, string currentModule, NodeDefinitionFunc function, ModuleGraph moduleGraph)
    {
        this.fileName = fileName;
        this.currentModule = currentModule;
        this.function = function;
        this.moduleGraph = moduleGraph;
    }

    private readonly string fileName;
    private readonly string currentModule;
    private readonly NodeDefinitionFunc function;
    private readonly ModuleGraph moduleGraph;
    private readonly Scope scope = new();

    private TypedNodeDefinitionFunc? CheckFunction(out List<Diagnostic> diagnostics)
    {
        diagnostics = [];

        var parameters = new List<TypedNodeDefinitionFunc.Param>();
        var invalidParameter = false;
        TypedNodeStatement? body = null;
        NubType? returnType = null;

        using (scope.EnterScope())
        {
            foreach (var parameter in function.Parameters)
            {
                NubType parameterType;

                try
                {
                    parameterType = ResolveType(parameter.Type);
                }
                catch (CompileException e)
                {
                    diagnostics.Add(e.Diagnostic);
                    invalidParameter = true;
                    continue;
                }

                scope.DeclareIdentifier(parameter.Name.Ident, parameterType);
                parameters.Add(new TypedNodeDefinitionFunc.Param(parameter.Tokens, parameter.Name, parameterType));
            }

            try
            {
                body = CheckStatement(function.Body);
            }
            catch (CompileException e)
            {
                diagnostics.Add(e.Diagnostic);
            }

            try
            {
                returnType = ResolveType(function.ReturnType);
            }
            catch (CompileException e)
            {
                diagnostics.Add(e.Diagnostic);
            }

            if (body == null || returnType is null || invalidParameter)
                return null;

            return new TypedNodeDefinitionFunc(function.Tokens, currentModule, function.Name, parameters, body, returnType);
        }
    }

    private TypedNodeStatement CheckStatement(NodeStatement node)
    {
        return node switch
        {
            NodeStatementAssignment statement => CheckStatementAssignment(statement),
            NodeStatementBlock statement => CheckStatementBlock(statement),
            NodeStatementExpression statement => CheckStatementExpression(statement),
            NodeStatementIf statement => CheckStatementIf(statement),
            NodeStatementReturn statement => CheckStatementReturn(statement),
            NodeStatementVariableDeclaration statement => CheckStatementVariableDeclaration(statement),
            NodeStatementWhile statement => CheckStatementWhile(statement),
            NodeStatementMatch statement => CheckStatementMatch(statement),
            _ => throw new ArgumentOutOfRangeException(nameof(node))
        };
    }

    private TypedNodeStatementAssignment CheckStatementAssignment(NodeStatementAssignment statement)
    {
        return new TypedNodeStatementAssignment(statement.Tokens, CheckExpression(statement.Target), CheckExpression(statement.Value));
    }

    private TypedNodeStatementBlock CheckStatementBlock(NodeStatementBlock statement)
    {
        using (scope.EnterScope())
        {
            var statements = statement.Statements.Select(CheckStatement).ToList();
            return new TypedNodeStatementBlock(statement.Tokens, statements);
        }
    }

    private TypedNodeStatementFuncCall CheckStatementExpression(NodeStatementExpression statement)
    {
        if (statement.Expression is not NodeExpressionFuncCall funcCall)
            throw new CompileException(Diagnostic.Error("Expected statement or function call").At(fileName, statement).Build());

        return new TypedNodeStatementFuncCall(statement.Tokens, CheckExpression(funcCall.Target), funcCall.Parameters.Select(CheckExpression).ToList());
    }

    private TypedNodeStatementIf CheckStatementIf(NodeStatementIf statement)
    {
        return new TypedNodeStatementIf(statement.Tokens, CheckExpression(statement.Condition), CheckStatement(statement.ThenBlock), statement.ElseBlock == null ? null : CheckStatement(statement.ElseBlock));
    }

    private TypedNodeStatementReturn CheckStatementReturn(NodeStatementReturn statement)
    {
        return new TypedNodeStatementReturn(statement.Tokens, CheckExpression(statement.Value));
    }

    private TypedNodeStatementVariableDeclaration CheckStatementVariableDeclaration(NodeStatementVariableDeclaration statement)
    {
        var type = ResolveType(statement.Type);
        var value = CheckExpression(statement.Value);

        if (!value.Type.IsAssignableTo(type))
            throw new CompileException(Diagnostic.Error("Type of variable does match type of assigned value").At(fileName, value).Build());

        scope.DeclareIdentifier(statement.Name.Ident, type);

        return new TypedNodeStatementVariableDeclaration(statement.Tokens, statement.Name, type, value);
    }

    private TypedNodeStatementWhile CheckStatementWhile(NodeStatementWhile statement)
    {
        return new TypedNodeStatementWhile(statement.Tokens, CheckExpression(statement.Condition), CheckStatement(statement.Body));
    }

    private TypedNodeStatementMatch CheckStatementMatch(NodeStatementMatch statement)
    {
        var cases = new List<TypedNodeStatementMatch.Case>();
        var target = CheckExpression(statement.Target);
        var enumType = (NubTypeEnum)target.Type;

        foreach (var @case in statement.Cases)
        {
            using (scope.EnterScope())
            {
                scope.DeclareIdentifier(@case.VariableName.Ident, NubTypeEnumVariant.Get(NubTypeEnum.Get(enumType.Module, enumType.Name), @case.Variant.Ident));
                var body = CheckStatement(@case.Body);
                cases.Add(new TypedNodeStatementMatch.Case(@case.Tokens, @case.Variant, @case.VariableName, body));
            }
        }

        return new TypedNodeStatementMatch(statement.Tokens, target, cases);
    }

    private TypedNodeExpression CheckExpression(NodeExpression node)
    {
        return node switch
        {
            NodeExpressionBinary expression => CheckExpressionBinary(expression),
            NodeExpressionUnary expression => CheckExpressionUnary(expression),
            NodeExpressionBoolLiteral expression => CheckExpressionBoolLiteral(expression),
            NodeExpressionLocalIdent expression => CheckExpressionIdent(expression),
            NodeExpressionModuleIdent expression => CheckExpressionModuleIdent(expression),
            NodeExpressionIntLiteral expression => CheckExpressionIntLiteral(expression),
            NodeExpressionMemberAccess expression => CheckExpressionMemberAccess(expression),
            NodeExpressionFuncCall expression => CheckExpressionFuncCall(expression),
            NodeExpressionStringLiteral expression => CheckExpressionStringLiteral(expression),
            NodeExpressionStructLiteral expression => CheckExpressionStructLiteral(expression),
            NodeExpressionEnumLiteral expression => CheckExpressionEnumLiteral(expression),
            _ => throw new ArgumentOutOfRangeException(nameof(node))
        };
    }

    private TypedNodeExpressionBinary CheckExpressionBinary(NodeExpressionBinary expression)
    {
        var left = CheckExpression(expression.Left);
        var right = CheckExpression(expression.Right);
        NubType type;

        switch (expression.Operation)
        {
            case NodeExpressionBinary.Op.Add:
            case NodeExpressionBinary.Op.Subtract:
            case NodeExpressionBinary.Op.Multiply:
            case NodeExpressionBinary.Op.Divide:
            case NodeExpressionBinary.Op.Modulo:
                {
                    if (left.Type is not NubTypeSInt and not NubTypeUInt)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for left hand side arithmetic operation: {left.Type}").At(fileName, left).Build());

                    if (right.Type is not NubTypeSInt and not NubTypeUInt)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for right hand side arithmetic operation: {right.Type}").At(fileName, right).Build());

                    type = left.Type;
                    break;
                }
            case NodeExpressionBinary.Op.LeftShift:
            case NodeExpressionBinary.Op.RightShift:
                {
                    if (left.Type is not NubTypeUInt)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for left hand side of left/right shift operation: {left.Type}").At(fileName, left).Build());

                    if (right.Type is not NubTypeUInt)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for right hand side of left/right shift operation: {right.Type}").At(fileName, right).Build());

                    type = left.Type;
                    break;
                }
            case NodeExpressionBinary.Op.Equal:
            case NodeExpressionBinary.Op.NotEqual:
            case NodeExpressionBinary.Op.LessThan:
            case NodeExpressionBinary.Op.LessThanOrEqual:
            case NodeExpressionBinary.Op.GreaterThan:
            case NodeExpressionBinary.Op.GreaterThanOrEqual:
                {
                    if (left.Type is not NubTypeSInt and not NubTypeUInt)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for left hand side of comparison: {left.Type}").At(fileName, left).Build());

                    if (right.Type is not NubTypeSInt and not NubTypeUInt)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for right hand side of comparison: {right.Type}").At(fileName, right).Build());

                    type = NubTypeBool.Instance;
                    break;
                }
            case NodeExpressionBinary.Op.LogicalAnd:
            case NodeExpressionBinary.Op.LogicalOr:
                {
                    if (left.Type is not NubTypeBool)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for left hand side of logical operation: {left.Type}").At(fileName, left).Build());

                    if (right.Type is not NubTypeBool)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for right hand side of logical operation: {right.Type}").At(fileName, right).Build());

                    type = NubTypeBool.Instance;
                    break;
                }
            default:
                throw new ArgumentOutOfRangeException();
        }

        return new TypedNodeExpressionBinary(expression.Tokens, type, left, CheckExpressionBinaryOperation(expression.Operation), right);
    }

    private static TypedNodeExpressionBinary.Op CheckExpressionBinaryOperation(NodeExpressionBinary.Op op)
    {
        return op switch
        {
            NodeExpressionBinary.Op.Add => TypedNodeExpressionBinary.Op.Add,
            NodeExpressionBinary.Op.Subtract => TypedNodeExpressionBinary.Op.Subtract,
            NodeExpressionBinary.Op.Multiply => TypedNodeExpressionBinary.Op.Multiply,
            NodeExpressionBinary.Op.Divide => TypedNodeExpressionBinary.Op.Divide,
            NodeExpressionBinary.Op.Modulo => TypedNodeExpressionBinary.Op.Modulo,
            NodeExpressionBinary.Op.Equal => TypedNodeExpressionBinary.Op.Equal,
            NodeExpressionBinary.Op.NotEqual => TypedNodeExpressionBinary.Op.NotEqual,
            NodeExpressionBinary.Op.LessThan => TypedNodeExpressionBinary.Op.LessThan,
            NodeExpressionBinary.Op.LessThanOrEqual => TypedNodeExpressionBinary.Op.LessThanOrEqual,
            NodeExpressionBinary.Op.GreaterThan => TypedNodeExpressionBinary.Op.GreaterThan,
            NodeExpressionBinary.Op.GreaterThanOrEqual => TypedNodeExpressionBinary.Op.GreaterThanOrEqual,
            NodeExpressionBinary.Op.LeftShift => TypedNodeExpressionBinary.Op.LeftShift,
            NodeExpressionBinary.Op.RightShift => TypedNodeExpressionBinary.Op.RightShift,
            NodeExpressionBinary.Op.LogicalAnd => TypedNodeExpressionBinary.Op.LogicalAnd,
            NodeExpressionBinary.Op.LogicalOr => TypedNodeExpressionBinary.Op.LogicalOr,
            _ => throw new ArgumentOutOfRangeException(nameof(op), op, null)
        };
    }

    private TypedNodeExpressionUnary CheckExpressionUnary(NodeExpressionUnary expression)
    {
        var target = CheckExpression(expression.Target);
        NubType type;

        switch (expression.Operation)
        {
            case NodeExpressionUnary.Op.Negate:
                {
                    if (target.Type is not NubTypeSInt and not NubTypeUInt)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for negation: {target.Type}").At(fileName, target).Build());

                    type = target.Type;
                    break;
                }
            case NodeExpressionUnary.Op.Invert:
                {
                    if (target.Type is not NubTypeBool)
                        throw new CompileException(Diagnostic.Error($"Unsupported type for inversion: {target.Type}").At(fileName, target).Build());

                    type = NubTypeBool.Instance;
                    break;
                }
            default:
                throw new ArgumentOutOfRangeException();
        }

        return new TypedNodeExpressionUnary(expression.Tokens, type, target, CheckExpressionUnaryOperation(expression.Operation));
    }

    private static TypedNodeExpressionUnary.Op CheckExpressionUnaryOperation(NodeExpressionUnary.Op op)
    {
        return op switch
        {
            NodeExpressionUnary.Op.Negate => TypedNodeExpressionUnary.Op.Negate,
            NodeExpressionUnary.Op.Invert => TypedNodeExpressionUnary.Op.Invert,
            _ => throw new ArgumentOutOfRangeException(nameof(op), op, null)
        };
    }

    private TypedNodeExpressionBoolLiteral CheckExpressionBoolLiteral(NodeExpressionBoolLiteral expression)
    {
        return new TypedNodeExpressionBoolLiteral(expression.Tokens, NubTypeBool.Instance, expression.Value);
    }

    private TypedNodeExpressionLocalIdent CheckExpressionIdent(NodeExpressionLocalIdent expression)
    {
        var type = scope.GetIdentifierType(expression.Value.Ident);
        if (type is null)
            throw new CompileException(Diagnostic.Error($"Identifier '{expression.Value.Ident}' is not declared").At(fileName, expression.Value).Build());

        return new TypedNodeExpressionLocalIdent(expression.Tokens, type, expression.Value);
    }

    private TypedNodeExpressionModuleIdent CheckExpressionModuleIdent(NodeExpressionModuleIdent expression)
    {
        var module = ResolveModule(expression.Module);
        var info = ResolveModuleIdentifier(module, expression.Value);
        return new TypedNodeExpressionModuleIdent(expression.Tokens, info.Type, expression.Module, expression.Value);
    }

    private TypedNodeExpressionIntLiteral CheckExpressionIntLiteral(NodeExpressionIntLiteral expression)
    {
        return new TypedNodeExpressionIntLiteral(expression.Tokens, NubTypeSInt.Get(32), expression.Value);
    }

    private TypedNodeExpressionMemberAccess CheckExpressionMemberAccess(NodeExpressionMemberAccess expression)
    {
        var target = CheckExpression(expression.Target);

        switch (target.Type)
        {
            case NubTypeStruct structType:
            {
                if (!moduleGraph.TryResolveModule(structType.Module, out var module))
                    throw new CompileException(Diagnostic.Error($"Module '{structType.Module}' not found").At(fileName, expression.Target).Build());

                if (!module.TryResolveType(structType.Name, currentModule == structType.Module, out var typeDef))
                    throw new CompileException(Diagnostic.Error($"Type '{structType.Name}' not found in module '{structType.Module}'").At(fileName, expression.Target).Build());

                if (typeDef is not Module.TypeInfoStruct structDef)
                    throw new CompileException(Diagnostic.Error($"Type '{target.Type}' is not a struct").At(fileName, expression.Target).Build());

                var field = structDef.Fields.FirstOrDefault(x => x.Name == expression.Name.Ident);
                if (field == null)
                    throw new CompileException(Diagnostic.Error($"Struct '{target.Type}' does not have a field matching the name '{expression.Name.Ident}'").At(fileName, target).Build());

                return new TypedNodeExpressionMemberAccess(expression.Tokens, field.Type, target, expression.Name);
            }
            case NubTypeEnumVariant enumVariantType:
            {
                if (!moduleGraph.TryResolveModule(enumVariantType.EnumType.Module, out var module))
                    throw new CompileException(Diagnostic.Error($"Module '{enumVariantType.EnumType.Module}' not found").At(fileName, expression.Target).Build());

                if (!module.TryResolveType(enumVariantType.EnumType.Name, currentModule == enumVariantType.EnumType.Module, out var typeDef))
                    throw new CompileException(Diagnostic.Error($"Type '{enumVariantType.EnumType.Name}' not found in module '{enumVariantType.EnumType.Module}'").At(fileName, expression.Target).Build());

                if (typeDef is not Module.TypeInfoEnum enumDef)
                    throw new CompileException(Diagnostic.Error($"Type '{enumVariantType.EnumType.Module}::{enumVariantType.EnumType.Name}' is not an enum").At(fileName, expression.Target).Build());

                var variant = enumDef.Variants.FirstOrDefault(x => x.Name == enumVariantType.Variant);
                if (variant == null)
                    throw new CompileException(Diagnostic.Error($"Type '{target.Type}' does not have a variant named '{enumVariantType.Variant}'").At(fileName, expression.Target).Build());

                var field = variant.Fields.FirstOrDefault(x => x.Name == expression.Name.Ident);
                if (field == null)
                    throw new CompileException(Diagnostic.Error($"Enum variant '{target.Type}' does not have a field matching the name '{expression.Name.Ident}'").At(fileName, target).Build());

                return new TypedNodeExpressionMemberAccess(expression.Tokens, field.Type, target, expression.Name);
            }
            default:
                throw new CompileException(Diagnostic.Error($"{target.Type} has no member '{expression.Name.Ident}'").At(fileName, target).Build());
        }
    }

    private TypedNodeExpressionFuncCall CheckExpressionFuncCall(NodeExpressionFuncCall expression)
    {
        var target = CheckExpression(expression.Target);
        if (target.Type is not NubTypeFunc funcType)
            throw new CompileException(Diagnostic.Error($"Cannot invoke function call on type '{target.Type}'").At(fileName, target).Build());

        var parameters = expression.Parameters.Select(CheckExpression).ToList();

        return new TypedNodeExpressionFuncCall(expression.Tokens, funcType.ReturnType, target, parameters);
    }

    private TypedNodeExpressionStringLiteral CheckExpressionStringLiteral(NodeExpressionStringLiteral expression)
    {
        return new TypedNodeExpressionStringLiteral(expression.Tokens, NubTypeString.Instance, expression.Value);
    }

    private TypedNodeExpressionStructLiteral CheckExpressionStructLiteral(NodeExpressionStructLiteral expression)
    {
        var info = ResolveModuleStruct(ResolveModule(expression.Module), expression.Name);

        var initializers = new List<TypedNodeExpressionStructLiteral.Initializer>();
        foreach (var initializer in expression.Initializers)
        {
            var field = info.Fields.FirstOrDefault(x => x.Name == initializer.Name.Ident);
            if (field == null)
                throw new CompileException(Diagnostic.Error($"Field '{initializer.Name.Ident}' does not exist on struct '{expression.Module.Ident}::{expression.Name.Ident}'").At(fileName, initializer.Name).Build());

            var value = CheckExpression(initializer.Value);
            if (!value.Type.IsAssignableTo(field.Type))
                throw new CompileException(Diagnostic.Error($"Type of assignment ({value.Type}) does not match expected type of field '{field.Name}' ({field.Type})").At(fileName, initializer.Name).Build());

            initializers.Add(new TypedNodeExpressionStructLiteral.Initializer(initializer.Tokens, initializer.Name, value));
        }

        return new TypedNodeExpressionStructLiteral(expression.Tokens, NubTypeStruct.Get(expression.Module.Ident, expression.Name.Ident), initializers);
    }

    private TypedNodeExpressionEnumLiteral CheckExpressionEnumLiteral(NodeExpressionEnumLiteral expression)
    {
        var info = ResolveModuleEnum(ResolveModule(expression.Module), expression.EnumName);

        var variant = info.Variants.FirstOrDefault(x => x.Name == expression.VariantName.Ident);
        if (variant == null)
            throw new CompileException(Diagnostic.Error($"Enum '{expression.Module.Ident}::{expression.EnumName.Ident}' does not have a variant '{expression.VariantName.Ident}'").At(fileName, expression.VariantName).Build());

        var initializers = new List<TypedNodeExpressionEnumLiteral.Initializer>();
        foreach (var initializer in expression.Initializers)
        {
            var field = variant.Fields.FirstOrDefault(x => x.Name == initializer.Name.Ident);
            if (field == null)
                throw new CompileException(Diagnostic.Error($"Field '{initializer.Name.Ident}' does not exist on enum variant '{expression.Module.Ident}::{expression.EnumName.Ident}.{expression.VariantName.Ident}'").At(fileName, initializer.Name).Build());

            var value = CheckExpression(initializer.Value);
            if (!value.Type.IsAssignableTo(field.Type))
                throw new CompileException(Diagnostic.Error($"Type of assignment ({value.Type}) does not match expected type of field '{field.Name}' ({field.Type})").At(fileName, initializer.Name).Build());

            initializers.Add(new TypedNodeExpressionEnumLiteral.Initializer(initializer.Tokens, initializer.Name, value));
        }

        return new TypedNodeExpressionEnumLiteral(expression.Tokens, NubTypeEnumVariant.Get(NubTypeEnum.Get(expression.Module.Ident, expression.EnumName.Ident), expression.VariantName.Ident), initializers);
    }

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

    private NubType ResolveNamedType(NodeTypeNamed type)
    {
        return type.Sections.Count switch
        {
            3 => ResolveThreePartType(type.Sections[0], type.Sections[1], type.Sections[2]),
            2 => ResolveTwoPartType(type.Sections[0], type.Sections[1]),
            1 => ResolveOnePartType(type.Sections[0]),
            _ => throw BasicError("Invalid type name", type)
        };
    }

    private NubType ResolveThreePartType(TokenIdent first, TokenIdent second, TokenIdent third)
    {
        if (TryResolveEnumVariant(currentModule, second.Ident, third.Ident, out var variantType))
            return variantType;

        throw BasicError($"Enum '{first.Ident}::{second.Ident}::{third.Ident}' does not have a variant named '{third.Ident}'", third);
    }

    private NubType ResolveTwoPartType(TokenIdent first, TokenIdent second)
    {
        if (TryResolveEnumVariant(currentModule, first.Ident, second.Ident, out var variantType))
            return variantType;

        var typeInfo = ResolveModuleTypeInfo(ResolveModule(first), second);
        return typeInfo switch
        {
            Module.TypeInfoStruct => NubTypeStruct.Get(first.Ident, second.Ident),
            Module.TypeInfoEnum => NubTypeEnum.Get(first.Ident, second.Ident),
            _ => throw new ArgumentOutOfRangeException(nameof(typeInfo))
        };
    }

    private NubType ResolveOnePartType(TokenIdent name)
    {
        if (!moduleGraph.TryResolveModule(currentModule, out var module))
            throw BasicError($"Module '{currentModule}' not found", name);

        var typeInfo = ResolveModuleTypeInfo(module, name);
        return typeInfo switch
        {
            Module.TypeInfoStruct => NubTypeStruct.Get(currentModule, name.Ident),
            Module.TypeInfoEnum => NubTypeEnum.Get(currentModule, name.Ident),
            _ => throw new ArgumentOutOfRangeException(nameof(typeInfo))
        };
    }

    private Module ResolveModule(TokenIdent name)
    {
        if (!moduleGraph.TryResolveModule(name.Ident, out var module))
            throw BasicError($"Module '{name.Ident}' not found", name);

        return module;
    }

    private Module.IdentifierInfo ResolveModuleIdentifier(Module module, TokenIdent name)
    {
        if (!module.TryResolveIdentifier(name.Ident, currentModule == module.Name, out var type))
            throw BasicError($"Identifier '{module.Name}::{name.Ident}' not found", name);

        return type;
    }

    private Module.TypeInfo ResolveModuleTypeInfo(Module module, TokenIdent name)
    {
        if (!module.TryResolveType(name.Ident, currentModule == module.Name, out var type))
            throw BasicError($"Named type '{module.Name}::{name.Ident}' not found", name);

        return type;
    }

    private Module.TypeInfoEnum ResolveModuleEnum(Module module, TokenIdent name)
    {
        var type = ResolveModuleTypeInfo(module, name);
        if (type is not Module.TypeInfoEnum info)
            throw BasicError($"'{module.Name}::{name.Ident}' is not an enum", name);

        return info;
    }

    private Module.TypeInfoStruct ResolveModuleStruct(Module module, TokenIdent name)
    {
        var type = ResolveModuleTypeInfo(module, name);
        if (type is not Module.TypeInfoStruct info)
            throw BasicError($"'{module.Name}::{name.Ident}' is not a struct", name);

        return info;
    }

    private bool TryResolveEnumVariant(string moduleName, string enumName, string variantName, [NotNullWhen(true)] out NubType? result)
    {
        result = null;

        if (!moduleGraph.TryResolveModule(moduleName, out var module))
            return false;

        if (!module.TryResolveType(enumName, true, out var type))
            return false;

        if (type is not Module.TypeInfoEnum enumInfo)
            return false;

        var variant = enumInfo.Variants.FirstOrDefault(v => v.Name == variantName);

        if (variant == null)
            return false;

        result = NubTypeEnumVariant.Get(NubTypeEnum.Get(moduleName, enumName), variantName);
        return true;
    }

    private CompileException BasicError(string message, TokenIdent ident)
    {
        return new CompileException(Diagnostic.Error(message).At(fileName, ident).Build());
    }

    private CompileException BasicError(string message, Node node)
    {
        return new CompileException(Diagnostic.Error(message).At(fileName, node).Build());
    }

    private sealed class Scope
    {
        private readonly Stack<Dictionary<string, NubType>> scopes = new();

        public IDisposable EnterScope()
        {
            scopes.Push([]);
            return new ScopeGuard(this);
        }

        public void DeclareIdentifier(string name, NubType type)
        {
            scopes.Peek().Add(name, type);
        }

        public NubType? GetIdentifierType(string name)
        {
            foreach (var scope in scopes)
            {
                if (scope.TryGetValue(name, out var type))
                {
                    return type;
                }
            }

            return null;
        }

        private void ExitScope()
        {
            scopes.Pop();
        }

        private sealed class ScopeGuard(Scope owner) : IDisposable
        {
            public void Dispose()
            {
                owner.ExitScope();
            }
        }
    }
}

public abstract class TypedNode(List<Token> tokens)
{
    public List<Token> Tokens { get; } = tokens;
}

public abstract class TypedNodeDefinition(List<Token> tokens, string module) : TypedNode(tokens)
{
    public string Module { get; } = module;
}

public class TypedNodeDefinitionFunc(List<Token> tokens, string module, TokenIdent name, List<TypedNodeDefinitionFunc.Param> parameters, TypedNodeStatement body, NubType returnType) : TypedNodeDefinition(tokens, module)
{
    public TokenIdent Name { get; } = name;
    public List<Param> Parameters { get; } = parameters;
    public TypedNodeStatement Body { get; } = body;
    public NubType ReturnType { get; } = returnType;

    public NubTypeFunc GetNubType()
    {
        return NubTypeFunc.Get(Parameters.Select(x => x.Type).ToList(), ReturnType);
    }

    public class Param(List<Token> tokens, TokenIdent name, NubType type) : TypedNode(tokens)
    {
        public TokenIdent Name { get; } = name;
        public NubType Type { get; } = type;
    }
}

public abstract class TypedNodeStatement(List<Token> tokens) : TypedNode(tokens);

public class TypedNodeStatementBlock(List<Token> tokens, List<TypedNodeStatement> statements) : TypedNodeStatement(tokens)
{
    public List<TypedNodeStatement> Statements { get; } = statements;
}

public class TypedNodeStatementFuncCall(List<Token> tokens, TypedNodeExpression target, List<TypedNodeExpression> parameters) : TypedNodeStatement(tokens)
{
    public TypedNodeExpression Target { get; } = target;
    public List<TypedNodeExpression> Parameters { get; } = parameters;
}

public class TypedNodeStatementReturn(List<Token> tokens, TypedNodeExpression value) : TypedNodeStatement(tokens)
{
    public TypedNodeExpression Value { get; } = value;
}

public class TypedNodeStatementVariableDeclaration(List<Token> tokens, TokenIdent name, NubType type, TypedNodeExpression value) : TypedNodeStatement(tokens)
{
    public TokenIdent Name { get; } = name;
    public NubType Type { get; } = type;
    public TypedNodeExpression Value { get; } = value;
}

public class TypedNodeStatementAssignment(List<Token> tokens, TypedNodeExpression target, TypedNodeExpression value) : TypedNodeStatement(tokens)
{
    public TypedNodeExpression Target { get; } = target;
    public TypedNodeExpression Value { get; } = value;
}

public class TypedNodeStatementIf(List<Token> tokens, TypedNodeExpression condition, TypedNodeStatement thenBlock, TypedNodeStatement? elseBlock) : TypedNodeStatement(tokens)
{
    public TypedNodeExpression Condition { get; } = condition;
    public TypedNodeStatement ThenBlock { get; } = thenBlock;
    public TypedNodeStatement? ElseBlock { get; } = elseBlock;
}

public class TypedNodeStatementWhile(List<Token> tokens, TypedNodeExpression condition, TypedNodeStatement body) : TypedNodeStatement(tokens)
{
    public TypedNodeExpression Condition { get; } = condition;
    public TypedNodeStatement Body { get; } = body;
}

public class TypedNodeStatementMatch(List<Token> tokens, TypedNodeExpression target, List<TypedNodeStatementMatch.Case> cases) : TypedNodeStatement(tokens)
{
    public TypedNodeExpression Target { get; } = target;
    public List<Case> Cases { get; } = cases;

    public class Case(List<Token> tokens, TokenIdent type, TokenIdent variableName, TypedNodeStatement body) : Node(tokens)
    {
        public TokenIdent Type { get; } = type;
        public TokenIdent VariableName { get; } = variableName;
        public TypedNodeStatement Body { get; } = body;
    }
}

public abstract class TypedNodeExpression(List<Token> tokens, NubType type) : TypedNode(tokens)
{
    public NubType Type { get; } = type;
}

public class TypedNodeExpressionIntLiteral(List<Token> tokens, NubType type, TokenIntLiteral value) : TypedNodeExpression(tokens, type)
{
    public TokenIntLiteral Value { get; } = value;
}

public class TypedNodeExpressionStringLiteral(List<Token> tokens, NubType type, TokenStringLiteral value) : TypedNodeExpression(tokens, type)
{
    public TokenStringLiteral Value { get; } = value;
}

public class TypedNodeExpressionBoolLiteral(List<Token> tokens, NubType type, TokenBoolLiteral value) : TypedNodeExpression(tokens, type)
{
    public TokenBoolLiteral Value { get; } = value;
}

public class TypedNodeExpressionStructLiteral(List<Token> tokens, NubType type, List<TypedNodeExpressionStructLiteral.Initializer> initializers) : TypedNodeExpression(tokens, type)
{
    public List<Initializer> Initializers { get; } = initializers;

    public class Initializer(List<Token> tokens, TokenIdent name, TypedNodeExpression value) : Node(tokens)
    {
        public TokenIdent Name { get; } = name;
        public TypedNodeExpression Value { get; } = value;
    }
}

public class TypedNodeExpressionEnumLiteral(List<Token> tokens, NubType type, List<TypedNodeExpressionEnumLiteral.Initializer> initializers) : TypedNodeExpression(tokens, type)
{
    public List<Initializer> Initializers { get; } = initializers;

    public class Initializer(List<Token> tokens, TokenIdent name, TypedNodeExpression value) : Node(tokens)
    {
        public TokenIdent Name { get; } = name;
        public TypedNodeExpression Value { get; } = value;
    }
}

public class TypedNodeExpressionMemberAccess(List<Token> tokens, NubType type, TypedNodeExpression target, TokenIdent name) : TypedNodeExpression(tokens, type)
{
    public TypedNodeExpression Target { get; } = target;
    public TokenIdent Name { get; } = name;
}

public class TypedNodeExpressionFuncCall(List<Token> tokens, NubType type, TypedNodeExpression target, List<TypedNodeExpression> parameters) : TypedNodeExpression(tokens, type)
{
    public TypedNodeExpression Target { get; } = target;
    public List<TypedNodeExpression> Parameters { get; } = parameters;
}

public class TypedNodeExpressionLocalIdent(List<Token> tokens, NubType type, TokenIdent value) : TypedNodeExpression(tokens, type)
{
    public TokenIdent Value { get; } = value;
}

public class TypedNodeExpressionModuleIdent(List<Token> tokens, NubType type, TokenIdent module, TokenIdent value) : TypedNodeExpression(tokens, type)
{
    public TokenIdent Module { get; } = module;
    public TokenIdent Value { get; } = value;
}

public class TypedNodeExpressionBinary(List<Token> tokens, NubType type, TypedNodeExpression left, TypedNodeExpressionBinary.Op operation, TypedNodeExpression right) : TypedNodeExpression(tokens, type)
{
    public TypedNodeExpression Left { get; } = left;
    public Op Operation { get; } = operation;
    public TypedNodeExpression Right { get; } = right;

    public enum Op
    {
        Add,
        Subtract,
        Multiply,
        Divide,
        Modulo,

        Equal,
        NotEqual,
        LessThan,
        LessThanOrEqual,
        GreaterThan,
        GreaterThanOrEqual,

        LeftShift,
        RightShift,

        // BitwiseAnd,
        // BitwiseXor,
        // BitwiseOr,

        LogicalAnd,
        LogicalOr,
    }
}

public class TypedNodeExpressionUnary(List<Token> tokens, NubType type, TypedNodeExpression target, TypedNodeExpressionUnary.Op op) : TypedNodeExpression(tokens, type)
{
    public TypedNodeExpression Target { get; } = target;
    public Op Operation { get; } = op;

    public enum Op
    {
        Negate,
        Invert,
    }
}