using NubLang.Syntax;

namespace NubLang.Ast;

public abstract record Node(List<Token> Tokens)
{
    public abstract IEnumerable<Node> Children();

    public IEnumerable<Node> Descendants()
    {
        foreach (var child in Children())
        {
            foreach (var descendant in child.DescendantsAndSelf())
            {
                yield return descendant;
            }
        }
    }

    public IEnumerable<Node> DescendantsAndSelf()
    {
        yield return this;
        foreach (var descendant in Descendants())
        {
            yield return descendant;
        }
    }
}

#region Definitions

public abstract record DefinitionNode(List<Token> Tokens, string Module, string Name) : Node(Tokens);

public record FuncParameterNode(List<Token> Tokens, string Name, NubType Type) : Node(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record FuncPrototypeNode(List<Token> Tokens, string Module, string Name, string? ExternSymbol, List<FuncParameterNode> Parameters, NubType ReturnType) : Node(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        return Parameters;
    }
}

public record FuncNode(List<Token> Tokens, FuncPrototypeNode Prototype, BlockNode? Body) : DefinitionNode(Tokens, Prototype.Module, Prototype.Name)
{
    public override IEnumerable<Node> Children()
    {
        yield return Prototype;
        if (Body != null)
        {
            yield return Body;
        }
    }
}

#endregion

#region Statements

public abstract record StatementNode(List<Token> Tokens) : Node(Tokens);

public abstract record TerminalStatementNode(List<Token> Tokens) : StatementNode(Tokens);

public record BlockNode(List<Token> Tokens, List<StatementNode> Statements) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        return Statements;
    }
}

public record StatementFuncCallNode(List<Token> Tokens, FuncCallNode FuncCall) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        yield return FuncCall;
    }
}

public record ReturnNode(List<Token> Tokens, ExpressionNode? Value) : TerminalStatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        if (Value != null) yield return Value;
    }
}

public record AssignmentNode(List<Token> Tokens, LValueExpressionNode Target, ExpressionNode Value) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Value;
    }
}

public record IfNode(List<Token> Tokens, ExpressionNode Condition, BlockNode Body, Variant<IfNode, BlockNode>? Else) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        yield return Condition;
        yield return Body;
        if (Else.HasValue)
        {
            yield return Else.Value.Match<Node>(x => x, x => x);
        }
    }
}

public record VariableDeclarationNode(List<Token> Tokens, string Name, ExpressionNode? Assignment, NubType Type) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        if (Assignment != null) yield return Assignment;
    }
}

public record ContinueNode(List<Token> Tokens) : TerminalStatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record BreakNode(List<Token> Tokens) : TerminalStatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record WhileNode(List<Token> Tokens, ExpressionNode Condition, BlockNode Body) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        yield return Condition;
        yield return Body;
    }
}

public record ForSliceNode(List<Token> Tokens, string ElementName, string? IndexName, ExpressionNode Target, BlockNode Body) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Body;
    }
}

public record ForConstArrayNode(List<Token> Tokens, string ElementName, string? IndexName, ExpressionNode Target, BlockNode Body) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Body;
    }
}

public record DeferNode(List<Token> Tokens, StatementNode Statement) : StatementNode(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        yield return Statement;
    }
}

#endregion

#region Expressions

public enum UnaryOperator
{
    Negate,
    Invert
}

public enum BinaryOperator
{
    Equal,
    NotEqual,
    GreaterThan,
    GreaterThanOrEqual,
    LessThan,
    LessThanOrEqual,
    LogicalAnd,
    LogicalOr,
    Plus,
    Minus,
    Multiply,
    Divide,
    Modulo,
    LeftShift,
    RightShift,
    BitwiseAnd,
    BitwiseXor,
    BitwiseOr
}

public abstract record ExpressionNode(List<Token> Tokens, NubType Type) : Node(Tokens);

public abstract record LValueExpressionNode(List<Token> Tokens, NubType Type) : ExpressionNode(Tokens, Type);

public abstract record RValueExpressionNode(List<Token> Tokens, NubType Type) : ExpressionNode(Tokens, Type);

public abstract record IntermediateExpression(List<Token> Tokens) : ExpressionNode(Tokens, new NubVoidType());

public record StringLiteralNode(List<Token> Tokens, string Value) : RValueExpressionNode(Tokens, new NubStringType())
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record CStringLiteralNode(List<Token> Tokens, string Value) : RValueExpressionNode(Tokens, new NubCStringType())
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record I8LiteralNode(List<Token> Tokens, sbyte Value) : RValueExpressionNode(Tokens, new NubIntType(true, 8))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record I16LiteralNode(List<Token> Tokens, short Value) : RValueExpressionNode(Tokens, new NubIntType(true, 16))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record I32LiteralNode(List<Token> Tokens, int Value) : RValueExpressionNode(Tokens, new NubIntType(true, 32))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record I64LiteralNode(List<Token> Tokens, long Value) : RValueExpressionNode(Tokens, new NubIntType(true, 64))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record U8LiteralNode(List<Token> Tokens, byte Value) : RValueExpressionNode(Tokens, new NubIntType(false, 8))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record U16LiteralNode(List<Token> Tokens, ushort Value) : RValueExpressionNode(Tokens, new NubIntType(false, 16))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record U32LiteralNode(List<Token> Tokens, uint Value) : RValueExpressionNode(Tokens, new NubIntType(false, 32))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record U64LiteralNode(List<Token> Tokens, ulong Value) : RValueExpressionNode(Tokens, new NubIntType(false, 64))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record Float32LiteralNode(List<Token> Tokens, float Value) : RValueExpressionNode(Tokens, new NubFloatType(32))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record Float64LiteralNode(List<Token> Tokens, double Value) : RValueExpressionNode(Tokens, new NubFloatType(64))
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record BoolLiteralNode(List<Token> Tokens, NubType Type, bool Value) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record BinaryExpressionNode(List<Token> Tokens, NubType Type, ExpressionNode Left, BinaryOperator Operator, ExpressionNode Right) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Left;
        yield return Right;
    }
}

public record UnaryExpressionNode(List<Token> Tokens, NubType Type, UnaryOperator Operator, ExpressionNode Operand) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Operand;
    }
}

public record FuncCallNode(List<Token> Tokens, NubType Type, ExpressionNode Expression, List<ExpressionNode> Parameters) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Expression;
        foreach (var expressionNode in Parameters)
        {
            yield return expressionNode;
        }
    }
}

public record VariableIdentifierNode(List<Token> Tokens, NubType Type, string Name) : LValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record FuncIdentifierNode(List<Token> Tokens, NubType Type, string Module, string Name, string? ExternSymbol) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record ArrayInitializerNode(List<Token> Tokens, NubType Type, List<ExpressionNode> Values) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        return Values;
    }
}

public record ConstArrayInitializerNode(List<Token> Tokens, NubType Type, List<ExpressionNode> Values) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        return Values;
    }
}

public record ArrayIndexAccessNode(List<Token> Tokens, NubType Type, ExpressionNode Target, ExpressionNode Index) : LValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Index;
    }
}

public record ConstArrayIndexAccessNode(List<Token> Tokens, NubType Type, ExpressionNode Target, ExpressionNode Index) : LValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Index;
    }
}

public record SliceIndexAccessNode(List<Token> Tokens, NubType Type, ExpressionNode Target, ExpressionNode Index) : LValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Index;
    }
}

public record AddressOfNode(List<Token> Tokens, NubType Type, LValueExpressionNode LValue) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return LValue;
    }
}

public record StructFieldAccessNode(List<Token> Tokens, NubType Type, ExpressionNode Target, string Field) : LValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Target;
    }
}

public record StructInitializerNode(List<Token> Tokens, NubType Type, Dictionary<string, ExpressionNode> Initializers) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        foreach (var initializer in Initializers)
        {
            yield return initializer.Value;
        }
    }
}

public record DereferenceNode(List<Token> Tokens, NubType Type, ExpressionNode Target) : LValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Target;
    }
}

public record SizeNode(List<Token> Tokens, NubType Type, NubType TargetType) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public record CastNode(List<Token> Tokens, NubType Type, ExpressionNode Value) : RValueExpressionNode(Tokens, Type)
{
    public override IEnumerable<Node> Children()
    {
        yield return Value;
    }
}

public record EnumReferenceIntermediateNode(List<Token> Tokens, string Module, string Name) : IntermediateExpression(Tokens)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

#endregion