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

namespace NubLang.Ast;

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

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

    private Scope Scope => _scopes.Peek();

    public List<Diagnostic> Diagnostics { get; } = [];

    public TypeChecker(SyntaxTree syntaxTree, Dictionary<string, Module> modules)
    {
        _syntaxTree = syntaxTree;
        _importedModules = modules
            .Where(x => syntaxTree.Imports.Contains(x.Key) || _syntaxTree.ModuleName == x.Key)
            .ToDictionary();
    }

    public CompilationUnit Check()
    {
        _scopes.Clear();
        _typeCache.Clear();
        _resolvingTypes.Clear();

        var functions = new List<FuncNode>();

        using (BeginRootScope(_syntaxTree.ModuleName))
        {
            foreach (var funcSyntax in _syntaxTree.Definitions.OfType<FuncSyntax>())
            {
                try
                {
                    functions.Add(CheckFuncDefinition(funcSyntax));
                }
                catch (TypeCheckerException e)
                {
                    Diagnostics.Add(e.Diagnostic);
                }
            }
        }

        var importedStructTypes = new List<NubStructType>();
        var importedFunctions = new List<FuncPrototypeNode>();

        foreach (var (name, module) in _importedModules)
        {
            using (BeginRootScope(name))
            {
                foreach (var structSyntax in module.Structs(true))
                {
                    try
                    {
                        var fields = structSyntax.Fields
                            .Select(f => new NubStructFieldType(f.Name, ResolveType(f.Type), f.Value != null))
                            .ToList();

                        importedStructTypes.Add(new NubStructType(name, structSyntax.Name, fields));
                    }
                    catch (TypeCheckerException e)
                    {
                        Diagnostics.Add(e.Diagnostic);
                    }
                }

                foreach (var funcSyntax in module.Functions(true))
                {
                    try
                    {
                        importedFunctions.Add(CheckFuncPrototype(funcSyntax.Prototype));
                    }
                    catch (TypeCheckerException e)
                    {
                        Diagnostics.Add(e.Diagnostic);
                    }
                }
            }
        }

        return new CompilationUnit(functions, importedStructTypes, importedFunctions);
    }

    private ScopeDisposer BeginScope()
    {
        _scopes.Push(Scope.SubScope());
        return new ScopeDisposer(this);
    }

    private ScopeDisposer BeginRootScope(string moduleName)
    {
        _scopes.Push(new Scope(moduleName));
        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 FuncNode CheckFuncDefinition(FuncSyntax node)
    {
        using (BeginScope())
        {
            var prototype = CheckFuncPrototype(node.Prototype);

            Scope.SetReturnType(prototype.ReturnType);
            foreach (var parameter in prototype.Parameters)
            {
                Scope.DeclareVariable(new Variable(parameter.Name, parameter.Type));
            }

            var body = node.Body == null ? null : CheckBlock(node.Body);
            return new FuncNode(node.Tokens, prototype, body);
        }
    }

    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);

        if (value.Type != lValue.Type)
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Cannot assign {value.Type} to {lValue.Type}")
                .At(statement.Value)
                .Build());
        }

        return new AssignmentNode(statement.Tokens, lValue, value);
    }

    private IfNode CheckIf(IfSyntax statement)
    {
        var condition = CheckExpression(statement.Condition);
        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(statement.Tokens, condition, body, elseStatement);
    }

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

        if (statement.Value != null)
        {
            var expectedReturnType = Scope.GetReturnType();
            value = CheckExpression(statement.Value, expectedReturnType);
        }

        return new ReturnNode(statement.Tokens, value);
    }

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

        return expression switch
        {
            FuncCallNode funcCall => new StatementFuncCallNode(statement.Tokens, funcCall),
            _ => 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);

            if (type == null)
            {
                type = assignmentNode.Type;
            }
            else if (assignmentNode.Type != type)
            {
                throw new TypeCheckerException(Diagnostic
                    .Error($"Cannot assign {assignmentNode.Type} to variable of type {type}")
                    .At(statement.Assignment)
                    .Build());
            }
        }

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

        Scope.DeclareVariable(new Variable(statement.Name, type));

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

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

    private StatementNode CheckFor(ForSyntax forSyntax)
    {
        var target = CheckExpression(forSyntax.Target);


        switch (target.Type)
        {
            case NubSliceType sliceType:
            {
                using (BeginScope())
                {
                    Scope.DeclareVariable(new Variable(forSyntax.ElementName, sliceType.ElementType));
                    if (forSyntax.IndexName != null)
                    {
                        Scope.DeclareVariable(new Variable(forSyntax.IndexName, new NubIntType(false, 64)));
                    }

                    var body = CheckBlock(forSyntax.Body);
                    return new ForSliceNode(forSyntax.Tokens, forSyntax.ElementName, forSyntax.IndexName, target, body);
                }
            }
            case NubConstArrayType constArrayType:
            {
                using (BeginScope())
                {
                    Scope.DeclareVariable(new Variable(forSyntax.ElementName, constArrayType.ElementType));
                    if (forSyntax.IndexName != null)
                    {
                        Scope.DeclareVariable(new Variable(forSyntax.IndexName, new NubIntType(false, 64)));
                    }

                    var body = CheckBlock(forSyntax.Body);
                    return new ForConstArrayNode(forSyntax.Tokens, forSyntax.ElementName, forSyntax.IndexName, target, body);
                }
            }
            default:
            {
                throw new TypeCheckerException(Diagnostic
                    .Error($"Cannot iterate over type {target.Type} which does not have size information")
                    .At(forSyntax.Target)
                    .Build());
            }
        }
    }

    private FuncPrototypeNode CheckFuncPrototype(FuncPrototypeSyntax statement)
    {
        var parameters = new List<FuncParameterNode>();
        foreach (var parameter in statement.Parameters)
        {
            parameters.Add(new FuncParameterNode(parameter.Tokens, parameter.Name, ResolveType(parameter.Type)));
        }

        return new FuncPrototypeNode(statement.Tokens, Scope.Module, statement.Name, statement.ExternSymbol, parameters, ResolveType(statement.ReturnType));
    }

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

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

            if (IsCastAllowed(result.Type, expectedType))
            {
                return new CastNode(result.Tokens, expectedType, result);
            }
        }

        return result;
    }

    private ExpressionNode CheckCast(CastSyntax expression, NubType? expectedType)
    {
        if (expectedType == null)
        {
            throw new TypeCheckerException(Diagnostic
                .Error("Unable to infer target type of cast")
                .At(expression)
                .WithHelp("Specify target type where value is used")
                .Build());
        }

        var value = CheckExpression(expression.Value, expectedType);

        if (value.Type == expectedType)
        {
            Diagnostics.Add(Diagnostic
                .Warning("Target type of cast is same as the value. Cast is unnecessary")
                .At(expression)
                .Build());

            return value;
        }

        if (!IsCastAllowed(value.Type, expectedType, false))
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Cannot cast from {value.Type} to {expectedType}")
                .Build());
        }

        return new CastNode(expression.Tokens, expectedType, value);
    }

    private static bool IsCastAllowed(NubType from, NubType to, bool strict = true)
    {
        // note(nub31): Implicit casts
        switch (from)
        {
            case NubIntType fromInt when to is NubIntType toInt && fromInt.Width < toInt.Width:
            case NubPointerType when to is NubPointerType { BaseType: NubVoidType }:
            case NubConstArrayType constArrayType1 when to is NubArrayType arrayType && constArrayType1.ElementType == arrayType.ElementType:
            case NubConstArrayType constArrayType3 when to is NubSliceType sliceType2 && constArrayType3.ElementType == sliceType2.ElementType:
            case NubCStringType when to is NubStringType:
            {
                return true;
            }
        }

        if (!strict)
        {
            // note(nub31): Explicit casts
            switch (from)
            {
                case NubIntType when to is NubIntType:
                case NubIntType when to is NubFloatType:
                case NubFloatType when to is NubIntType:
                case NubFloatType when to is NubFloatType:
                case NubPointerType when to is NubPointerType:
                case NubPointerType when to is NubIntType:
                case NubIntType when to is NubPointerType:
                case NubCStringType when to is NubPointerType { BaseType: NubIntType { Width: 8 } }:
                {
                    return true;
                }
            }
        }


        return false;
    }

    private AddressOfNode CheckAddressOf(AddressOfSyntax expression, NubType? expectedType)
    {
        var target = CheckExpression(expression.Target, (expectedType as NubPointerType)?.BaseType);
        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(expression.Tokens, type, lvalue);
    }

    private ExpressionNode CheckArrayIndexAccess(ArrayIndexAccessSyntax expression, NubType? _)
    {
        var index = CheckExpression(expression.Index);
        if (index.Type is not NubIntType)
        {
            throw new TypeCheckerException(Diagnostic
                .Error("Array indexer must be of type int")
                .At(expression.Index)
                .Build());
        }

        var target = CheckExpression(expression.Target);

        return target.Type switch
        {
            NubArrayType arrayType => new ArrayIndexAccessNode(expression.Tokens, arrayType.ElementType, target, index),
            NubConstArrayType constArrayType => new ConstArrayIndexAccessNode(expression.Tokens, constArrayType.ElementType, target, index),
            NubSliceType sliceType => new SliceIndexAccessNode(expression.Tokens, sliceType.ElementType, target, index),
            _ => throw new TypeCheckerException(Diagnostic.Error($"Cannot use array indexer on type {target.Type}").At(expression).Build())
        };
    }

    private ExpressionNode CheckArrayInitializer(ArrayInitializerSyntax expression, NubType? expectedType)
    {
        var elementType = expectedType switch
        {
            NubArrayType arrayType => arrayType.ElementType,
            NubConstArrayType constArrayType => constArrayType.ElementType,
            NubSliceType sliceType => sliceType.ElementType,
            _ => null
        };

        if (elementType == null)
        {
            var firstValue = expression.Values.FirstOrDefault();
            if (firstValue != null)
            {
                elementType = CheckExpression(firstValue).Type;
            }
        }

        if (elementType == null)
        {
            throw new TypeCheckerException(Diagnostic
                .Error("Unable to infer type of array initializer")
                .At(expression)
                .WithHelp("Provide a type for a variable assignment")
                .Build());
        }

        var values = new List<ExpressionNode>();
        foreach (var valueExpression in expression.Values)
        {
            var value = CheckExpression(valueExpression, elementType);
            if (value.Type != elementType)
            {
                throw new TypeCheckerException(Diagnostic
                    .Error("Value in array initializer is not the same as the array type")
                    .At(valueExpression)
                    .Build());
            }

            values.Add(value);
        }

        return expectedType switch
        {
            NubArrayType => new ArrayInitializerNode(expression.Tokens, new NubArrayType(elementType), values),
            NubConstArrayType constArrayType => new ConstArrayInitializerNode(expression.Tokens, constArrayType, values),
            _ => new ConstArrayInitializerNode(expression.Tokens, new NubConstArrayType(elementType, expression.Values.Count), values)
        };
    }

    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:
            {
                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);
                if (right.Type != left.Type)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error($"Expected type {left.Type} from left side of binary expression, but got {right.Type}")
                        .At(expression.Right)
                        .Build());
                }

                return new BinaryExpressionNode(expression.Tokens, new NubBoolType(), left, op, right);
            }
            case BinaryOperatorSyntax.GreaterThan:
            case BinaryOperatorSyntax.GreaterThanOrEqual:
            case BinaryOperatorSyntax.LessThan:
            case BinaryOperatorSyntax.LessThanOrEqual:
            {
                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);
                if (right.Type != left.Type)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error($"Expected type {left.Type} from left side of binary expression, but got {right.Type}")
                        .At(expression.Right)
                        .Build());
                }

                return new BinaryExpressionNode(expression.Tokens, new NubBoolType(), left, op, right);
            }
            case BinaryOperatorSyntax.LogicalAnd:
            case BinaryOperatorSyntax.LogicalOr:
            {
                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);
                if (right.Type != left.Type)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error($"Expected type {left.Type} from left side of binary expression, but got {right.Type}")
                        .At(expression.Right)
                        .Build());
                }

                return new BinaryExpressionNode(expression.Tokens, new NubBoolType(), left, op, right);
            }
            case BinaryOperatorSyntax.Plus:
            {
                var left = CheckExpression(expression.Left, expectedType);
                if (left.Type is not NubIntType and not NubFloatType)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error("The plus operator must only be used with int and float types")
                        .At(expression.Left)
                        .Build());
                }

                var right = CheckExpression(expression.Right, left.Type);
                if (right.Type != left.Type)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error($"Expected type {left.Type} from left side of binary expression, but got {right.Type}")
                        .At(expression.Right)
                        .Build());
                }

                return new BinaryExpressionNode(expression.Tokens, left.Type, left, op, right);
            }
            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);
                if (right.Type != left.Type)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error($"Expected type {left.Type} from left side of binary expression, but got {right.Type}")
                        .At(expression.Right)
                        .Build());
                }

                return new BinaryExpressionNode(expression.Tokens, 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);
                if (right.Type != left.Type)
                {
                    throw new TypeCheckerException(Diagnostic
                        .Error($"Expected type {left.Type} from left side of binary expression, but got {right.Type}")
                        .At(expression.Right)
                        .Build());
                }

                return new BinaryExpressionNode(expression.Tokens, 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(expression.Tokens, 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(expression.Tokens, operand.Type, UnaryOperator.Invert, operand);
            }
            default:
            {
                throw new ArgumentOutOfRangeException();
            }
        }
    }

    private DereferenceNode CheckDereference(DereferenceSyntax expression, NubType? _)
    {
        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(expression.Tokens, pointerType.BaseType, target);
    }

    private FuncCallNode CheckFuncCall(FuncCallSyntax expression, NubType? _)
    {
        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 expectedParameterType = funcType.Parameters[i];

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

            parameters.Add(parameterExpression);
        }

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

    private ExpressionNode? CheckIdentifier(ExpressionSyntax expression, string moduleName, string name)
    {
        if (!_importedModules.TryGetValue(moduleName, out var module))
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Module {moduleName} not found")
                .WithHelp($"import \"{moduleName}\"")
                .At(expression)
                .Build());
        }

        var function = module.Functions(IsCurretModule(moduleName)).FirstOrDefault(x => x.Name == name);
        if (function != null)
        {
            using (BeginRootScope(moduleName))
            {
                var parameters = function.Prototype.Parameters.Select(x => ResolveType(x.Type)).ToList();
                var type = new NubFuncType(parameters, ResolveType(function.Prototype.ReturnType));
                return new FuncIdentifierNode(expression.Tokens, type, moduleName, name, function.Prototype.ExternSymbol);
            }
        }

        var enumDef = module.Enums(IsCurretModule(moduleName)).FirstOrDefault(x => x.Name == name);
        if (enumDef != null)
        {
            return new EnumReferenceIntermediateNode(expression.Tokens, moduleName, name);
        }

        return null;
    }

    private ExpressionNode CheckLocalIdentifier(LocalIdentifierSyntax expression, NubType? _)
    {
        // note(nub31): Local identifiers can be variables or a symbol in a module
        var scopeIdent = Scope.LookupVariable(expression.Name);
        if (scopeIdent != null)
        {
            return new VariableIdentifierNode(expression.Tokens, scopeIdent.Type, expression.Name);
        }

        var ident = CheckIdentifier(expression, Scope.Module, expression.Name);
        if (ident == null)
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"There is no identifier named {expression.Name}")
                .At(expression)
                .Build());
        }

        return ident;
    }

    private ExpressionNode CheckModuleIdentifier(ModuleIdentifierSyntax expression, NubType? _)
    {
        // note(nub31): Unlike local identifiers, module identifiers does not look for local variables
        var ident = CheckIdentifier(expression, expression.Module, expression.Name);
        if (ident == null)
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Module {expression.Module} does not export a member named {expression.Name}")
                .At(expression)
                .Build());
        }

        return ident;
    }

    private ExpressionNode CheckStringLiteral(StringLiteralSyntax expression, NubType? expectedType)
    {
        if (expectedType is NubCStringType)
        {
            return new CStringLiteralNode(expression.Tokens, expression.Value);
        }

        return new StringLiteralNode(expression.Tokens, expression.Value);
    }

    private ExpressionNode CheckIntLiteral(IntLiteralSyntax expression, NubType? expectedType)
    {
        if (expectedType is NubIntType intType)
        {
            return intType.Width switch
            {
                8 => intType.Signed ? new I8LiteralNode(expression.Tokens, Convert.ToSByte(expression.Value, expression.Base)) : new U8LiteralNode(expression.Tokens, Convert.ToByte(expression.Value, expression.Base)),
                16 => intType.Signed ? new I16LiteralNode(expression.Tokens, Convert.ToInt16(expression.Value, expression.Base)) : new U16LiteralNode(expression.Tokens, Convert.ToUInt16(expression.Value, expression.Base)),
                32 => intType.Signed ? new I32LiteralNode(expression.Tokens, Convert.ToInt32(expression.Value, expression.Base)) : new U32LiteralNode(expression.Tokens, Convert.ToUInt32(expression.Value, expression.Base)),
                64 => intType.Signed ? new I64LiteralNode(expression.Tokens, Convert.ToInt64(expression.Value, expression.Base)) : new U64LiteralNode(expression.Tokens, Convert.ToUInt64(expression.Value, expression.Base)),
                _ => throw new ArgumentOutOfRangeException()
            };
        }

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

        return new I64LiteralNode(expression.Tokens, Convert.ToInt64(expression.Value, expression.Base));
    }

    private ExpressionNode CheckFloatLiteral(FloatLiteralSyntax expression, NubType? expectedType)
    {
        if (expectedType is NubFloatType floatType)
        {
            return floatType.Width switch
            {
                32 => new Float32LiteralNode(expression.Tokens, Convert.ToSingle(expression.Value)),
                64 => new Float64LiteralNode(expression.Tokens, Convert.ToDouble(expression.Value)),
                _ => throw new ArgumentOutOfRangeException()
            };
        }

        return new Float64LiteralNode(expression.Tokens, Convert.ToDouble(expression.Value));
    }

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

    private ExpressionNode CheckMemberAccess(MemberAccessSyntax expression, NubType? _)
    {
        var target = CheckExpression(expression.Target);

        if (target is EnumReferenceIntermediateNode enumReferenceIntermediate)
        {
            var enumDef = _importedModules[enumReferenceIntermediate.Module]
                .Enums(IsCurretModule(enumReferenceIntermediate.Module))
                .First(x => x.Name == enumReferenceIntermediate.Name);

            var field = enumDef.Fields.FirstOrDefault(x => x.Name == expression.Member);
            if (field == null)
            {
                throw new TypeCheckerException(Diagnostic
                    .Error($"Enum {Scope.Module}::{enumReferenceIntermediate.Name} does not have a field named {expression.Member}")
                    .At(enumDef)
                    .Build());
            }

            var enumType = enumDef.Type != null ? ResolveType(enumDef.Type) : new NubIntType(false, 64);
            if (enumType is not NubIntType enumIntType)
            {
                throw new TypeCheckerException(Diagnostic.Error("Enum type must be an int type").At(enumDef.Type).Build());
            }

            return enumIntType.Width switch
            {
                8 => enumIntType.Signed ? new I8LiteralNode(expression.Tokens, (sbyte)field.Value) : new U8LiteralNode(expression.Tokens, (byte)field.Value),
                16 => enumIntType.Signed ? new I16LiteralNode(expression.Tokens, (short)field.Value) : new U16LiteralNode(expression.Tokens, (ushort)field.Value),
                32 => enumIntType.Signed ? new I32LiteralNode(expression.Tokens, (int)field.Value) : new U32LiteralNode(expression.Tokens, (uint)field.Value),
                64 => enumIntType.Signed ? new I64LiteralNode(expression.Tokens, field.Value) : new U64LiteralNode(expression.Tokens, (ulong)field.Value),
                _ => throw new ArgumentOutOfRangeException()
            };
        }

        if (target.Type is 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(expression.Tokens, field.Type, target, expression.Member);
        }

        throw new TypeCheckerException(Diagnostic
            .Error($"Cannot access struct member {expression.Member} on 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(expression.Tokens, structType, initializers);
    }

    private BlockNode CheckBlock(BlockSyntax node)
    {
        using (BeginScope())
        {
            var statements = new List<StatementNode>();
            foreach (var statement in node.Statements)
            {
                try
                {
                    statements.Add(CheckStatement(statement));
                }
                catch (TypeCheckerException e)
                {
                    Diagnostics.Add(e.Diagnostic);
                }
            }

            return new BlockNode(node.Tokens, statements);
        }
    }

    private StatementNode CheckStatement(StatementSyntax statement)
    {
        return statement switch
        {
            AssignmentSyntax assignmentStmt => CheckAssignment(assignmentStmt),
            BlockSyntax blockStmt => CheckBlock(blockStmt),
            BreakSyntax => new BreakNode(statement.Tokens),
            ContinueSyntax => new ContinueNode(statement.Tokens),
            IfSyntax ifStmt => CheckIf(ifStmt),
            ReturnSyntax retStmt => CheckReturn(retStmt),
            StatementExpressionSyntax stmtExpr => CheckStatementExpression(stmtExpr),
            VariableDeclarationSyntax varDeclStmt => CheckVariableDeclaration(varDeclStmt),
            WhileSyntax whileStmt => CheckWhile(whileStmt),
            DeferSyntax defer => new DeferNode(statement.Tokens, CheckStatement(defer.Statement)),
            ForSyntax forSyntax => CheckFor(forSyntax),
            _ => 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
        {
            ArrayTypeSyntax arr => new NubArrayType(ResolveType(arr.BaseType)),
            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)),
            SliceTypeSyntax slice => new NubSliceType(ResolveType(slice.BaseType)),
            ConstArrayTypeSyntax arr => new NubConstArrayType(ResolveType(arr.BaseType), arr.Size),
            PointerTypeSyntax ptr => new NubPointerType(ResolveType(ptr.BaseType)),
            StringTypeSyntax => new NubStringType(),
            CustomTypeSyntax c => ResolveCustomType(c),
            VoidTypeSyntax => new NubVoidType(),
            _ => throw new NotSupportedException($"Unknown type syntax: {type}")
        };
    }

    private NubType ResolveCustomType(CustomTypeSyntax customType)
    {
        if (!_importedModules.TryGetValue(customType.Module ?? Scope.Module, out var module))
        {
            throw new TypeCheckerException(Diagnostic
                .Error($"Module {customType.Module ?? Scope.Module} not found")
                .WithHelp($"import \"{customType.Module ?? Scope.Module}\"")
                .At(customType)
                .Build());
        }

        var enumDef = module.Enums(IsCurretModule(customType.Module)).FirstOrDefault(x => x.Name == customType.Name);
        if (enumDef != null)
        {
            return enumDef.Type != null ? ResolveType(enumDef.Type) : new NubIntType(false, 64);
        }

        var structDef = module.Structs(IsCurretModule(customType.Module)).FirstOrDefault(x => x.Name == customType.Name);
        if (structDef != null)
        {
            var key = (customType.Module ?? Scope.Module, customType.Name);

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

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

            try
            {
                var result = new NubStructType(customType.Module ?? Scope.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);
                return result;
            }
            finally
            {
                _resolvingTypes.Remove(key);
            }
        }

        throw new TypeCheckerException(Diagnostic
            .Error($"Type {customType.Name} not found in module {customType.Module ?? Scope.Module}")
            .At(customType)
            .Build());
    }

    private bool IsCurretModule(string? module)
    {
        if (module == null)
        {
            return true;
        }

        return module == Scope.Module;
    }
}

public record Variable(string Name, NubType Type);

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

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

    public void SetReturnType(NubType returnType)
    {
        _returnType = returnType;
    }

    public NubType? GetReturnType()
    {
        return _returnType ?? parent?.GetReturnType();
    }

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

        return parent?.LookupVariable(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;
    }
}