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 Stack<NubType> _funcReturnTypes = [];
    private readonly Dictionary<(string Module, string Name), NubType> _typeCache = new();
    private readonly HashSet<(string Module, string Name)> _resolvingTypes = [];

    private Scope CurrentScope => _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();
        _funcReturnTypes.Clear();
        _typeCache.Clear();
        _resolvingTypes.Clear();

        using (BeginRootScope(_syntaxTree.ModuleName))
        {
            var functions = new List<FuncNode>();

            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)
            {
                foreach (var structSyntax in module.Structs(false))
                {
                    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));
                }

                foreach (var funcSyntax in module.Functions(false))
                {
                    importedFunctions.Add(CheckFuncPrototype(funcSyntax.Prototype));
                }
            }

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

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

        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)
    {
        foreach (var parameter in node.Prototype.Parameters)
        {
            CurrentScope.DeclareVariable(new Variable(parameter.Name, ResolveType(parameter.Type)));
        }

        var prototype = CheckFuncPrototype(node.Prototype);

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

            body = CheckBlock(node.Body);

            if (!AlwaysReturns(body))
            {
                if (prototype.ReturnType is NubVoidType)
                {
                    body.Statements.Add(new ReturnNode(node.Tokens.Skip(node.Tokens.Count - 1).ToList(), null));
                }
                else
                {
                    Diagnostics.Add(Diagnostic
                        .Error("Not all code paths return a value")
                        .At(node.Body)
                        .Build());
                }
            }

            _funcReturnTypes.Pop();
        }

        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);
        return new AssignmentNode(statement.Tokens, lValue, value);
    }

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

        return new IfNode(statement.Tokens, condition, body, elseStatement);
    }

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

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

        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);
            type ??= assignmentNode.Type;
        }

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

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

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

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

    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, CurrentScope.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),
            ArrayIndexAccessSyntax expression => CheckArrayIndexAccess(expression),
            ArrayInitializerSyntax expression => CheckArrayInitializer(expression),
            BinaryExpressionSyntax expression => CheckBinaryExpression(expression, expectedType),
            UnaryExpressionSyntax expression => CheckUnaryExpression(expression, expectedType),
            DereferenceSyntax expression => CheckDereference(expression),
            FuncCallSyntax expression => CheckFuncCall(expression),
            LocalIdentifierSyntax expression => CheckLocalIdentifier(expression),
            ModuleIdentifierSyntax expression => CheckModuleIdentifier(expression),
            BoolLiteralSyntax expression => CheckBoolLiteral(expression),
            StringLiteralSyntax expression => CheckStringLiteral(expression, expectedType),
            IntLiteralSyntax expression => CheckIntLiteral(expression, expectedType),
            FloatLiteralSyntax expression => CheckFloatLiteral(expression, expectedType),
            MemberAccessSyntax expression => CheckMemberAccess(expression),
            StructInitializerSyntax expression => CheckStructInitializer(expression, expectedType),
            InterpretBuiltinSyntax expression => CheckExpression(expression.Target) with { Type = ResolveType(expression.Type) },
            SizeBuiltinSyntax expression => new SizeBuiltinNode(node.Tokens, new NubIntType(false, 64), ResolveType(expression.Type)),
            FloatToIntBuiltinSyntax expression => CheckFloatToInt(expression),
            _ => throw new ArgumentOutOfRangeException(nameof(node))
        };

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

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

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

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

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

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

        return new FloatToIntBuiltinNode(expression.Tokens, targetIntType, value, sourceFloatType, targetIntType);
    }

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

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

    private ExpressionNode CheckArrayIndexAccess(ArrayIndexAccessSyntax expression)
    {
        var index = CheckExpression(expression.Index, new NubIntType(false, 64));
        var target = CheckExpression(expression.Target);

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

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

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

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

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

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

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

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

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

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

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

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

                return new BinaryExpressionNode(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);

                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)
    {
        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)
    {
        var accessor = CheckExpression(expression.Expression);
        if (accessor.Type is not NubFuncType funcType)
        {
            throw new TypeCheckerException(Diagnostic.Error($"Cannot call non-function type {accessor.Type}").At(expression.Expression).Build());
        }

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

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

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

            parameters.Add(parameterExpression);
        }

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

    private ExpressionNode CheckLocalIdentifier(LocalIdentifierSyntax expression)
    {
        var scopeIdent = CurrentScope.LookupVariable(expression.Name);
        if (scopeIdent != null)
        {
            return new VariableIdentifierNode(expression.Tokens, scopeIdent.Type, expression.Name);
        }

        var module = _importedModules[CurrentScope.Module];
        var function = module.Functions(true).FirstOrDefault(x => x.Name == expression.Name);

        if (function != null)
        {
            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, CurrentScope.Module, expression.Name, function.Prototype.ExternSymbol);
        }

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

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

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

        var function = module.Functions(includePrivate).FirstOrDefault(x => x.Name == expression.Name);
        if (function != null)
        {
            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, expression.Module, expression.Name, function.Prototype.ExternSymbol);
        }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

                continue;
            }

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

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

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

        return new StructInitializerNode(expression.Tokens, structType, initializers);
    }

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

            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)),
            _ => 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)
    {
        var key = (customType.Module, customType.Name);

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

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

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

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

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

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

                result.Fields.AddRange(fields);
                return result;
            }

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

public record Variable(string Name, NubType Type);

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

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

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

        return parent?.LookupVariable(name);
    }

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

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

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