using System.Diagnostics;
using System.Text;
using NubLang.Ast;
using NubLang.Syntax;

namespace NubLang.Generation;

public class Generator
{
    private readonly List<TopLevelNode> _compilationUnit;
    private readonly IndentedTextWriter _writer;
    private readonly Stack<Scope> _scopes = [];
    private int _tmpIndex;

    private Scope Scope => _scopes.Peek();

    public Generator(List<TopLevelNode> compilationUnit)
    {
        _compilationUnit = compilationUnit;
        _writer = new IndentedTextWriter();
    }

    // todo(nub31): Handle name collisions
    private string NewTmp()
    {
        return $"_t{++_tmpIndex}";
    }

    private static string FuncName(string module, string name, string? externSymbol)
    {
        return externSymbol ?? $"{module}_{name}";
    }

    private string GetModuleName()
    {
        return _compilationUnit.OfType<ModuleNode>().First().NameToken.Value;
    }

    public string Emit()
    {
        foreach (var structType in _compilationUnit.OfType<StructNode>())
        {
            _writer.WriteLine($"void {CType.Create(structType.StructType)}_create({CType.Create(structType.StructType)} *self)");
            _writer.WriteLine("{");
            using (_writer.Indent())
            {
                foreach (var field in structType.Fields)
                {
                    if (field.Value != null)
                    {
                        var value = EmitExpression(field.Value);
                        _writer.WriteLine($"self->{field.NameToken.Value} = {value}");
                    }
                }
            }

            _writer.WriteLine("}");
            _writer.WriteLine();

            _writer.WriteLine($"void {CType.Create(structType.StructType)}_destroy({CType.Create(structType.StructType)} *self)");
            _writer.WriteLine("{");
            using (_writer.Indent())
            {
                foreach (var field in structType.Fields)
                {
                    if (field.Type is NubRefType)
                    {
                        _writer.WriteLine($"rc_release(self->{field.NameToken.Value});");
                    }
                }
            }

            _writer.WriteLine("}");
            _writer.WriteLine();
        }

        // note(nub31): Normal functions
        foreach (var funcNode in _compilationUnit.OfType<FuncNode>())
        {
            if (funcNode.Body == null) continue;

            EmitLine(funcNode.Tokens.FirstOrDefault());
            var parameters = funcNode.Prototype.Parameters.Count != 0
                ? string.Join(", ", funcNode.Prototype.Parameters.Select(x => CType.Create(x.Type, x.NameToken.Value)))
                : "void";

            var name = FuncName(GetModuleName(), funcNode.NameToken.Value, funcNode.Prototype.ExternSymbolToken?.Value);
            _writer.WriteLine($"{CType.Create(funcNode.Prototype.ReturnType, name)}({parameters})");
            _writer.WriteLine("{");
            using (_writer.Indent())
            {
                using (BeginScope())
                {
                    foreach (var parameter in funcNode.Prototype.Parameters)
                    {
                        if (parameter.Type is NubRefType)
                        {
                            Scope.Defer(() => _writer.WriteLine($"rc_release({parameter.NameToken.Value});"));
                        }
                    }

                    EmitBlock(funcNode.Body);
                }
            }

            _writer.WriteLine("}");
            _writer.WriteLine();
        }

        return _writer.ToString();
    }

    private void EmitStatement(StatementNode statementNode)
    {
        EmitLine(statementNode.Tokens.FirstOrDefault());
        switch (statementNode)
        {
            case AssignmentNode assignmentNode:
                EmitAssignment(assignmentNode);
                break;
            case BlockNode blockNode:
                _writer.WriteLine("{");
                using (_writer.Indent())
                {
                    using (BeginScope())
                    {
                        EmitBlock(blockNode);
                    }
                }

                _writer.WriteLine("}");
                break;
            case BreakNode breakNode:
                EmitBreak(breakNode);
                break;
            case ContinueNode continueNode:
                EmitContinue(continueNode);
                break;
            case DeferNode deferNode:
                Scope.Defer(() => EmitStatement(deferNode.Statement));
                break;
            case ForConstArrayNode forConstArrayNode:
                EmitForConstArray(forConstArrayNode);
                break;
            case ForSliceNode forSliceNode:
                EmitForSlice(forSliceNode);
                break;
            case IfNode ifNode:
                EmitIf(ifNode);
                break;
            case ReturnNode returnNode:
                EmitReturn(returnNode);
                break;
            case StatementFuncCallNode statementFuncCallNode:
                EmitStatementFuncCall(statementFuncCallNode);
                break;
            case VariableDeclarationNode variableDeclarationNode:
                EmitVariableDeclaration(variableDeclarationNode);
                break;
            case WhileNode whileNode:
                EmitWhile(whileNode);
                break;
            default:
                throw new ArgumentOutOfRangeException(nameof(statementNode));
        }
    }

    private void EmitLine(Token? token)
    {
        if (token == null) return;
        var file = token.Span.FilePath;
        var line = token.Span.Start.Line;
        _writer.WriteLine($"#line {line} \"{file}\"");
    }

    private void EmitAssignment(AssignmentNode assignmentNode)
    {
        var target = EmitExpression(assignmentNode.Target);
        var value = EmitExpression(assignmentNode.Value);

        if (assignmentNode.Target.Type is NubRefType)
        {
            _writer.WriteLine($"rc_retain({value});");
            Scope.Defer(() => _writer.WriteLine($"rc_release({value});"));
            _writer.WriteLine($"rc_release({target});");
        }

        _writer.WriteLine($"{target} = {value};");
    }

    private void EmitBreak(BreakNode _)
    {
        // todo(nub31): Emit deferred statements
        _writer.WriteLine("break;");
    }

    private void EmitContinue(ContinueNode _)
    {
        // todo(nub31): Emit deferred statements
        _writer.WriteLine("continue;");
    }

    private void EmitForSlice(ForSliceNode forSliceNode)
    {
        var targetType = (NubSliceType)forSliceNode.Target.Type;
        var target = EmitExpression(forSliceNode.Target);
        var indexName = forSliceNode.IndexNameToken?.Value ?? NewTmp();

        _writer.WriteLine($"for (unsigned long long {indexName} = 0; {indexName} < {target}.length; ++{indexName})");
        _writer.WriteLine("{");
        using (_writer.Indent())
        {
            using (BeginScope())
            {
                _writer.WriteLine($"{CType.Create(targetType.ElementType, forSliceNode.ElementNameToken.Value)} = (({CType.Create(targetType.ElementType)}*){target}.data)[{indexName}];");
                EmitBlock(forSliceNode.Body);
            }
        }

        _writer.WriteLine("}");
    }

    private void EmitForConstArray(ForConstArrayNode forConstArrayNode)
    {
        var targetType = (NubConstArrayType)forConstArrayNode.Target.Type;
        var target = EmitExpression(forConstArrayNode.Target);
        var indexName = forConstArrayNode.IndexNameToken?.Value ?? NewTmp();

        _writer.WriteLine($"for (unsigned long long {indexName} = 0; {indexName} < {targetType.Size}; ++{indexName})");
        _writer.WriteLine("{");
        using (_writer.Indent())
        {
            using (BeginScope())
            {
                _writer.WriteLine($"{CType.Create(targetType.ElementType, forConstArrayNode.ElementNameToken.Value)} = {target}[{indexName}];");
                EmitBlock(forConstArrayNode.Body);
            }
        }

        _writer.WriteLine("}");
    }

    private void EmitIf(IfNode ifNode, bool elseIf = false)
    {
        var condition = EmitExpression(ifNode.Condition);
        _writer.WriteLine($"{(elseIf ? "else " : "")}if ({condition})");
        _writer.WriteLine("{");
        using (_writer.Indent())
        {
            using (BeginScope())
            {
                EmitBlock(ifNode.Body);
            }
        }

        _writer.WriteLine("}");
        ifNode.Else?.Match
        (
            elseIfNode => EmitIf(elseIfNode, true),
            elseNode =>
            {
                _writer.WriteLine("else");
                _writer.WriteLine("{");
                using (_writer.Indent())
                {
                    using (BeginScope())
                    {
                        EmitBlock(elseNode);
                    }
                }

                _writer.WriteLine("}");
            }
        );
    }

    private void EmitReturn(ReturnNode returnNode)
    {
        if (returnNode.Value == null)
        {
            EmitScopeCleanup();
            _writer.WriteLine("return;");
        }
        else
        {
            var returnValue = EmitExpression(returnNode.Value);
            var tmp = NewTmp();
            _writer.WriteLine($"{CType.Create(returnNode.Value.Type, tmp)} = {returnValue};");
            if (returnNode.Value.Type is NubRefType)
            {
                _writer.WriteLine($"rc_retain({tmp});");
            }

            EmitScopeCleanup();
            EmitLine(returnNode.Tokens.FirstOrDefault());
            _writer.WriteLine($"return {tmp};");
        }
    }

    private void EmitStatementFuncCall(StatementFuncCallNode statementFuncCallNode)
    {
        EmitFuncCall(statementFuncCallNode.FuncCall);
    }

    private void EmitVariableDeclaration(VariableDeclarationNode variableDeclarationNode)
    {
        if (variableDeclarationNode.Assignment != null)
        {
            var value = EmitExpression(variableDeclarationNode.Assignment);

            if (variableDeclarationNode.Type is NubRefType)
            {
                _writer.WriteLine($"rc_retain({value});");
                Scope.Defer(() => _writer.WriteLine($"rc_release({value});"));
            }

            _writer.WriteLine($"{CType.Create(variableDeclarationNode.Type, variableDeclarationNode.NameToken.Value)} = {value};");
        }
        else
        {
            _writer.WriteLine($"{CType.Create(variableDeclarationNode.Type, variableDeclarationNode.NameToken.Value)};");
        }
    }

    private void EmitWhile(WhileNode whileNode)
    {
        var condition = EmitExpression(whileNode.Condition);
        _writer.WriteLine($"while ({condition})");
        _writer.WriteLine("{");
        using (_writer.Indent())
        {
            using (BeginScope())
            {
                EmitBlock(whileNode.Body);
            }
        }

        _writer.WriteLine("}");
    }

    private string EmitExpression(ExpressionNode expressionNode)
    {
        if (expressionNode is IntermediateExpression)
        {
            throw new UnreachableException("Type checker fucked up");
        }

        var expr = expressionNode switch
        {
            ArrayIndexAccessNode arrayIndexAccessNode => EmitArrayIndexAccess(arrayIndexAccessNode),
            ArrayInitializerNode arrayInitializerNode => EmitArrayInitializer(arrayInitializerNode),
            BinaryExpressionNode binaryExpressionNode => EmitBinaryExpression(binaryExpressionNode),
            BoolLiteralNode boolLiteralNode => boolLiteralNode.Value ? "true" : "false",
            ConstArrayIndexAccessNode constArrayIndexAccessNode => EmitConstArrayIndexAccess(constArrayIndexAccessNode),
            ConstArrayInitializerNode constArrayInitializerNode => EmitConstArrayInitializer(constArrayInitializerNode),
            CStringLiteralNode cStringLiteralNode => $"\"{cStringLiteralNode.Value}\"",
            DereferenceNode dereferenceNode => EmitDereference(dereferenceNode),
            Float32LiteralNode float32LiteralNode => EmitFloat32Literal(float32LiteralNode),
            Float64LiteralNode float64LiteralNode => EmitFloat64Literal(float64LiteralNode),
            CastNode castNode => EmitCast(castNode),
            FuncCallNode funcCallNode => EmitFuncCall(funcCallNode),
            FuncIdentifierNode funcIdentifierNode => FuncName(funcIdentifierNode.ModuleToken.Value, funcIdentifierNode.NameToken.Value, funcIdentifierNode.ExternSymbolToken?.Value),
            AddressOfNode addressOfNode => EmitAddressOf(addressOfNode),
            RefDereferenceNode refDereferenceNode => EmitRefDereference(refDereferenceNode),
            RefStructInitializerNode refStructInitializerNode => EmitRefStructInitializer(refStructInitializerNode),
            SizeNode sizeBuiltinNode => $"sizeof({CType.Create(sizeBuiltinNode.TargetType)})",
            SliceIndexAccessNode sliceIndexAccessNode => EmitSliceArrayIndexAccess(sliceIndexAccessNode),
            StringLiteralNode stringLiteralNode => EmitStringLiteral(stringLiteralNode),
            StructFieldAccessNode structFieldAccessNode => EmitStructFieldAccess(structFieldAccessNode),
            StructInitializerNode structInitializerNode => EmitStructInitializer(structInitializerNode),
            I8LiteralNode i8LiteralNode => EmitI8Literal(i8LiteralNode),
            I16LiteralNode i16LiteralNode => EmitI16Literal(i16LiteralNode),
            I32LiteralNode i32LiteralNode => EmitI32Literal(i32LiteralNode),
            I64LiteralNode i64LiteralNode => EmitI64Literal(i64LiteralNode),
            U8LiteralNode u8LiteralNode => EmitU8Literal(u8LiteralNode),
            U16LiteralNode u16LiteralNode => EmitU16Literal(u16LiteralNode),
            U32LiteralNode u32LiteralNode => EmitU32Literal(u32LiteralNode),
            U64LiteralNode u64LiteralNode => EmitU64Literal(u64LiteralNode),
            UnaryExpressionNode unaryExpressionNode => EmitUnaryExpression(unaryExpressionNode),
            VariableIdentifierNode variableIdentifierNode => variableIdentifierNode.NameToken.Value,
            _ => throw new ArgumentOutOfRangeException(nameof(expressionNode))
        };

        return $"({expr})";
    }

    private string EmitArrayIndexAccess(ArrayIndexAccessNode arrayIndexAccessNode)
    {
        var target = EmitExpression(arrayIndexAccessNode.Target);
        var index = EmitExpression(arrayIndexAccessNode.Index);
        return $"{target}[{index}]";
    }

    private string EmitArrayInitializer(ArrayInitializerNode arrayInitializerNode)
    {
        var values = new List<string>();
        foreach (var value in arrayInitializerNode.Values)
        {
            values.Add(EmitExpression(value));
        }

        var arrayType = (NubArrayType)arrayInitializerNode.Type;
        return $"({CType.Create(arrayType.ElementType)}[]){{{string.Join(", ", values)}}}";
    }

    private string EmitBinaryExpression(BinaryExpressionNode binaryExpressionNode)
    {
        var left = EmitExpression(binaryExpressionNode.Left);
        var right = EmitExpression(binaryExpressionNode.Right);

        var op = binaryExpressionNode.Operator switch
        {
            BinaryOperator.Plus => "+",
            BinaryOperator.Minus => "-",
            BinaryOperator.Multiply => "*",
            BinaryOperator.Divide => "/",
            BinaryOperator.Modulo => "%",
            BinaryOperator.Equal => "==",
            BinaryOperator.NotEqual => "!=",
            BinaryOperator.LessThan => "<",
            BinaryOperator.LessThanOrEqual => "<=",
            BinaryOperator.GreaterThan => ">",
            BinaryOperator.GreaterThanOrEqual => ">=",
            BinaryOperator.LogicalAnd => "&&",
            BinaryOperator.LogicalOr => "||",
            BinaryOperator.BitwiseAnd => "&",
            BinaryOperator.BitwiseOr => "|",
            BinaryOperator.BitwiseXor => "^",
            BinaryOperator.LeftShift => "<<",
            BinaryOperator.RightShift => ">>",
            _ => throw new ArgumentOutOfRangeException()
        };

        return $"{left} {op} {right}";
    }

    private string EmitConstArrayIndexAccess(ConstArrayIndexAccessNode constArrayIndexAccessNode)
    {
        var target = EmitExpression(constArrayIndexAccessNode.Target);
        var index = EmitExpression(constArrayIndexAccessNode.Index);
        // todo(nub31): We can emit bounds checking here
        return $"{target}[{index}]";
    }

    private string EmitConstArrayInitializer(ConstArrayInitializerNode arrayInitializerNode)
    {
        var values = new List<string>();
        foreach (var value in arrayInitializerNode.Values)
        {
            values.Add(EmitExpression(value));
        }

        var arrayType = (NubConstArrayType)arrayInitializerNode.Type;
        return $"({CType.Create(arrayType.ElementType)}[{arrayType.Size}]){{{string.Join(", ", values)}}}";
    }

    private string EmitDereference(DereferenceNode dereferenceNode)
    {
        var pointer = EmitExpression(dereferenceNode.Target);
        return $"*{pointer}";
    }

    private string EmitFloat32Literal(Float32LiteralNode float32LiteralNode)
    {
        var str = float32LiteralNode.Value.ToString("G9", System.Globalization.CultureInfo.InvariantCulture);
        if (!str.Contains('.') && !str.Contains('e') && !str.Contains('E'))
        {
            str += ".0";
        }

        return str + "f";
    }

    private string EmitFloat64Literal(Float64LiteralNode float64LiteralNode)
    {
        var str = float64LiteralNode.Value.ToString("G17", System.Globalization.CultureInfo.InvariantCulture);
        if (!str.Contains('.') && !str.Contains('e') && !str.Contains('E'))
        {
            str += ".0";
        }

        return str;
    }

    private string EmitCast(CastNode castNode)
    {
        var value = EmitExpression(castNode.Value);

        if (castNode is { Type: NubSliceType sliceType, Value.Type: NubConstArrayType arrayType })
        {
            return $"({CType.Create(sliceType)}){{.length = {arrayType.Size}, .data = (void*){value}}}";
        }

        return $"({CType.Create(castNode.Type)}){value}";
    }

    private string EmitFuncCall(FuncCallNode funcCallNode)
    {
        var name = EmitExpression(funcCallNode.Expression);

        var parameterNames = new List<string>();
        foreach (var parameter in funcCallNode.Parameters)
        {
            var result = EmitExpression(parameter);
            if (parameter.Type is NubRefType)
            {
                _writer.WriteLine($"rc_retain({result});");
            }

            parameterNames.Add(result);
        }

        var tmp = NewTmp();

        _writer.WriteLine($"{CType.Create(funcCallNode.Type)} {tmp} = {name}({string.Join(", ", parameterNames)});");

        if (funcCallNode.Type is NubRefType)
        {
            Scope.Defer(() => _writer.WriteLine($"rc_release({tmp});"));
        }

        return tmp;
    }

    private string EmitAddressOf(AddressOfNode addressOfNode)
    {
        var value = EmitExpression(addressOfNode.LValue);
        return $"&{value}";
    }

    private string EmitRefDereference(RefDereferenceNode refDereferenceNode)
    {
        var pointer = EmitExpression(refDereferenceNode.Target);
        return $"*{pointer}";
    }

    private string EmitRefStructInitializer(RefStructInitializerNode refStructInitializerNode)
    {
        var type = (NubRefType)refStructInitializerNode.Type;
        var structType = (NubStructType)type.BaseType;

        var tmp = NewTmp();
        _writer.WriteLine($"{CType.Create(type)} {tmp} = ({CType.Create(type)})rc_alloc(sizeof({CType.Create(structType)}), (void (*)(void *)){CType.Create(structType)}_destroy);");
        Scope.Defer(() => _writer.WriteLine($"rc_release({tmp});"));

        _writer.WriteLine($"*{tmp} = ({CType.Create(structType)}){{{0}}};");
        _writer.WriteLine($"{CType.Create(structType)}_create({tmp});");

        foreach (var initializer in refStructInitializerNode.Initializers)
        {
            var value = EmitExpression(initializer.Value);
            _writer.WriteLine($"{tmp}->{initializer.Key} = {value};");
        }

        return tmp;
    }

    private string EmitSliceArrayIndexAccess(SliceIndexAccessNode sliceIndexAccessNode)
    {
        var targetType = (NubSliceType)sliceIndexAccessNode.Target.Type;
        var target = EmitExpression(sliceIndexAccessNode.Target);
        var index = EmitExpression(sliceIndexAccessNode.Index);
        // todo(nub31): We can emit bounds checking here
        return $"(({CType.Create(targetType.ElementType)}*){target}.data)[{index}]";
    }

    private string EmitStringLiteral(StringLiteralNode stringLiteralNode)
    {
        var length = Encoding.UTF8.GetByteCount(stringLiteralNode.Value);
        return $"({CType.Create(stringLiteralNode.Type)}){{.length = {length}, .data = \"{stringLiteralNode.Value}\"}}";
    }

    private string EmitStructFieldAccess(StructFieldAccessNode structFieldAccessNode)
    {
        var structExpr = EmitExpression(structFieldAccessNode.Target);
        return $"{structExpr}.{structFieldAccessNode.FieldToken.Value}";
    }

    private string EmitStructInitializer(StructInitializerNode structInitializerNode)
    {
        var structType = (NubStructType)structInitializerNode.Type;

        var tmp = NewTmp();
        _writer.WriteLine($"{CType.Create(structType)} {tmp} = ({CType.Create(structType)}){{0}};");
        _writer.WriteLine($"{CType.Create(structType)}_create(&{tmp});");

        foreach (var initializer in structInitializerNode.Initializers)
        {
            var value = EmitExpression(initializer.Value);
            _writer.WriteLine($"{tmp}.{initializer.Key} = {value};");
        }

        return tmp;
    }

    private string EmitI8Literal(I8LiteralNode i8LiteralNode)
    {
        return i8LiteralNode.Value.ToString();
    }

    private string EmitI16Literal(I16LiteralNode i16LiteralNode)
    {
        return i16LiteralNode.Value.ToString();
    }

    private string EmitI32Literal(I32LiteralNode i32LiteralNode)
    {
        return i32LiteralNode.Value.ToString();
    }

    private string EmitI64Literal(I64LiteralNode i64LiteralNode)
    {
        return i64LiteralNode.Value + "LL";
    }

    private string EmitU8Literal(U8LiteralNode u8LiteralNode)
    {
        return u8LiteralNode.Value.ToString();
    }

    private string EmitU16Literal(U16LiteralNode u16LiteralNode)
    {
        return u16LiteralNode.Value.ToString();
    }

    private string EmitU32Literal(U32LiteralNode u32LiteralNode)
    {
        return u32LiteralNode.Value.ToString();
    }

    private string EmitU64Literal(U64LiteralNode u64LiteralNode)
    {
        return u64LiteralNode.Value + "ULL";
    }

    private string EmitUnaryExpression(UnaryExpressionNode unaryExpressionNode)
    {
        var value = EmitExpression(unaryExpressionNode.Operand);

        return unaryExpressionNode.Operator switch
        {
            UnaryOperator.Negate => $"-{value}",
            UnaryOperator.Invert => $"!{value}",
            _ => throw new ArgumentOutOfRangeException()
        };
    }

    private void EmitBlock(BlockNode blockNode)
    {
        foreach (var statementNode in blockNode.Statements)
        {
            EmitStatement(statementNode);
        }
    }

    private void EmitScopeCleanup()
    {
        var deferredStack = Scope.GetDeferred();
        while (deferredStack.TryPop(out var deferred))
        {
            deferred.Invoke();
        }
    }

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

    private sealed class ScopeDisposer(Generator owner) : IDisposable
    {
        public void Dispose()
        {
            owner.EmitScopeCleanup();
            owner._scopes.Pop();
        }
    }
}

public class Scope
{
    private readonly List<Action> _deferred = [];

    public void Defer(Action action)
    {
        _deferred.Add(action);
    }

    public Stack<Action> GetDeferred() => new(_deferred);
}