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

    private readonly Stack<Scope> _scopes = [];

    private readonly TypeResolver _typeResolver;

    private Scope Scope => _scopes.Peek();

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

    public TypeChecker(SyntaxTree syntaxTree, Dictionary<string, Module> modules)
    {
        _syntaxTree = syntaxTree;
        _modules = modules;
        _typeResolver = new TypeResolver(_modules);
    }

    public List<TopLevelNode> Check()
    {
        _scopes.Clear();

        var moduleDeclarations = _syntaxTree.TopLevelSyntaxNodes.OfType<ModuleSyntax>().ToList();
        if (moduleDeclarations.Count == 0)
        {
            Diagnostics.Add(Diagnostic.Error("Missing module declaration").WithHelp("module \"main\"").Build());
            return [];
        }

        if (moduleDeclarations.Count > 1)
        {
            Diagnostics.Add(Diagnostic.Error("Multiple module declarations").WithHelp("Remove extra module declarations").Build());
        }

        var moduleName = moduleDeclarations[0].NameToken;

        var importDeclarations = _syntaxTree.TopLevelSyntaxNodes.OfType<ImportSyntax>().ToList();
        foreach (var importDeclaration in importDeclarations)
        {
            var name = importDeclaration.NameToken.Value;

            var last = importDeclarations.Last(x => x.NameToken.Value == name);
            if (importDeclaration != last)
            {
                Diagnostics.Add(Diagnostic
                    .Warning($"Module \"{last.NameToken.Value}\" is imported twice")
                    .WithHelp($"Remove duplicate import \"{last.NameToken.Value}\"")
                    .At(last)
                    .Build());
            }

            var exists = _modules.ContainsKey(name);
            if (!exists)
            {
                var suggestions = _modules.Keys
                    .Select(m => new { Name = m, Distance = Utils.LevenshteinDistance(name, m) })
                    .OrderBy(x => x.Distance)
                    .Take(3)
                    .Where(x => x.Distance <= 3)
                    .Select(x => $"\"{x.Name}\"")
                    .ToArray();

                var suggestionText = suggestions.Length != 0
                    ? $"Did you mean {string.Join(", ", suggestions)}?"
                    : "Check that the module name is correct.";

                Diagnostics.Add(Diagnostic
                    .Error($"Module \"{name}\" does not exist")
                    .WithHelp(suggestionText)
                    .At(last)
                    .Build());

                return [];
            }
        }

        var topLevelNodes = new List<TopLevelNode>();

        using (BeginRootScope(moduleName))
        {
            foreach (var topLevelSyntaxNode in _syntaxTree.TopLevelSyntaxNodes)
            {
                switch (topLevelSyntaxNode)
                {
                    case EnumSyntax:
                        break;
                    case FuncSyntax funcSyntax:
                        topLevelNodes.Add(CheckFuncDefinition(funcSyntax));
                        break;
                    case StructSyntax structSyntax:
                        topLevelNodes.Add(CheckStructDefinition(structSyntax));
                        break;
                    case ImportSyntax importSyntax:
                        topLevelNodes.Add(new ImportNode(importSyntax.Tokens, importSyntax.NameToken));
                        break;
                    case ModuleSyntax moduleSyntax:
                        topLevelNodes.Add(new ModuleNode(moduleSyntax.Tokens, moduleSyntax.NameToken));
                        break;
                    default:
                        throw new ArgumentOutOfRangeException(nameof(topLevelSyntaxNode));
                }
            }
        }

        return topLevelNodes;
    }

    private (IdentifierToken Name, Module Module) GetCurrentModule()
    {
        var currentModule = _syntaxTree.TopLevelSyntaxNodes.OfType<ModuleSyntax>().First().NameToken;
        return (currentModule, _modules[currentModule.Value]);
    }

    private List<(IdentifierToken Name, Module Module)> GetImportedModules()
    {
        var currentModule = _syntaxTree.TopLevelSyntaxNodes.OfType<ModuleSyntax>().First().NameToken;
        return _syntaxTree.TopLevelSyntaxNodes
            .OfType<ImportSyntax>()
            .Select(x => (Name: x.NameToken, Module: _modules[x.NameToken.Value]))
            .Concat([(Name: currentModule, Module: _modules[currentModule.Value])])
            .ToList();
    }

    private bool IsCurrentModule(IdentifierToken? module)
    {
        if (module == null)
        {
            return true;
        }

        return module.Value == Scope.Module.Value;
    }

    private Module? GetImportedModule(string module)
    {
        var currentModule = _syntaxTree.TopLevelSyntaxNodes.OfType<ModuleSyntax>().First().NameToken;
        if (module == currentModule.Value)
        {
            return _modules[currentModule.Value];
        }

        var import = _syntaxTree.TopLevelSyntaxNodes.OfType<ImportSyntax>().FirstOrDefault(x => x.NameToken.Value == module);
        if (import != null)
        {
            return _modules[import.NameToken.Value];
        }

        return null;
    }

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

    private ScopeDisposer BeginRootScope(IdentifierToken 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.NameToken, parameter.Type));
            }

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

    private StructNode CheckStructDefinition(StructSyntax structSyntax)
    {
        var fields = new List<StructFieldNode>();

        foreach (var field in structSyntax.Fields)
        {
            var fieldType = _typeResolver.ResolveType(field.Type, Scope.Module.Value);
            ExpressionNode? value = null;
            if (field.Value != null)
            {
                value = CheckExpression(field.Value, fieldType);
                if (value.Type != fieldType)
                {
                    throw new CompileException(Diagnostic
                        .Error($"Type {value.Type} is not assignable to {field.Type} for field {field.NameToken.Value}")
                        .At(field)
                        .Build());
                }
            }

            fields.Add(new StructFieldNode(field.Tokens, field.NameToken, fieldType, value));
        }

        var currentModule = GetCurrentModule();
        var type = new NubStructType(currentModule.Name.Value, structSyntax.NameToken.Value, structSyntax.Packed, fields.Select(x => new NubStructFieldType(x.NameToken.Value, x.Type, x.Value != null)).ToList());

        return new StructNode(structSyntax.Tokens, structSyntax.NameToken, type, structSyntax.Packed, fields);
    }

    private AssignmentNode CheckAssignment(AssignmentSyntax statement)
    {
        var target = CheckExpression(statement.Target);
        var value = CheckExpression(statement.Value, target.Type);

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

        return new AssignmentNode(statement.Tokens, target, 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 CompileException(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 = _typeResolver.ResolveType(statement.ExplicitType, Scope.Module.Value);
        }

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

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

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

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

        return new VariableDeclarationNode(statement.Tokens, statement.NameToken, 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.ElementNameToken, sliceType.ElementType));
                    if (forSyntax.IndexNameToken != null)
                    {
                        Scope.DeclareVariable(new Variable(forSyntax.IndexNameToken, new NubIntType(false, 64)));
                    }

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

                    var body = CheckBlock(forSyntax.Body);
                    return new ForConstArrayNode(forSyntax.Tokens, forSyntax.ElementNameToken, forSyntax.IndexNameToken, target, body);
                }
            }
            default:
            {
                throw new CompileException(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.NameToken, _typeResolver.ResolveType(parameter.Type, Scope.Module.Value)));
        }

        return new FuncPrototypeNode(statement.Tokens, statement.NameToken, statement.ExternSymbolToken, parameters, _typeResolver.ResolveType(statement.ReturnType, Scope.Module.Value));
    }

    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),
            SizeSyntax expression => new SizeNode(node.Tokens, _typeResolver.ResolveType(expression.Type, Scope.Module.Value)),
            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 CompileException(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 CompileException(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:
            {
                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:
                {
                    return true;
                }
            }
        }


        return false;
    }

    private AddressOfNode CheckAddressOf(AddressOfSyntax expression, NubType? expectedType)
    {
        var target = CheckExpression(expression.Target, (expectedType as NubPointerType)?.BaseType);
        var type = new NubPointerType(target.Type);
        return new AddressOfNode(expression.Tokens, type, target);
    }

    private ExpressionNode CheckArrayIndexAccess(ArrayIndexAccessSyntax expression, NubType? _)
    {
        var index = CheckExpression(expression.Index);
        if (index.Type is not NubIntType)
        {
            throw new CompileException(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 CompileException(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 CompileException(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 CompileException(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, (ulong)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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(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 CompileException(Diagnostic.Error($"Cannot call non-function type {accessor.Type}").At(expression.Expression).Build());
        }

        if (expression.Parameters.Count != funcType.Parameters.Count)
        {
            throw new CompileException(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 CompileException(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 CheckLocalIdentifier(LocalIdentifierSyntax expression, NubType? _)
    {
        // note(nub31): Local identifiers can be variables or a symbol in a module
        var scopeIdent = Scope.LookupVariable(expression.NameToken);
        if (scopeIdent != null)
        {
            return new VariableIdentifierNode(expression.Tokens, scopeIdent.Type, expression.NameToken);
        }

        var module = GetImportedModule(Scope.Module.Value)!;

        var function = module.Functions(true).FirstOrDefault(x => x.NameToken.Value == expression.NameToken.Value);
        if (function != null)
        {
            var parameters = function.Prototype.Parameters.Select(x => _typeResolver.ResolveType(x.Type, Scope.Module.Value)).ToList();
            var type = new NubFuncType(parameters, _typeResolver.ResolveType(function.Prototype.ReturnType, Scope.Module.Value));
            return new FuncIdentifierNode(expression.Tokens, type, Scope.Module, expression.NameToken, function.Prototype.ExternSymbolToken);
        }

        var enumDef = module.Enums(true).FirstOrDefault(x => x.NameToken.Value == expression.NameToken.Value);
        if (enumDef != null)
        {
            return new EnumReferenceIntermediateNode(expression.Tokens, Scope.Module, expression.NameToken);
        }

        throw new CompileException(Diagnostic
            .Error($"There is no identifier named {expression.NameToken.Value}")
            .At(expression)
            .Build());
    }

    private ExpressionNode CheckModuleIdentifier(ModuleIdentifierSyntax expression, NubType? _)
    {
        var module = GetImportedModule(expression.ModuleToken.Value);
        if (module == null)
        {
            throw new CompileException(Diagnostic
                .Error($"Module {expression.ModuleToken.Value} not found")
                .WithHelp($"import \"{expression.ModuleToken.Value}\"")
                .At(expression.ModuleToken)
                .Build());
        }

        var function = module.Functions(false).FirstOrDefault(x => x.NameToken.Value == expression.NameToken.Value);
        if (function != null)
        {
            using (BeginRootScope(expression.ModuleToken))
            {
                var parameters = function.Prototype.Parameters.Select(x => _typeResolver.ResolveType(x.Type, Scope.Module.Value)).ToList();
                var type = new NubFuncType(parameters, _typeResolver.ResolveType(function.Prototype.ReturnType, Scope.Module.Value));
                return new FuncIdentifierNode(expression.Tokens, type, expression.ModuleToken, expression.NameToken, function.Prototype.ExternSymbolToken);
            }
        }

        var enumDef = module.Enums(false).FirstOrDefault(x => x.NameToken.Value == expression.NameToken.Value);
        if (enumDef != null)
        {
            return new EnumReferenceIntermediateNode(expression.Tokens, expression.ModuleToken, expression.NameToken);
        }

        throw new CompileException(Diagnostic
            .Error($"Module {expression.ModuleToken.Value} does not export a member named {expression.NameToken.Value}")
            .At(expression)
            .Build());
    }

    private ExpressionNode CheckStringLiteral(StringLiteralSyntax expression, NubType? expectedType)
    {
        if (expectedType is NubPointerType { BaseType: NubIntType { Signed: true, Width: 8 } })
        {
            return new CStringLiteralNode(expression.Tokens, expression.Token.Value);
        }

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

    private ExpressionNode CheckIntLiteral(IntLiteralSyntax expression, NubType? expectedType)
    {
        if (expectedType is NubIntType intType)
        {
            return intType.Width switch
            {
                8 => intType.Signed ? new I8LiteralNode(expression.Tokens, expression.Token.AsI8) : new U8LiteralNode(expression.Tokens, expression.Token.AsU8),
                16 => intType.Signed ? new I16LiteralNode(expression.Tokens, expression.Token.AsI16) : new U16LiteralNode(expression.Tokens, expression.Token.AsU16),
                32 => intType.Signed ? new I32LiteralNode(expression.Tokens, expression.Token.AsI32) : new U32LiteralNode(expression.Tokens, expression.Token.AsU32),
                64 => intType.Signed ? new I64LiteralNode(expression.Tokens, expression.Token.AsI64) : new U64LiteralNode(expression.Tokens, expression.Token.AsU64),
                _ => throw new ArgumentOutOfRangeException()
            };
        }

        if (expectedType is NubFloatType floatType)
        {
            return floatType.Width switch
            {
                32 => new Float32LiteralNode(expression.Tokens, expression.Token.AsF32),
                64 => new Float64LiteralNode(expression.Tokens, expression.Token.AsF64),
                _ => throw new ArgumentOutOfRangeException()
            };
        }

        return new I64LiteralNode(expression.Tokens, expression.Token.AsI64);
    }

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

        return new Float64LiteralNode(expression.Tokens, expression.Token.AsF64);
    }

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

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

        if (target is EnumReferenceIntermediateNode enumReferenceIntermediate)
        {
            var enumDef = GetImportedModules()
                .First(x => x.Name.Value == enumReferenceIntermediate.ModuleToken.Value)
                .Module
                .Enums(IsCurrentModule(enumReferenceIntermediate.ModuleToken))
                .First(x => x.NameToken.Value == enumReferenceIntermediate.NameToken.Value);

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

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

            if (enumIntType.Signed)
            {
                var fieldValue = CalculateSignedEnumFieldValue(enumDef, field);
                return enumIntType.Width switch
                {
                    8 => new I8LiteralNode(expression.Tokens, (sbyte)fieldValue),
                    16 => new I16LiteralNode(expression.Tokens, (short)fieldValue),
                    32 => new I32LiteralNode(expression.Tokens, (int)fieldValue),
                    64 => new I64LiteralNode(expression.Tokens, fieldValue),
                    _ => throw new ArgumentOutOfRangeException()
                };
            }
            else
            {
                var fieldValue = CalculateUnsignedEnumFieldValue(enumDef, field);
                return enumIntType.Width switch
                {
                    8 => new U8LiteralNode(expression.Tokens, (byte)fieldValue),
                    16 => new U16LiteralNode(expression.Tokens, (ushort)fieldValue),
                    32 => new U32LiteralNode(expression.Tokens, (uint)fieldValue),
                    64 => new U64LiteralNode(expression.Tokens, fieldValue),
                    _ => throw new ArgumentOutOfRangeException()
                };
            }
        }

        switch (target.Type)
        {
            case NubStructType structType:
            {
                var field = structType.Fields.FirstOrDefault(x => x.Name == expression.MemberToken.Value);
                if (field == null)
                {
                    throw new CompileException(Diagnostic
                        .Error($"Struct {target.Type} does not have a field with the name {expression.MemberToken.Value}")
                        .At(expression)
                        .Build());
                }

                return new StructFieldAccessNode(expression.Tokens, field.Type, target, expression.MemberToken);
            }
            default:
            {
                throw new CompileException(Diagnostic
                    .Error($"Cannot access struct member {expression.MemberToken.Value} on type {target.Type}")
                    .At(expression)
                    .Build());
            }
        }
    }

    private static long CalculateSignedEnumFieldValue(EnumSyntax enumDef, EnumFieldSyntax field)
    {
        long currentValue = 0;

        foreach (var f in enumDef.Fields)
        {
            if (f.ValueToken != null)
            {
                currentValue = f.ValueToken.AsI64;
            }

            if (f == field)
            {
                return currentValue;
            }

            currentValue++;
        }

        throw new UnreachableException();
    }

    private static ulong CalculateUnsignedEnumFieldValue(EnumSyntax enumDef, EnumFieldSyntax field)
    {
        ulong currentValue = 0;

        foreach (var f in enumDef.Fields)
        {
            if (f.ValueToken != null)
            {
                currentValue = f.ValueToken.AsU64;
            }

            if (f == field)
            {
                return currentValue;
            }

            currentValue++;
        }

        throw new UnreachableException();
    }

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

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

            structType = checkedStructType;
        }

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

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

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

        foreach (var initializer in expression.Initializers)
        {
            var typeField = structType.Fields.FirstOrDefault(x => x.Name == initializer.Key.Value);
            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.All(y => y.Key.Value != 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 (CompileException 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))
        };
    }
}

public record Variable(IdentifierToken Name, NubType Type);

public class Scope(IdentifierToken module, Scope? parent = null)
{
    private NubType? _returnType;
    private readonly List<Variable> _variables = [];
    public IdentifierToken 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(IdentifierToken name)
    {
        var variable = _variables.FirstOrDefault(x => x.Name.Value == name.Value);
        if (variable != null)
        {
            return variable;
        }

        return parent?.LookupVariable(name);
    }

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