using NubLang.Syntax;

namespace NubLang.Ast;

public abstract class Node(List<Token> tokens)
{
    public List<Token> Tokens { get; } = 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;
        }
    }
}

public abstract class TopLevelNode(List<Token> tokens) : Node(tokens);

public class ImportNode(List<Token> tokens, IdentifierToken nameToken) : TopLevelNode(tokens)
{
    public IdentifierToken NameToken { get; } = nameToken;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class ModuleNode(List<Token> tokens, IdentifierToken nameToken) : TopLevelNode(tokens)
{
    public IdentifierToken NameToken { get; } = nameToken;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

#region Definitions

public abstract class DefinitionNode(List<Token> tokens, IdentifierToken nameToken) : TopLevelNode(tokens)
{
    public IdentifierToken NameToken { get; } = nameToken;
}

public class FuncParameterNode(List<Token> tokens, IdentifierToken nameToken, NubType type) : Node(tokens)
{
    public IdentifierToken NameToken { get; } = nameToken;
    public NubType Type { get; } = type;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class FuncPrototypeNode(List<Token> tokens, IdentifierToken nameToken, StringLiteralToken? externSymbolToken, List<FuncParameterNode> parameters, NubType returnType) : Node(tokens)
{
    public IdentifierToken NameToken { get; } = nameToken;
    public StringLiteralToken? ExternSymbolToken { get; } = externSymbolToken;
    public List<FuncParameterNode> Parameters { get; } = parameters;
    public NubType ReturnType { get; } = returnType;

    public override IEnumerable<Node> Children()
    {
        return Parameters;
    }
}

public class FuncNode(List<Token> tokens, FuncPrototypeNode prototype, BlockNode? body) : DefinitionNode(tokens, prototype.NameToken)
{
    public FuncPrototypeNode Prototype { get; } = prototype;
    public BlockNode? Body { get; } = body;

    public override IEnumerable<Node> Children()
    {
        yield return Prototype;
        if (Body != null)
        {
            yield return Body;
        }
    }
}

public class StructFieldNode(List<Token> tokens, IdentifierToken nameToken, NubType type, ExpressionNode? value) : Node(tokens)
{
    public IdentifierToken NameToken { get; } = nameToken;
    public NubType Type { get; } = type;
    public ExpressionNode? Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        if (Value != null)
        {
            yield return Value;
        }
    }
}

public class StructNode(List<Token> tokens, IdentifierToken name, NubStructType structType, bool packed, List<StructFieldNode> fields) : DefinitionNode(tokens, name)
{
    public NubStructType StructType { get; } = structType;
    public bool Packed { get; } = packed;
    public List<StructFieldNode> Fields { get; } = fields;

    public override IEnumerable<Node> Children()
    {
        foreach (var field in Fields)
        {
            yield return field;
        }
    }
}

#endregion

#region Statements

public abstract class StatementNode(List<Token> tokens) : Node(tokens);

public class BlockNode(List<Token> tokens, List<StatementNode> statements) : StatementNode(tokens)
{
    public List<StatementNode> Statements { get; } = statements;

    public override IEnumerable<Node> Children()
    {
        return Statements;
    }
}

public class StatementFuncCallNode(List<Token> tokens, FuncCallNode funcCall) : StatementNode(tokens)
{
    public FuncCallNode FuncCall { get; } = funcCall;

    public override IEnumerable<Node> Children()
    {
        yield return FuncCall;
    }
}

public class ReturnNode(List<Token> tokens, ExpressionNode? value) : StatementNode(tokens)
{
    public ExpressionNode? Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        if (Value != null) yield return Value;
    }
}

public class AssignmentNode(List<Token> tokens, ExpressionNode target, ExpressionNode value) : StatementNode(tokens)
{
    public ExpressionNode Target { get; } = target;
    public ExpressionNode Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Value;
    }
}

public class IfNode(List<Token> tokens, ExpressionNode condition, BlockNode body, Variant<IfNode, BlockNode>? @else) : StatementNode(tokens)
{
    public ExpressionNode Condition { get; } = condition;
    public BlockNode Body { get; } = body;
    public Variant<IfNode, BlockNode>? Else { get; } = @else;

    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 class VariableDeclarationNode(List<Token> tokens, IdentifierToken nameToken, ExpressionNode? assignment, NubType type) : StatementNode(tokens)
{
    public IdentifierToken NameToken { get; } = nameToken;
    public ExpressionNode? Assignment { get; } = assignment;
    public NubType Type { get; } = type;

    public override IEnumerable<Node> Children()
    {
        if (Assignment != null) yield return Assignment;
    }
}

public class ContinueNode(List<Token> tokens) : StatementNode(tokens)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class BreakNode(List<Token> tokens) : StatementNode(tokens)
{
    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class WhileNode(List<Token> tokens, ExpressionNode condition, BlockNode body) : StatementNode(tokens)
{
    public ExpressionNode Condition { get; } = condition;
    public BlockNode Body { get; } = body;

    public override IEnumerable<Node> Children()
    {
        yield return Condition;
        yield return Body;
    }
}

public class ForSliceNode(List<Token> tokens, IdentifierToken elementNameToken, IdentifierToken? indexNameToken, ExpressionNode target, BlockNode body) : StatementNode(tokens)
{
    public IdentifierToken ElementNameToken { get; } = elementNameToken;
    public IdentifierToken? IndexNameToken { get; } = indexNameToken;
    public ExpressionNode Target { get; } = target;
    public BlockNode Body { get; } = body;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Body;
    }
}

public class ForConstArrayNode(List<Token> tokens, IdentifierToken elementNameToken, IdentifierToken? indexNameToken, ExpressionNode target, BlockNode body) : StatementNode(tokens)
{
    public IdentifierToken ElementNameToken { get; } = elementNameToken;
    public IdentifierToken? IndexNameToken { get; } = indexNameToken;
    public ExpressionNode Target { get; } = target;
    public BlockNode Body { get; } = body;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Body;
    }
}

public class DeferNode(List<Token> tokens, StatementNode statement) : StatementNode(tokens)
{
    public StatementNode Statement { get; } = statement;

    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 class ExpressionNode(List<Token> tokens, NubType type) : Node(tokens)
{
    public NubType Type { get; } = type;
}

public class StringLiteralNode(List<Token> tokens, string value) : ExpressionNode(tokens, new NubStringType())
{
    public string Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class CStringLiteralNode(List<Token> tokens, string value) : ExpressionNode(tokens, new NubPointerType(new NubIntType(true, 8)))
{
    public string Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class I8LiteralNode(List<Token> tokens, sbyte value) : ExpressionNode(tokens, new NubIntType(true, 8))
{
    public sbyte Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class I16LiteralNode(List<Token> tokens, short value) : ExpressionNode(tokens, new NubIntType(true, 16))
{
    public short Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class I32LiteralNode(List<Token> tokens, int value) : ExpressionNode(tokens, new NubIntType(true, 32))
{
    public int Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class I64LiteralNode(List<Token> tokens, long value) : ExpressionNode(tokens, new NubIntType(true, 64))
{
    public long Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class U8LiteralNode(List<Token> tokens, byte value) : ExpressionNode(tokens, new NubIntType(false, 8))
{
    public byte Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class U16LiteralNode(List<Token> tokens, ushort value) : ExpressionNode(tokens, new NubIntType(false, 16))
{
    public ushort Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class U32LiteralNode(List<Token> tokens, uint value) : ExpressionNode(tokens, new NubIntType(false, 32))
{
    public uint Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class U64LiteralNode(List<Token> tokens, ulong value) : ExpressionNode(tokens, new NubIntType(false, 64))
{
    public ulong Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class Float32LiteralNode(List<Token> tokens, float value) : ExpressionNode(tokens, new NubFloatType(32))
{
    public float Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class Float64LiteralNode(List<Token> tokens, double value) : ExpressionNode(tokens, new NubFloatType(64))
{
    public double Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class BoolLiteralNode(List<Token> tokens, bool value) : ExpressionNode(tokens, new NubBoolType())
{
    public bool Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class BinaryExpressionNode(List<Token> tokens, NubType type, ExpressionNode left, BinaryOperator @operator, ExpressionNode right) : ExpressionNode(tokens, type)
{
    public ExpressionNode Left { get; } = left;
    public BinaryOperator Operator { get; } = @operator;
    public ExpressionNode Right { get; } = right;

    public override IEnumerable<Node> Children()
    {
        yield return Left;
        yield return Right;
    }
}

public class UnaryExpressionNode(List<Token> tokens, NubType type, UnaryOperator @operator, ExpressionNode operand) : ExpressionNode(tokens, type)
{
    public UnaryOperator Operator { get; } = @operator;
    public ExpressionNode Operand { get; } = operand;

    public override IEnumerable<Node> Children()
    {
        yield return Operand;
    }
}

public class FuncCallNode(List<Token> tokens, NubType type, ExpressionNode expression, List<ExpressionNode> parameters) : ExpressionNode(tokens, type)
{
    public ExpressionNode Expression { get; } = expression;
    public List<ExpressionNode> Parameters { get; } = parameters;

    public override IEnumerable<Node> Children()
    {
        yield return Expression;
        foreach (var expressionNode in Parameters)
        {
            yield return expressionNode;
        }
    }
}

public class VariableIdentifierNode(List<Token> tokens, NubType type, IdentifierToken nameToken) : ExpressionNode(tokens, type)
{
    public IdentifierToken NameToken { get; } = nameToken;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class FuncIdentifierNode(List<Token> tokens, NubType type, IdentifierToken moduleToken, IdentifierToken nameToken, StringLiteralToken? externSymbolToken) : ExpressionNode(tokens, type)
{
    public IdentifierToken ModuleToken { get; } = moduleToken;
    public IdentifierToken NameToken { get; } = nameToken;
    public StringLiteralToken? ExternSymbolToken { get; } = externSymbolToken;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class ArrayInitializerNode(List<Token> tokens, NubType type, List<ExpressionNode> values) : ExpressionNode(tokens, type)
{
    public List<ExpressionNode> Values { get; } = values;

    public override IEnumerable<Node> Children()
    {
        return Values;
    }
}

public class ConstArrayInitializerNode(List<Token> tokens, NubType type, List<ExpressionNode> values) : ExpressionNode(tokens, type)
{
    public List<ExpressionNode> Values { get; } = values;

    public override IEnumerable<Node> Children()
    {
        return Values;
    }
}

public class ArrayIndexAccessNode(List<Token> tokens, NubType type, ExpressionNode target, ExpressionNode index) : ExpressionNode(tokens, type)
{
    public ExpressionNode Target { get; } = target;
    public ExpressionNode Index { get; } = index;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Index;
    }
}

public class ConstArrayIndexAccessNode(List<Token> tokens, NubType type, ExpressionNode target, ExpressionNode index) : ExpressionNode(tokens, type)
{
    public ExpressionNode Target { get; } = target;
    public ExpressionNode Index { get; } = index;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Index;
    }
}

public class SliceIndexAccessNode(List<Token> tokens, NubType type, ExpressionNode target, ExpressionNode index) : ExpressionNode(tokens, type)
{
    public ExpressionNode Target { get; } = target;
    public ExpressionNode Index { get; } = index;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
        yield return Index;
    }
}

public class AddressOfNode(List<Token> tokens, NubType type, ExpressionNode target) : ExpressionNode(tokens, type)
{
    public ExpressionNode Target { get; } = target;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
    }
}

public class StructFieldAccessNode(List<Token> tokens, NubType type, ExpressionNode target, IdentifierToken fieldToken) : ExpressionNode(tokens, type)
{
    public ExpressionNode Target { get; } = target;
    public IdentifierToken FieldToken { get; } = fieldToken;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
    }
}

public class StructInitializerNode(List<Token> tokens, NubType type, Dictionary<IdentifierToken, ExpressionNode> initializers) : ExpressionNode(tokens, type)
{
    public Dictionary<IdentifierToken, ExpressionNode> Initializers { get; } = initializers;

    public override IEnumerable<Node> Children()
    {
        foreach (var initializer in Initializers)
        {
            yield return initializer.Value;
        }
    }
}

public class DereferenceNode(List<Token> tokens, NubType type, ExpressionNode target) : ExpressionNode(tokens, type)
{
    public ExpressionNode Target { get; } = target;

    public override IEnumerable<Node> Children()
    {
        yield return Target;
    }
}

public class SizeNode(List<Token> tokens, NubType TargetType) : ExpressionNode(tokens, new NubIntType(false, 64))
{
    public NubType TargetType { get; } = TargetType;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

public class CastNode(List<Token> tokens, NubType type, ExpressionNode value) : ExpressionNode(tokens, type)
{
    public ExpressionNode Value { get; } = value;

    public override IEnumerable<Node> Children()
    {
        yield return Value;
    }
}

public abstract class IntermediateExpression(List<Token> tokens) : ExpressionNode(tokens, new NubVoidType());

public class EnumReferenceIntermediateNode(List<Token> tokens, IdentifierToken moduleToken, IdentifierToken nameToken) : IntermediateExpression(tokens)
{
    public IdentifierToken ModuleToken { get; } = moduleToken;
    public IdentifierToken NameToken { get; } = nameToken;

    public override IEnumerable<Node> Children()
    {
        return [];
    }
}

#endregion