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

namespace NubLang.Generation;

public class Generator
{
    private readonly List<DefinitionNode> _definitions;
    private readonly IndentedTextWriter _writer;
    private readonly Stack<List<DeferNode>> _deferStack = [];
    private readonly Stack<(string Name, NubFuncType FuncType)> _funcDefs = [];
    private readonly List<NubStructType> _structTypes = [];
    private int _tmpIndex;
    private int _funcDefIndex;

    public Generator(List<DefinitionNode> definitions)
    {
        _definitions = definitions;
        _writer = new IndentedTextWriter();
    }

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

    private string MapType(NubType nubType)
    {
        return nubType switch
        {
            NubArrayType arrayType => MapType(arrayType.ElementType),
            NubConstArrayType arrayType => MapType(arrayType.ElementType),
            NubBoolType => "bool",
            NubCStringType => "char",
            NubFloatType floatType => floatType.Width switch
            {
                32 => "f32",
                64 => "f64",
                _ => throw new ArgumentOutOfRangeException()
            },
            NubFuncType funcType => MapFuncType(funcType),
            NubIntType intType => (intType.Signed, intType.Width) switch
            {
                (false, 8) => "u8",
                (false, 16) => "u16",
                (false, 32) => "u32",
                (false, 64) => "u64",
                (true, 8) => "i8",
                (true, 16) => "i16",
                (true, 32) => "i32",
                (true, 64) => "i64",
                _ => throw new ArgumentOutOfRangeException()
            },
            NubPointerType pointerType => MapType(pointerType.BaseType),
            NubStringType => throw new NotImplementedException(),
            NubStructType structType => MapStructType(structType),
            NubVoidType => "void",
            _ => throw new ArgumentOutOfRangeException(nameof(nubType))
        };
    }

    private string MapStructType(NubStructType structType)
    {
        _structTypes.Add(structType);
        return StructName(structType.Module, structType.Name);
    }

    private string MapFuncType(NubFuncType funcType)
    {
        var name = $"_func_type_def{++_funcDefIndex}";
        _funcDefs.Push((name, funcType));
        return name;
    }

    private string MapNameWithType(NubType nubType, string name)
    {
        var prefix = "";
        var postfix = "";

        switch (nubType)
        {
            case NubCStringType or NubPointerType:
                prefix = "*";
                break;
            case NubArrayType:
                postfix = "[]";
                break;
            case NubConstArrayType constArrayType:
                postfix = $"[{constArrayType.Size}]";
                break;
        }

        return $"{MapType(nubType)} {prefix}{name}{postfix}";
    }

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

    private static string StructName(string module, string name)
    {
        return $"{module}_{name}";
    }

    private static string StructFuncName(string module, string name, string function)
    {
        return $"{module}_{name}_{function}";
    }

    public string Emit()
    {
        const string header = """
                              typedef __builtin_va_list va_list;

                              #define va_start(ap, last) __builtin_va_start(ap, last)
                              #define va_arg(ap, type) __builtin_va_arg(ap, type)
                              #define va_end(ap) __builtin_va_end(ap)
                              #define va_copy(dest, src) __builtin_va_copy(dest, src)

                              #define NULL ((void*)0)

                              typedef unsigned long size_t;
                              typedef unsigned long uintptr_t;

                              #define offsetof(type, member) __builtin_offsetof(type, member)

                              typedef unsigned char u8;
                              typedef signed char i8;
                              typedef unsigned short u16;
                              typedef signed short i16;
                              typedef unsigned int u32;
                              typedef signed int i32;
                              typedef unsigned long long u64;
                              typedef signed long long i64;

                              typedef float f32;
                              typedef double f64;

                              #define I8_C(x) x
                              #define U8_C(x) x##U

                              #define I16_C(x) x
                              #define U16_C(x) x##U

                              #define I32_C(x) x
                              #define U32_C(x) x##U

                              #define I64_C(x) x##LL
                              #define U64_C(x) x##ULL
                              """;

        _writer.WriteLine();

        var appendNewLine = false;

        foreach (var funcNode in _definitions.OfType<FuncNode>())
        {
            EmitLine(funcNode.Tokens.FirstOrDefault());
            appendNewLine = true;
            var parameters = funcNode.Signature.Parameters.Count != 0
                ? string.Join(", ", funcNode.Signature.Parameters.Select(x => MapNameWithType(x.Type, x.Name)))
                : "void";

            var name = FuncName(funcNode.Module, funcNode.Name, funcNode.ExternSymbol);
            _writer.WriteLine($"{MapNameWithType(funcNode.Signature.ReturnType, name)}({parameters});");
        }

        if (appendNewLine)
        {
            _writer.WriteLine();
        }

        foreach (var structNode in _definitions.OfType<StructNode>())
        {
            foreach (var structFuncNode in structNode.Functions)
            {
                EmitLine(structFuncNode.Tokens.FirstOrDefault());
                var parameters = structFuncNode.Signature.Parameters.Count != 0
                    ? string.Join(", ", structFuncNode.Signature.Parameters.Select(x => MapNameWithType(x.Type, x.Name)))
                    : "void";

                var name = StructFuncName(structNode.Module, structNode.Name, structFuncNode.Name);
                _writer.WriteLine($"{MapNameWithType(structFuncNode.Signature.ReturnType, name)}({parameters})");
                EmitBlock(structFuncNode.Body);
                _writer.WriteLine();
            }
        }

        foreach (var funcNode in _definitions.OfType<FuncNode>())
        {
            if (funcNode.Body == null) continue;

            EmitLine(funcNode.Tokens.FirstOrDefault());
            var parameters = funcNode.Signature.Parameters.Count != 0
                ? string.Join(", ", funcNode.Signature.Parameters.Select(x => MapNameWithType(x.Type, x.Name)))
                : "void";

            if (funcNode.ExternSymbol == null)
            {
                _writer.Write("static ");
            }

            var name = FuncName(funcNode.Module, funcNode.Name, funcNode.ExternSymbol);
            _writer.WriteLine($"{MapNameWithType(funcNode.Signature.ReturnType, name)}({parameters})");
            EmitBlock(funcNode.Body);
            _writer.WriteLine();
        }

        List<string> typedefs = [];

        while (_funcDefs.TryPop(out var funcTypeDef))
        {
            var returnType = MapType(funcTypeDef.FuncType.ReturnType);
            var paramList = string.Join(", ", funcTypeDef.FuncType.Parameters.Select((type, i) => MapNameWithType(type, $"arg{i}")));
            if (funcTypeDef.FuncType.Parameters.Count == 0)
            {
                paramList = "void";
            }

            typedefs.Add($"typedef {returnType} (*{funcTypeDef.Name})({paramList});");
        }

        var structDefSb = new StringBuilder();
        foreach (var structType in _structTypes)
        {
            structDefSb.AppendLine("typedef struct");
            structDefSb.AppendLine("{");
            foreach (var field in structType.Fields)
            {
                structDefSb.AppendLine($"    {MapNameWithType(field.Type, field.Name)};");
            }

            structDefSb.AppendLine($"}} {StructName(structType.Module, structType.Name)};");
            structDefSb.AppendLine();
        }

        return header + structDefSb + "\n\n" + string.Join('\n', typedefs) + "\n\n" + _writer;
    }

    private void EmitStatement(StatementNode statementNode)
    {
        EmitLine(statementNode.Tokens.FirstOrDefault());
        switch (statementNode)
        {
            case AssignmentNode assignmentNode:
                EmitAssignment(assignmentNode);
                break;
            case BlockNode blockNode:
                EmitBlock(blockNode);
                break;
            case BreakNode breakNode:
                EmitBreak(breakNode);
                break;
            case ContinueNode continueNode:
                EmitContinue(continueNode);
                break;
            case DeferNode deferNode:
                EmitDefer(deferNode);
                break;
            case IfNode ifNode:
                EmitIf(ifNode);
                break;
            case ReturnNode returnNode:
                EmitReturn(returnNode);
                break;
            case StatementFuncCallNode statementFuncCallNode:
                EmitStatementFuncCall(statementFuncCallNode);
                break;
            case StatementStructFuncCallNode statementStructFuncCallNode:
                EmitStatementStructFuncCall(statementStructFuncCallNode);
                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);
        _writer.WriteLine($"{target} = {value};");
    }

    private void EmitBreak(BreakNode _)
    {
        _writer.WriteLine("break;");
    }

    private void EmitContinue(ContinueNode _)
    {
        _writer.WriteLine("continue;");
    }

    private void EmitDefer(DeferNode deferNode)
    {
        _deferStack.Peek().Add(deferNode);
    }

    private void EmitIf(IfNode ifNode, bool elseIf = false)
    {
        var condition = EmitExpression(ifNode.Condition);
        _writer.WriteLine($"{(elseIf ? "else" : "")} if ({condition})");
        EmitBlock(ifNode.Body);
        ifNode.Else?.Match
        (
            elseIfNode => EmitIf(elseIfNode, true),
            elseNode =>
            {
                _writer.WriteLine("else");
                EmitBlock(elseNode);
            }
        );
    }

    private void EmitReturn(ReturnNode returnNode)
    {
        if (returnNode.Value == null)
        {
            var blockDefers = _deferStack.Peek();
            for (var i = blockDefers.Count - 1; i >= 0; i--)
            {
                EmitStatement(blockDefers[i].Statement);
            }

            _writer.WriteLine("return;");
        }
        else
        {
            var returnValue = EmitExpression(returnNode.Value);

            if (_deferStack.Peek().Count != 0)
            {
                var tmp = NewTmp();
                _writer.WriteLine($"{MapNameWithType(returnNode.Value.Type, tmp)} = {returnValue};");

                var blockDefers = _deferStack.Peek();
                for (var i = blockDefers.Count - 1; i >= 0; i--)
                {
                    EmitStatement(blockDefers[i].Statement);
                }

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

    private void EmitStatementFuncCall(StatementFuncCallNode statementFuncCallNode)
    {
        var funcCall = EmitFuncCall(statementFuncCallNode.FuncCall);
        _writer.WriteLine($"{funcCall};");
    }

    private void EmitStatementStructFuncCall(StatementStructFuncCallNode statementStructFuncCallNode)
    {
        var structFuncCall = EmitStructFuncCall(statementStructFuncCallNode.StructFuncCall);
        _writer.WriteLine($"{structFuncCall};");
    }

    private void EmitVariableDeclaration(VariableDeclarationNode variableDeclarationNode)
    {
        if (variableDeclarationNode.Assignment != null)
        {
            var value = EmitExpression(variableDeclarationNode.Assignment);
            _writer.WriteLine($"{MapNameWithType(variableDeclarationNode.Type, variableDeclarationNode.Name)} = {value};");
        }
        else
        {
            _writer.WriteLine($"{MapNameWithType(variableDeclarationNode.Type, variableDeclarationNode.Name)};");
        }
    }

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

    private string EmitExpression(ExpressionNode expressionNode)
    {
        var expr = expressionNode switch
        {
            ArrayIndexAccessNode arrayIndexAccessNode => EmitArrayIndexAccess(arrayIndexAccessNode),
            ArrayInitializerNode arrayInitializerNode => EmitArrayInitializer(arrayInitializerNode),
            BinaryExpressionNode binaryExpressionNode => EmitBinaryExpression(binaryExpressionNode),
            BoolLiteralNode boolLiteralNode => EmitBoolLiteral(boolLiteralNode),
            ConvertFloatNode convertFloatNode => EmitConvertFloat(convertFloatNode),
            ConvertIntNode convertIntNode => EmitConvertInt(convertIntNode),
            CStringLiteralNode cStringLiteralNode => EmitCStringLiteral(cStringLiteralNode),
            DereferenceNode dereferenceNode => EmitDereference(dereferenceNode),
            Float32LiteralNode float32LiteralNode => EmitFloat32Literal(float32LiteralNode),
            Float64LiteralNode float64LiteralNode => EmitFloat64Literal(float64LiteralNode),
            FloatToIntBuiltinNode floatToIntBuiltinNode => EmitFloatToIntBuiltin(floatToIntBuiltinNode),
            FuncCallNode funcCallNode => EmitFuncCall(funcCallNode),
            FuncIdentifierNode funcIdentifierNode => EmitFuncIdentifier(funcIdentifierNode),
            IntLiteralNode intLiteralNode => EmitIntLiteral(intLiteralNode),
            AddressOfNode addressOfNode => EmitAddressOf(addressOfNode),
            LValueIdentifierNode lValueIdentifierNode => EmitLValueIdentifier(lValueIdentifierNode),
            RValueIdentifierNode rValueIdentifierNode => EmitRValueIdentifier(rValueIdentifierNode),
            SizeBuiltinNode sizeBuiltinNode => EmitSizeBuiltin(sizeBuiltinNode),
            StringLiteralNode stringLiteralNode => EmitStringLiteral(stringLiteralNode),
            StructFieldAccessNode structFieldAccessNode => EmitStructFieldAccess(structFieldAccessNode),
            StructFuncCallNode structFuncCallNode => EmitStructFuncCall(structFuncCallNode),
            StructInitializerNode structInitializerNode => EmitStructInitializer(structInitializerNode),
            UIntLiteralNode uIntLiteralNode => EmitUIntLiteral(uIntLiteralNode),
            UnaryExpressionNode unaryExpressionNode => EmitUnaryExpression(unaryExpressionNode),
            _ => throw new ArgumentOutOfRangeException(nameof(expressionNode))
        };

        return $"({expr})";
    }

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

    private string EmitArrayInitializer(ArrayInitializerNode arrayInitializerNode)
    {
        throw new NotImplementedException();
    }

    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 EmitBoolLiteral(BoolLiteralNode boolLiteralNode)
    {
        return boolLiteralNode.Value ? "true" : "false";
    }

    private string EmitConvertFloat(ConvertFloatNode convertFloatNode)
    {
        var value = EmitExpression(convertFloatNode.Value);
        var targetCast = convertFloatNode.TargetType.Width switch
        {
            32 => "f32",
            64 => "f64",
            _ => throw new ArgumentOutOfRangeException()
        };

        return $"({targetCast}){value}";
    }

    private string EmitConvertInt(ConvertIntNode convertIntNode)
    {
        var value = EmitExpression(convertIntNode.Value);
        var targetType = (convertIntNode.TargetType.Signed, convertIntNode.TargetType.Width) switch
        {
            (false, 8) => "u8",
            (false, 16) => "u16",
            (false, 32) => "u32",
            (false, 64) => "u64",
            (true, 8) => "i8",
            (true, 16) => "i16",
            (true, 32) => "i32",
            (true, 64) => "i64",
            _ => throw new ArgumentOutOfRangeException()
        };
        return $"({targetType}){value}";
    }

    private string EmitCStringLiteral(CStringLiteralNode cStringLiteralNode)
    {
        return $"\"{cStringLiteralNode.Value}\"";
    }

    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 EmitFloatToIntBuiltin(FloatToIntBuiltinNode floatToIntBuiltinNode)
    {
        var value = EmitExpression(floatToIntBuiltinNode.Value);
        var targetType = (floatToIntBuiltinNode.TargetType.Signed, floatToIntBuiltinNode.TargetType.Width) switch
        {
            (false, 8) => "u8",
            (false, 16) => "u16",
            (false, 32) => "u32",
            (false, 64) => "u64",
            (true, 8) => "i8",
            (true, 16) => "i16",
            (true, 32) => "i32",
            (true, 64) => "i64",
            _ => throw new ArgumentOutOfRangeException()
        };
        return $"({targetType}){value}";
    }

    private string EmitFuncCall(FuncCallNode funcCallNode)
    {
        var name = EmitExpression(funcCallNode.Expression);
        var parameterNames = funcCallNode.Parameters.Select(EmitExpression).ToList();
        return $"{name}({string.Join(", ", parameterNames)})";
    }

    private string EmitFuncIdentifier(FuncIdentifierNode funcIdentifierNode)
    {
        return FuncName(funcIdentifierNode.Module, funcIdentifierNode.Name, funcIdentifierNode.ExternSymbol);
    }

    private string EmitIntLiteral(IntLiteralNode intLiteralNode)
    {
        var type = (NubIntType)intLiteralNode.Type;
        return type.Width switch
        {
            8 => $"I8_C({intLiteralNode.Value})",
            16 => $"I16_C({intLiteralNode.Value})",
            32 => $"I32_C({intLiteralNode.Value})",
            64 => $"I64_C({intLiteralNode.Value})",
            _ => throw new ArgumentOutOfRangeException()
        };
    }

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

    private string EmitLValueIdentifier(LValueIdentifierNode lValueIdentifierNode)
    {
        return lValueIdentifierNode.Name;
    }

    private string EmitRValueIdentifier(RValueIdentifierNode rValueIdentifierNode)
    {
        return rValueIdentifierNode.Name;
    }

    private string EmitSizeBuiltin(SizeBuiltinNode sizeBuiltinNode)
    {
        if (sizeBuiltinNode.TargetType is NubConstArrayType constArrayType)
        {
            return $"sizeof({MapType(constArrayType.ElementType)}) * {constArrayType.Size}";
        }

        return $"sizeof({MapType(sizeBuiltinNode.TargetType)})";
    }

    private string EmitStringLiteral(StringLiteralNode stringLiteralNode)
    {
        throw new NotImplementedException();
    }

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

    private string EmitStructFuncCall(StructFuncCallNode structFuncCallNode)
    {
        var thisParameter = EmitExpression(structFuncCallNode.StructExpression);
        var name = StructFuncName(structFuncCallNode.Module, structFuncCallNode.StructName, structFuncCallNode.FuncName);
        var parameterNames = structFuncCallNode.Parameters.Select(EmitExpression).ToList();
        return $"{name}({thisParameter}, {string.Join(", ", parameterNames)})";
    }

    private string EmitStructInitializer(StructInitializerNode structInitializerNode)
    {
        var initValues = new List<string>();
        foreach (var initializer in structInitializerNode.Initializers)
        {
            var value = EmitExpression(initializer.Value);
            initValues.Add($".{initializer.Key} = {value}");
        }

        var initString = initValues.Count == 0
            ? "0"
            : string.Join(", ", initValues);

        return $"({StructName(structInitializerNode.StructType.Module, structInitializerNode.StructType.Name)}){{{initString}}}";
    }

    private string EmitUIntLiteral(UIntLiteralNode uIntLiteralNode)
    {
        var type = (NubIntType)uIntLiteralNode.Type;
        return type.Width switch
        {
            8 => $"U8_C({uIntLiteralNode.Value})",
            16 => $"U16_C({uIntLiteralNode.Value})",
            32 => $"U32_C({uIntLiteralNode.Value})",
            64 => $"U64_C({uIntLiteralNode.Value})",
            _ => throw new ArgumentOutOfRangeException()
        };
    }

    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)
    {
        EmitLine(blockNode.Tokens.FirstOrDefault());
        _writer.WriteLine("{");
        using (_writer.Indent())
        {
            _deferStack.Push([]);

            foreach (var statementNode in blockNode.Statements)
            {
                EmitStatement(statementNode);
            }

            var blockDefers = _deferStack.Pop();
            for (var i = blockDefers.Count - 1; i >= 0; i--)
            {
                EmitStatement(blockDefers[i].Statement);
            }
        }

        EmitLine(blockNode.Tokens.LastOrDefault());
        _writer.WriteLine("}");
    }
}