using System.Diagnostics;
using NubLang.Diagnostics;
using NubLang.Modules;
using NubLang.Syntax;

namespace NubLang.Ast;

public sealed class TypeChecker
{
    private readonly SyntaxTree _syntaxTree;
    private readonly Dictionary<string, Module> _visibleModules;

    private readonly Stack<Scope> _scopes = [];
    private readonly Stack<NubType> _funcReturnTypes = [];
    private readonly Dictionary<(string Module, string Name), NubType> _typeCache = new();
    private readonly HashSet<(string Module, string Name)> _resolvingTypes = [];
    private readonly HashSet<string> _checkedTemplateStructs = [];

    private Scope CurrentScope => _scopes.Peek();

    public TypeChecker(SyntaxTree syntaxTree, ModuleRepository moduleRepository)
    {
        _syntaxTree = syntaxTree;
        _visibleModules = moduleRepository
            .Modules()
            .Where(x => syntaxTree.Metadata.Imports.Contains(x.Key) || _syntaxTree.Metadata.ModuleName == x.Key)
            .ToDictionary();
    }

    public List<DefinitionNode> Definitions { get; } = [];
    public List<Diagnostic> Diagnostics { get; } = [];
    public List<NubStructType> ReferencedStructTypes { get; } = [];

    public void Check()
    {
        _scopes.Clear();
        _funcReturnTypes.Clear();
        _typeCache.Clear();
        _resolvingTypes.Clear();
        _checkedTemplateStructs.Clear();

        Diagnostics.Clear();
        Definitions.Clear();
        ReferencedStructTypes.Clear();

        using (BeginRootScope(_syntaxTree.Metadata.ModuleName))
        {
            foreach (var definition in _syntaxTree.Definitions)
            {
                try
                {
                    switch (definition)
                    {
                        case FuncSyntax funcSyntax:
                            Definitions.Add(CheckFuncDefinition(funcSyntax));
                            break;
                        case StructSyntax structSyntax:
                            Definitions.Add(CheckStructDefinition(structSyntax));
                            break;
                        case GlobalVariableSyntax globalVariableSyntax:
                            Definitions.Add(CheckGlobalVariable(globalVariableSyntax));
                            break;
                        case StructTemplateSyntax:
                            break;
                        default:
                            throw new ArgumentOutOfRangeException();
                    }
                }
                catch (TypeCheckerException e)
                {
                    Diagnostics.Add(e.Diagnostic);
                }
            }
        }
    }

    private ScopeDisposer BeginScope()
    {
        if (_scopes.TryPeek(out var scope))
        {
            _scopes.Push(scope.SubScope());
        }
        else
        {
            _scopes.Push(new Scope(_syntaxTree.Metadata.ModuleName));
        }

        return new ScopeDisposer(this);
    }

    private ScopeDisposer BeginRootScope(string moduleName)
    {
        if (!_visibleModules.TryGetValue(moduleName, out var moduleScope))
        {
            throw new TypeCheckerException(Diagnostic.Error($"Module with name {moduleName} not found").Build());
        }

        var scope = new Scope(moduleName);
        _scopes.Push(scope);

        foreach (var globalVariable in moduleScope.GlobalVariables(true))
        {
            NubType? type;

            if (globalVariable.ExplicitType != null)
            {
                type = ResolveType(globalVariable.ExplicitType);
                var valueExpression = CheckExpression(globalVariable.Value, type);

                if (valueExpression.Type != ResolveType(globalVariable.ExplicitType))
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("Value does not match explicit type of global variable")
                        .At(globalVariable.Value)
                        .Build());
                }
            }
            else
            {
                type = CheckExpression(globalVariable.Value).Type;
            }

            scope.DeclareVariable(new Variable(globalVariable.Name, type, VariableKind.LValue));
        }

        return new ScopeDisposer(this);
    }

    private sealed class ScopeDisposer(TypeChecker owner) : IDisposable
    {
        private bool _disposed;

        public void Dispose()
        {
            if (_disposed) return;
            owner._scopes.Pop();
            _disposed = true;
        }
    }

    private StructNode CheckStructDefinition(StructSyntax node)
    {
        var fieldTypes = node.Fields
            .Select(x => new NubStructFieldType(x.Name, ResolveType(x.Type), x.Value != null))
            .ToList();

        var fieldFunctions = node.Functions
            .Select(x =>
            {
                var parameters = x.Signature.Parameters.Select(y => ResolveType(y.Type)).ToList();
                var returnType = ResolveType(x.Signature.ReturnType);
                return new NubStructFuncType(x.Name, parameters, returnType);
            })
            .ToList();

        var structType = new NubStructType(CurrentScope.Module, node.Name, fieldTypes, fieldFunctions);

        CurrentScope.DeclareVariable(new Variable("this", structType, VariableKind.RValue));

        var fields = node.Fields.Select(CheckStructField).ToList();
        var functions = node.Functions.Select(CheckStructFunc).ToList();

        return new StructNode(CurrentScope.Module, node.Name, fields, functions);
    }

    private StructFuncNode CheckStructFunc(StructFuncSyntax function)
    {
        foreach (var parameter in function.Signature.Parameters)
        {
            CurrentScope.DeclareVariable(new Variable(parameter.Name, ResolveType(parameter.Type), VariableKind.RValue));
        }

        _funcReturnTypes.Push(ResolveType(function.Signature.ReturnType));
        var body = CheckBlock(function.Body);
        _funcReturnTypes.Pop();
        return new StructFuncNode(function.Name, CheckFuncSignature(function.Signature), body);
    }

    private StructFieldNode CheckStructField(StructFieldSyntax field)
    {
        ExpressionNode? value = null;
        if (field.Value != null)
        {
            value = CheckExpression(field.Value, ResolveType(field.Type));
        }

        return new StructFieldNode(field.Name, ResolveType(field.Type), value);
    }

    private FuncNode CheckFuncDefinition(FuncSyntax node)
    {
        foreach (var parameter in node.Signature.Parameters)
        {
            CurrentScope.DeclareVariable(new Variable(parameter.Name, ResolveType(parameter.Type), VariableKind.RValue));
        }

        var signature = CheckFuncSignature(node.Signature);

        BlockNode? body = null;
        if (node.Body != null)
        {
            _funcReturnTypes.Push(signature.ReturnType);

            body = CheckBlock(node.Body);

            if (!AlwaysReturns(body))
            {
                if (signature.ReturnType is NubVoidType)
                {
                    body.Statements.Add(new ReturnNode(null));
                }
                else
                {
                    Diagnostics.Add(Diagnostic
                        .Error("Not all code paths return a value")
                        .At(node.Body)
                        .Build());
                }
            }

            _funcReturnTypes.Pop();
        }

        return new FuncNode(CurrentScope.Module, node.Name, node.ExternSymbol, signature, body);
    }

    private GlobalVariableNode CheckGlobalVariable(GlobalVariableSyntax node)
    {
        return new GlobalVariableNode(CurrentScope.Module, node.Name, CheckExpression(node.Value));
    }

    private AssignmentNode CheckAssignment(AssignmentSyntax statement)
    {
        var target = CheckExpression(statement.Target);
        if (target is not LValueExpressionNode lValue)
        {
            throw new TypeCheckerException(Diagnostic.Error("Cannot assign to an rvalue").At(statement).Build());
        }

        var value = CheckExpression(statement.Value, lValue.Type);
        return new AssignmentNode(lValue, value);
    }

    private IfNode CheckIf(IfSyntax statement)
    {
        var condition = CheckExpression(statement.Condition, new NubBoolType());
        var body = CheckBlock(statement.Body);
        Variant<IfNode, BlockNode>? elseStatement = null;
        if (statement.Else.HasValue)
        {
            elseStatement = statement.Else.Value.Match<Variant<IfNode, BlockNode>>(elif => CheckIf(elif), el => CheckBlock(el));
        }

        return new IfNode(condition, body, elseStatement);
    }

    private ReturnNode CheckReturn(ReturnSyntax statement)
    {
        ExpressionNode? value = null;

        if (statement.Value != null)
        {
            value = CheckExpression(statement.Value, _funcReturnTypes.Peek());
        }

        return new ReturnNode(value);
    }

    private StatementNode CheckStatementExpression(StatementExpressionSyntax statement)
    {
        var expression = CheckExpression(statement.Expression);

        return expression switch
        {
            FuncCallNode funcCall => new StatementFuncCallNode(funcCall),
            StructFuncCallNode structFuncCall => new StatementStructFuncCallNode(structFuncCall),
            _ => throw new TypeCheckerException(Diagnostic.Error("Expressions statements can only be function calls").At(statement).Build())
        };
    }

    private VariableDeclarationNode CheckVariableDeclaration(VariableDeclarationSyntax statement)
    {
        NubType? type = null;
        ExpressionNode? assignmentNode = null;

        if (statement.ExplicitType != null)
        {
            type = ResolveType(statement.ExplicitType);
        }

        if (statement.Assignment != null)
        {
            assignmentNode = CheckExpression(statement.Assignment, type);
            type ??= assignmentNode.Type;
        }

        if (type == null)
        {
            throw new TypeCheckerException(Diagnostic.Error($"Cannot infer type of variable {statement.Name}").At(statement).Build());
        }

        CurrentScope.DeclareVariable(new Variable(statement.Name, type, VariableKind.LValue));

        return new VariableDeclarationNode(statement.Name, assignmentNode, type);
    }

    private WhileNode CheckWhile(WhileSyntax statement)
    {
        var condition = CheckExpression(statement.Condition, new NubBoolType());
        var body = CheckBlock(statement.Body);
        return new WhileNode(condition, body);
    }

    private FuncSignatureNode CheckFuncSignature(FuncSignatureSyntax statement)
    {
        var parameters = new List<FuncParameterNode>();
        foreach (var parameter in statement.Parameters)
        {
            parameters.Add(new FuncParameterNode(parameter.Name, ResolveType(parameter.Type)));
        }

        return new FuncSignatureNode(parameters, ResolveType(statement.ReturnType));
    }

    private ExpressionNode CheckExpression(ExpressionSyntax node, NubType? expectedType = null)
    {
        var result = node switch
        {
            AddressOfSyntax expression => CheckAddressOf(expression),
            ArrayIndexAccessSyntax expression => CheckArrayIndexAccess(expression),
            ArrayInitializerSyntax expression => CheckArrayInitializer(expression),
            BinaryExpressionSyntax expression => CheckBinaryExpression(expression, expectedType),
            UnaryExpressionSyntax expression => CheckUnaryExpression(expression, expectedType),
            DereferenceSyntax expression => CheckDereference(expression),
            MemberFuncCallSyntax expression => CheckMemberFuncCall(expression),
            FuncCallSyntax expression => CheckFuncCall(expression),
            LocalIdentifierSyntax expression => CheckLocalIdentifier(expression),
            ModuleIdentifierSyntax expression => CheckModuleIdentifier(expression),
            BoolLiteralSyntax expression => CheckBoolLiteral(expression),
            StringLiteralSyntax expression => CheckStringLiteral(expression, expectedType),
            IntLiteralSyntax expression => CheckIntLiteral(expression, expectedType),
            FloatLiteralSyntax expression => CheckFloatLiteral(expression, expectedType),
            MemberAccessSyntax expression => CheckMemberAccess(expression),
            StructInitializerSyntax expression => CheckStructInitializer(expression, expectedType),
            InterpretBuiltinSyntax expression => CheckExpression(expression.Target) with { Type = ResolveType(expression.Type) },
            SizeBuiltinSyntax expression => new SizeBuiltinNode(new NubIntType(false, 64), ResolveType(expression.Type)),
            FloatToIntBuiltinSyntax expression => CheckFloatToInt(expression),
            _ => throw new ArgumentOutOfRangeException(nameof(node))
        };

        if (expectedType == null || result.Type == expectedType)
        {
            return result;
        }

        if (result.Type is NubIntType sourceIntType && expectedType is NubIntType targetIntType)
        {
            if (sourceIntType.Signed == targetIntType.Signed && sourceIntType.Width < targetIntType.Width)
            {
                return new ConvertIntNode(targetIntType, result, sourceIntType, targetIntType);
            }
        }

        if (result.Type is NubFloatType sourceFloatType && expectedType is NubFloatType targetFloatType)
        {
            if (sourceFloatType.Width < targetFloatType.Width)
            {
                return new ConvertFloatNode(targetFloatType, result, sourceFloatType, targetFloatType);
            }
        }

        throw new TypeCheckerException(Diagnostic.Error($"Cannot convert {result.Type} to {expectedType}").At(node).Build());
    }

    private FloatToIntBuiltinNode CheckFloatToInt(FloatToIntBuiltinSyntax expression)
    {
        var value = CheckExpression(expression.Value);
        if (value.Type is not NubFloatType sourceFloatType)
        {
            throw new TypeCheckerException(Diagnostic
                .Error("Source type of float to int conversion must be an float")
                .At(expression.Value)
                .Build());
        }

        var targetType = ResolveType(expression.Type);
        if (targetType is not NubIntType targetIntType)
        {
            throw new TypeCheckerException(Diagnostic
                .Error("Target type of float to int conversion must be an integer")
                .At(expression.Type)
                .Build());
        }

        return new FloatToIntBuiltinNode(targetIntType, value, sourceFloatType, targetIntType);
    }

    private AddressOfNode CheckAddressOf(AddressOfSyntax expression)
    {
        var target = CheckExpression(expression.Target);
        if (target is not LValueExpressionNode lvalue)
        {
            throw new TypeCheckerException(Diagnostic.Error("Cannot take address of an rvalue").At(expression).Build());
        }

        var type = new NubPointerType(target.Type);
        return new AddressOfNode(type, lvalue);
    }

    private ArrayIndexAccessNode CheckArrayIndexAccess(ArrayIndexAccessSyntax expression)
    {
        var index = CheckExpression(expression.Index, new NubIntType(false, 64));
        var target = CheckExpression(expression.Target);
        if (target.Type is not NubArrayType arrayType)
        {
            throw new TypeCheckerException(Diagnostic.Error($"Cannot use array indexer on type {target.Type}").At(expression).Build());
        }

        return new ArrayIndexAccessNode(arrayType.ElementType, target, index);
    }

    private ArrayInitializerNode CheckArrayInitializer(ArrayInitializerSyntax expression)
    {
        var elementType = ResolveType(expression.ElementType);
        var type = new NubArrayType(elementType);
        var capacity = CheckExpression(expression.Capacity);
        return new ArrayInitializerNode(type, capacity, elementType);
    }

    private BinaryExpressionNode CheckBinaryExpression(BinaryExpressionSyntax expression, NubType? expectedType)
    {
        var op = expression.Operator switch
        {
            BinaryOperatorSyntax.Equal => BinaryOperator.Equal,
            BinaryOperatorSyntax.NotEqual => BinaryOperator.NotEqual,
            BinaryOperatorSyntax.GreaterThan => BinaryOperator.GreaterThan,
            BinaryOperatorSyntax.GreaterThanOrEqual => BinaryOperator.GreaterThanOrEqual,
            BinaryOperatorSyntax.LessThan => BinaryOperator.LessThan,
            BinaryOperatorSyntax.LessThanOrEqual => BinaryOperator.LessThanOrEqual,
            BinaryOperatorSyntax.LogicalAnd => BinaryOperator.LogicalAnd,
            BinaryOperatorSyntax.LogicalOr => BinaryOperator.LogicalOr,
            BinaryOperatorSyntax.Plus => BinaryOperator.Plus,
            BinaryOperatorSyntax.Minus => BinaryOperator.Minus,
            BinaryOperatorSyntax.Multiply => BinaryOperator.Multiply,
            BinaryOperatorSyntax.Divide => BinaryOperator.Divide,
            BinaryOperatorSyntax.Modulo => BinaryOperator.Modulo,
            BinaryOperatorSyntax.LeftShift => BinaryOperator.LeftShift,
            BinaryOperatorSyntax.RightShift => BinaryOperator.RightShift,
            BinaryOperatorSyntax.BitwiseAnd => BinaryOperator.BitwiseAnd,
            BinaryOperatorSyntax.BitwiseXor => BinaryOperator.BitwiseXor,
            BinaryOperatorSyntax.BitwiseOr => BinaryOperator.BitwiseOr,
            _ => throw new ArgumentOutOfRangeException()
        };

        switch (expression.Operator)
        {
            case BinaryOperatorSyntax.Equal:
            case BinaryOperatorSyntax.NotEqual:
            {
                // note(nub31): Expected type should not be passed down since the resulting type is different than operands
                var left = CheckExpression(expression.Left);
                if (left.Type is not NubIntType and not NubFloatType and not NubBoolType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("Equal and not equal operators must must be used with int, float or bool types")
                        .At(expression.Left)
                        .Build());
                }

                var right = CheckExpression(expression.Right, left.Type);

                return new BinaryExpressionNode(new NubBoolType(), left, op, right);
            }
            case BinaryOperatorSyntax.GreaterThan:
            case BinaryOperatorSyntax.GreaterThanOrEqual:
            case BinaryOperatorSyntax.LessThan:
            case BinaryOperatorSyntax.LessThanOrEqual:
            {
                // note(nub31): Expected type should not be passed down since the resulting type is different than operands
                var left = CheckExpression(expression.Left);
                if (left.Type is not NubIntType and not NubFloatType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("Greater than and less than operators must must be used with int or float types")
                        .At(expression.Left)
                        .Build());
                }

                var right = CheckExpression(expression.Right, left.Type);

                return new BinaryExpressionNode(new NubBoolType(), left, op, right);
            }
            case BinaryOperatorSyntax.LogicalAnd:
            case BinaryOperatorSyntax.LogicalOr:
            {
                // note(nub31): Expected type should not be passed down since the resulting type is different than operands
                var left = CheckExpression(expression.Left);
                if (left.Type is not NubBoolType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("Logical and/or must must be used with bool types")
                        .At(expression.Left)
                        .Build());
                }

                var right = CheckExpression(expression.Right, left.Type);

                return new BinaryExpressionNode(new NubBoolType(), left, op, right);
            }
            case BinaryOperatorSyntax.Plus:
            {
                var left = CheckExpression(expression.Left, expectedType);
                if (left.Type is NubIntType or NubFloatType or NubStringType or NubCStringType)
                {
                    var right = CheckExpression(expression.Right, left.Type);
                    return new BinaryExpressionNode(left.Type, left, op, right);
                }

                throw new TypeCheckerException(Diagnostic
                    .Error("The plus operator must be used with int, float, cstring or string types")
                    .At(expression.Left)
                    .Build());
            }
            case BinaryOperatorSyntax.Minus:
            case BinaryOperatorSyntax.Multiply:
            case BinaryOperatorSyntax.Divide:
            case BinaryOperatorSyntax.Modulo:
            {
                var left = CheckExpression(expression.Left, expectedType);
                if (left.Type is not NubIntType and not NubFloatType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("Math operators must be used with int or float types")
                        .At(expression.Left)
                        .Build());
                }

                var right = CheckExpression(expression.Right, left.Type);

                return new BinaryExpressionNode(left.Type, left, op, right);
            }
            case BinaryOperatorSyntax.LeftShift:
            case BinaryOperatorSyntax.RightShift:
            case BinaryOperatorSyntax.BitwiseAnd:
            case BinaryOperatorSyntax.BitwiseXor:
            case BinaryOperatorSyntax.BitwiseOr:
            {
                var left = CheckExpression(expression.Left, expectedType);
                if (left.Type is not NubIntType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("Bitwise operators must be used with int types")
                        .At(expression.Left)
                        .Build());
                }

                var right = CheckExpression(expression.Right, left.Type);

                return new BinaryExpressionNode(left.Type, left, op, right);
            }
            default:
            {
                throw new ArgumentOutOfRangeException();
            }
        }
    }

    private UnaryExpressionNode CheckUnaryExpression(UnaryExpressionSyntax expression, NubType? expectedType)
    {
        switch (expression.Operator)
        {
            case UnaryOperatorSyntax.Negate:
            {
                var operand = CheckExpression(expression.Operand, expectedType);
                if (operand.Type is not NubIntType { Signed: true } and not NubFloatType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("Negation operator must be used with signed integer or float types")
                        .At(expression)
                        .Build());
                }

                return new UnaryExpressionNode(operand.Type, UnaryOperator.Negate, operand);
            }
            case UnaryOperatorSyntax.Invert:
            {
                var operand = CheckExpression(expression.Operand, expectedType);
                if (operand.Type is not NubBoolType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("Invert operator must be used with booleans")
                        .At(expression)
                        .Build());
                }

                return new UnaryExpressionNode(operand.Type, UnaryOperator.Invert, operand);
            }
            default:
            {
                throw new ArgumentOutOfRangeException();
            }
        }
    }

    private DereferenceNode CheckDereference(DereferenceSyntax expression)
    {
        var target = CheckExpression(expression.Target);
        if (target.Type is not NubPointerType pointerType)
        {
            throw new TypeCheckerException(Diagnostic.Error($"Cannot dereference non-pointer type {target.Type}").At(expression).Build());
        }

        return new DereferenceNode(pointerType.BaseType, target);
    }

    private FuncCallNode CheckFuncCall(FuncCallSyntax expression)
    {
        var accessor = CheckExpression(expression.Expression);
        if (accessor.Type is not NubFuncType funcType)
        {
            throw new TypeCheckerException(Diagnostic.Error($"Cannot call non-function type {accessor.Type}").At(expression.Expression).Build());
        }

        if (expression.Parameters.Count != funcType.Parameters.Count)
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Function {funcType} expects {funcType.Parameters.Count} parameters but got {expression.Parameters.Count}")
                .At(expression.Parameters.LastOrDefault(expression))
                .Build());
        }

        var parameters = new List<ExpressionNode>();
        for (var i = 0; i < expression.Parameters.Count; i++)
        {
            var parameter = expression.Parameters[i];
            var expectedType = funcType.Parameters[i];

            var parameterExpression = CheckExpression(parameter, expectedType);
            if (parameterExpression.Type != expectedType)
            {
                throw new TypeCheckerException(Diagnostic
                    .Error($"Parameter {i + 1} does not match the type {expectedType} for function {funcType}")
                    .At(parameter)
                    .Build());
            }

            parameters.Add(parameterExpression);
        }

        return new FuncCallNode(funcType.ReturnType, accessor, parameters);
    }

    private StructFuncCallNode CheckMemberFuncCall(MemberFuncCallSyntax expression)
    {
        // todo(nub31): When adding interfaces, also support other types than structs
        var target = CheckExpression(expression.Target);
        if (target.Type is NubStructType structType)
        {
            var function = structType.Functions.FirstOrDefault(x => x.Name == expression.Name);
            if (function == null)
            {
                throw new TypeCheckerException(Diagnostic
                    .Error($"Function {expression.Name} not found on struct {structType}")
                    .At(expression)
                    .Build());
            }

            if (expression.Parameters.Count != function.Parameters.Count)
            {
                throw new TypeCheckerException(Diagnostic
                    .Error($"Function {function.Name} expects {function.Parameters.Count} parameters but got {expression.Parameters.Count}")
                    .At(expression.Parameters.LastOrDefault(expression))
                    .Build());
            }

            var parameters = new List<ExpressionNode>();
            for (var i = 0; i < expression.Parameters.Count; i++)
            {
                var parameter = expression.Parameters[i];
                var expectedType = function.Parameters[i];

                var parameterExpression = CheckExpression(parameter, expectedType);
                if (parameterExpression.Type != expectedType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error($"Parameter {i + 1} does not match the type {expectedType} for function {function}")
                        .At(parameter)
                        .Build());
                }

                parameters.Add(parameterExpression);
            }

            return new StructFuncCallNode(function.ReturnType, structType.Module, structType.Name, expression.Name, target, parameters);
        }

        throw new TypeCheckerException(Diagnostic
            .Error($"No function {expression.Name} exists on type {target.Type}")
            .At(expression)
            .Build());
    }

    private ExpressionNode CheckLocalIdentifier(LocalIdentifierSyntax expression)
    {
        // First, look in the current scope for a matching identifier
        var scopeIdent = CurrentScope.LookupVariable(expression.Name);
        if (scopeIdent != null)
        {
            return scopeIdent.Kind switch
            {
                VariableKind.LValue => new LValueIdentifierNode(scopeIdent.Type, expression.Name),
                VariableKind.RValue => new RValueIdentifierNode(scopeIdent.Type, expression.Name),
                _ => throw new ArgumentOutOfRangeException()
            };
        }

        // Second, look in the current module for a function matching the identifier
        var module = _visibleModules[CurrentScope.Module];
        var function = module.Functions(true).FirstOrDefault(x => x.Name == expression.Name);

        if (function != null)
        {
            var parameters = function.Signature.Parameters.Select(x => ResolveType(x.Type)).ToList();
            var type = new NubFuncType(parameters, ResolveType(function.Signature.ReturnType));
            return new FuncIdentifierNode(type, CurrentScope.Module, expression.Name, function.ExternSymbol);
        }

        throw new TypeCheckerException(Diagnostic.Error($"Symbol {expression.Name} not found").At(expression).Build());
    }

    private ExpressionNode CheckModuleIdentifier(ModuleIdentifierSyntax expression)
    {
        if (!_visibleModules.TryGetValue(expression.Module, out var module))
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Module {expression.Module} not found")
                .WithHelp($"import \"{expression.Module}\"")
                .At(expression)
                .Build());
        }

        var includePrivate = expression.Module == CurrentScope.Module;

        var globalVariable = module.GlobalVariables(includePrivate).FirstOrDefault(x => x.Name == expression.Name);
        if (globalVariable != null)
        {
            // todo(nub31): This should be done in the global scope
            NubType? type = null;
            if (globalVariable.ExplicitType != null)
            {
                type = ResolveType(globalVariable.ExplicitType);
            }

            return CheckExpression(globalVariable.Value, type);
        }

        // First, look for the exported function in the specified module
        var function = module.Functions(includePrivate).FirstOrDefault(x => x.Name == expression.Name);
        if (function != null)
        {
            var parameters = function.Signature.Parameters.Select(x => ResolveType(x.Type)).ToList();
            var type = new NubFuncType(parameters, ResolveType(function.Signature.ReturnType));
            return new FuncIdentifierNode(type, expression.Module, expression.Name, function.ExternSymbol);
        }

        throw new TypeCheckerException(Diagnostic
            .Error($"No exported symbol {expression.Name} not found in module {expression.Module}")
            .At(expression)
            .Build());
    }

    private ExpressionNode CheckStringLiteral(StringLiteralSyntax expression, NubType? expectedType)
    {
        return expectedType switch
        {
            NubCStringType => new CStringLiteralNode(expectedType, expression.Value),
            NubStringType => new StringLiteralNode(expectedType, expression.Value),
            _ => new CStringLiteralNode(new NubCStringType(), expression.Value)
        };
    }

    private ExpressionNode CheckIntLiteral(IntLiteralSyntax expression, NubType? expectedType)
    {
        if (expectedType is NubIntType intType)
        {
            return intType.Signed
                ? new IntLiteralNode(intType, Convert.ToInt64(expression.Value, expression.Base))
                : new UIntLiteralNode(intType, Convert.ToUInt64(expression.Value, expression.Base));
        }

        if (expectedType is NubFloatType floatType)
        {
            return floatType.Width switch
            {
                32 => new Float32LiteralNode(floatType, Convert.ToSingle(Convert.ToInt64(expression.Value, expression.Base))),
                64 => new Float64LiteralNode(floatType, Convert.ToDouble(Convert.ToInt64(expression.Value, expression.Base))),
                _ => throw new ArgumentOutOfRangeException()
            };
        }

        var type = new NubIntType(true, 64);
        return new IntLiteralNode(type, Convert.ToInt64(expression.Value, expression.Base));
    }

    private ExpressionNode CheckFloatLiteral(FloatLiteralSyntax expression, NubType? expectedType)
    {
        var type = expectedType as NubFloatType ?? new NubFloatType(64);
        return type.Width == 32
            ? new Float32LiteralNode(type, Convert.ToSingle(expression.Value))
            : new Float64LiteralNode(type, Convert.ToDouble(expression.Value));
    }

    private BoolLiteralNode CheckBoolLiteral(BoolLiteralSyntax expression)
    {
        return new BoolLiteralNode(new NubBoolType(), expression.Value);
    }

    private ExpressionNode CheckMemberAccess(MemberAccessSyntax expression)
    {
        var target = CheckExpression(expression.Target);
        switch (target.Type)
        {
            case NubStructType structType:
            {
                var field = structType.Fields.FirstOrDefault(x => x.Name == expression.Member);
                if (field == null)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error($"Struct {target.Type} does not have a field with the name {expression.Member}")
                        .At(expression)
                        .Build());
                }

                return new StructFieldAccessNode(field.Type, target, expression.Member);
            }
            default:
            {
                throw new TypeCheckerException(Diagnostic
                    .Error($"Cannot access struct member on non-struct type {target.Type}")
                    .At(expression)
                    .Build());
            }
        }
    }

    private StructInitializerNode CheckStructInitializer(StructInitializerSyntax expression, NubType? expectedType)
    {
        NubStructType? structType = null;

        if (expression.StructType != null)
        {
            var checkedType = ResolveType(expression.StructType);
            if (checkedType is not NubStructType checkedStructType)
            {
                throw new UnreachableException("Parser fucked up");
            }

            structType = checkedStructType;
        }
        else if (expectedType is NubStructType expectedStructType)
        {
            structType = expectedStructType;
        }

        if (structType == null)
        {
            throw new TypeCheckerException(Diagnostic
                .Error("Cannot get implicit type of struct")
                .WithHelp("Specify struct type with struct {type_name} syntax")
                .At(expression)
                .Build());
        }

        var initializers = new Dictionary<string, ExpressionNode>();

        foreach (var initializer in expression.Initializers)
        {
            var typeField = structType.Fields.FirstOrDefault(x => x.Name == initializer.Key);
            if (typeField == null)
            {
                Diagnostics.AddRange(Diagnostic
                    .Error($"Struct {structType.Name} does not have a field named {initializer.Key}")
                    .At(initializer.Value)
                    .Build());

                continue;
            }

            initializers.Add(initializer.Key, CheckExpression(initializer.Value, typeField.Type));
        }

        var missingFields = structType.Fields
            .Where(x => !x.HasDefaultValue && !initializers.ContainsKey(x.Name))
            .Select(x => x.Name)
            .ToArray();

        if (missingFields.Length != 0)
        {
            Diagnostics.Add(Diagnostic
                .Warning($"Fields {string.Join(", ", missingFields)} are not initialized")
                .At(expression)
                .Build());
        }

        return new StructInitializerNode(structType, initializers);
    }

    private BlockNode CheckBlock(BlockSyntax node)
    {
        var statements = new List<StatementNode>();

        var reachable = true;
        var warnedUnreachable = false;

        using (BeginScope())
        {
            foreach (var statement in node.Statements)
            {
                var checkedStatement = CheckStatement(statement);

                if (reachable)
                {
                    if (checkedStatement is TerminalStatementNode)
                    {
                        reachable = false;
                    }

                    statements.Add(checkedStatement);
                }
                else
                {
                    if (!warnedUnreachable)
                    {
                        Diagnostics.Add(Diagnostic.Warning("Statement is unreachable").At(statement).Build());
                        warnedUnreachable = true;
                    }
                }
            }

            return new BlockNode(statements);
        }
    }

    private StatementNode CheckStatement(StatementSyntax statement)
    {
        return statement switch
        {
            AssignmentSyntax assignmentStmt => CheckAssignment(assignmentStmt),
            BlockSyntax blockStmt => CheckBlock(blockStmt),
            BreakSyntax => new BreakNode(),
            ContinueSyntax => new ContinueNode(),
            IfSyntax ifStmt => CheckIf(ifStmt),
            ReturnSyntax retStmt => CheckReturn(retStmt),
            StatementExpressionSyntax stmtExpr => CheckStatementExpression(stmtExpr),
            VariableDeclarationSyntax varDeclStmt => CheckVariableDeclaration(varDeclStmt),
            WhileSyntax whileStmt => CheckWhile(whileStmt),
            DeferSyntax defer => new DeferNode(CheckStatement(defer.Statement)),
            _ => throw new ArgumentOutOfRangeException(nameof(statement))
        };
    }

    private bool AlwaysReturns(StatementNode statement)
    {
        switch (statement)
        {
            case ReturnNode:
                return true;
            case BlockNode block:
                return block.Statements.Count != 0 && AlwaysReturns(block.Statements.Last());
            case IfNode ifNode:
            {
                if (!AlwaysReturns(ifNode.Body))
                {
                    return false;
                }

                return !ifNode.Else.HasValue || ifNode.Else.Value.Match(AlwaysReturns, AlwaysReturns);
            }
            default:
                return false;
        }
    }

    private NubType ResolveType(TypeSyntax type)
    {
        return type switch
        {
            BoolTypeSyntax => new NubBoolType(),
            CStringTypeSyntax => new NubCStringType(),
            IntTypeSyntax i => new NubIntType(i.Signed, i.Width),
            FloatTypeSyntax f => new NubFloatType(f.Width),
            FuncTypeSyntax func => new NubFuncType(func.Parameters.Select(ResolveType).ToList(), ResolveType(func.ReturnType)),
            ArrayTypeSyntax arr => new NubArrayType(ResolveType(arr.BaseType)),
            ConstArrayTypeSyntax arr => new NubConstArrayType(ResolveType(arr.BaseType), arr.Size),
            PointerTypeSyntax ptr => new NubPointerType(ResolveType(ptr.BaseType)),
            StringTypeSyntax => new NubStringType(),
            CustomTypeSyntax c => ResolveCustomType(c),
            StructTemplateTypeSyntax t => ResolveStructTemplateType(t),
            SubstitutionTypeSyntax s => ResolveTypeSubstitution(s),
            VoidTypeSyntax => new NubVoidType(),
            _ => throw new NotSupportedException($"Unknown type syntax: {type}")
        };
    }

    private NubType ResolveTypeSubstitution(SubstitutionTypeSyntax substitution)
    {
        var type = CurrentScope.LookupTypeSubstitution(substitution.Name);
        if (type == null)
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Template argument {substitution.Name} does not exist")
                .At(substitution)
                .Build());
        }

        return type;
    }

    private NubType ResolveCustomType(CustomTypeSyntax customType)
    {
        var key = (customType.Module, customType.Name);

        if (_typeCache.TryGetValue(key, out var cachedType))
        {
            return cachedType;
        }

        if (!_resolvingTypes.Add(key))
        {
            var placeholder = new NubStructType(customType.Module, customType.Name, [], []);
            _typeCache[key] = placeholder;
            return placeholder;
        }

        try
        {
            if (!_visibleModules.TryGetValue(customType.Module, out var module))
            {
                throw new TypeCheckerException(Diagnostic
                    .Error($"Module {customType.Module} not found")
                    .WithHelp($"import \"{customType.Module}\"")
                    .At(customType)
                    .Build());
            }

            var includePrivate = customType.Module == CurrentScope.Module;

            var structDef = module.Structs(includePrivate).FirstOrDefault(x => x.Name == customType.Name);
            if (structDef != null)
            {
                var result = new NubStructType(customType.Module, structDef.Name, [], []);
                _typeCache[key] = result;

                var fields = structDef.Fields
                    .Select(x => new NubStructFieldType(x.Name, ResolveType(x.Type), x.Value != null))
                    .ToList();

                result.Fields.AddRange(fields);

                var functions = structDef.Functions
                    .Select(x =>
                    {
                        var parameters = x.Signature.Parameters
                            .Select(y => ResolveType(y.Type))
                            .ToList();

                        return new NubStructFuncType(x.Name, parameters, ResolveType(x.Signature.ReturnType));
                    })
                    .ToList();

                result.Functions.AddRange(functions);

                ReferencedStructTypes.Add(result);
                return result;
            }

            throw new TypeCheckerException(Diagnostic
                .Error($"Type {customType.Name} not found in module {customType.Module}")
                .At(customType)
                .Build());
        }
        finally
        {
            _resolvingTypes.Remove(key);
        }
    }

    private NubStructType ResolveStructTemplateType(StructTemplateTypeSyntax structTemplate)
    {
        if (!_visibleModules.TryGetValue(structTemplate.Module, out var module))
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Module {structTemplate.Module} not found")
                .WithHelp($"import \"{structTemplate.Module}\"")
                .At(structTemplate)
                .Build());
        }

        var includePrivate = structTemplate.Module == CurrentScope.Module;

        var templateDef = module
            .StructTemplates(includePrivate)
            .FirstOrDefault(x => x.Name == structTemplate.Name);

        if (templateDef == null)
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Template type {structTemplate.Name} not found in module {structTemplate.Module}")
                .At(structTemplate)
                .Build());
        }

        var templateParameterTypes = structTemplate.TemplateParameters.Select(ResolveType).ToList();
        var mangledName = $"{structTemplate.Name}.{NameMangler.Mangle(templateParameterTypes)}";

        using (BeginRootScope(structTemplate.Module))
        {
            for (var i = 0; i < templateParameterTypes.Count; i++)
            {
                var parameterName = templateDef.TemplateArguments[i];
                var parameterType = templateParameterTypes[i];
                CurrentScope.DeclareTypeSubstitution(parameterName, parameterType);
            }

            var fieldTypes = templateDef.Fields
                .Select(x => new NubStructFieldType(x.Name, ResolveType(x.Type), x.Value != null))
                .ToList();

            var fieldFunctions = templateDef.Functions
                .Select(x =>
                {
                    var parameters = x.Signature.Parameters.Select(y => ResolveType(y.Type)).ToList();
                    var returnType = ResolveType(x.Signature.ReturnType);
                    return new NubStructFuncType(x.Name, parameters, returnType);
                })
                .ToList();

            var structType = new NubStructType(structTemplate.Module, mangledName, fieldTypes, fieldFunctions);

            if (!_checkedTemplateStructs.Contains($"{structTemplate.Module}.{mangledName}"))
            {
                CurrentScope.DeclareVariable(new Variable("this", structType, VariableKind.RValue));
                var fields = templateDef.Fields.Select(CheckStructField).ToList();
                var functions = templateDef.Functions.Select(CheckStructFunc).ToList();
                Definitions.Add(new StructNode(structTemplate.Module, mangledName, fields, functions));
                _checkedTemplateStructs.Add($"{structTemplate.Module}.{mangledName}");
            }

            ReferencedStructTypes.Add(structType);

            return structType;
        }
    }
}

public enum VariableKind
{
    LValue,
    RValue
}

public record Variable(string Name, NubType Type, VariableKind Kind);

public class Scope(string module, Scope? parent = null)
{
    private readonly List<Variable> _variables = [];
    private readonly Dictionary<string, NubType> _typeSubstitutions = [];
    public string Module { get; } = module;

    public void DeclareVariable(Variable variable)
    {
        _variables.Add(variable);
    }

    public Variable? LookupVariable(string name)
    {
        var variable = _variables.FirstOrDefault(x => x.Name == name);
        if (variable != null)
        {
            return variable;
        }

        return parent?.LookupVariable(name);
    }

    public void DeclareTypeSubstitution(string name, NubType type)
    {
        _typeSubstitutions[name] = type;
    }

    public NubType? LookupTypeSubstitution(string name)
    {
        if (_typeSubstitutions.TryGetValue(name, out var type))
        {
            return type;
        }

        return parent?.LookupTypeSubstitution(name);
    }

    public Scope SubScope()
    {
        return new Scope(Module, this);
    }
}

public class TypeCheckerException : Exception
{
    public Diagnostic Diagnostic { get; }

    public TypeCheckerException(Diagnostic diagnostic) : base(diagnostic.Message)
    {
        Diagnostic = diagnostic;
    }
}