const std = @import("std");
const tokens = @import("tokens.zig");
const Token = tokens.Token;

pub const NodeList = std.ArrayList(Node);
pub const StmtList = std.ArrayList(Stmt);
pub const Block = std.ArrayList(Stmt);
pub const ExprList = std.ArrayList(Expr);
pub const TokenList = std.ArrayList(Token);
pub const ParamList = std.ArrayList(ParamDecl);
pub const ConstList = std.ArrayList(SingleConst);

pub const NodeType = enum {
    Root,
    FnDecl,
    ConstDecl,
    Struct,
    Enum,
    Block,
    Stmt,
};

pub const ParamDecl = struct {
    name: Token,
    typ: Token,
};

pub const MethodData = struct {
    variable: Token,
    typ: Token,
    mutable: bool,
};

pub const FnDecl = struct {
    func_name: Token,
    params: ParamList,
    return_type: Token,
    body: StmtList,
    method: ?*MethodData,
};

pub const SingleConst = struct {
    name: Token,
    expr: *Expr,
};

pub const BinaryOperator = enum {

    // numeric
    Add,
    Sub,
    Mul,
    Div,
    Mod,

    // numeric -> bool
    Greater,
    GreaterEqual,
    Less,
    LessEqual,

    // booleans
    Equal,
    And,
    Or,
};

pub const BinaryExpr = struct {
    left: *Expr,
    op: BinaryOperator,
    right: *Expr,
};

pub const UnaryOperator = enum {
    Not,
    Negate,
};

pub const UnaryExpr = struct {
    op: UnaryOperator,
    right: *Expr,
};

pub const LiteralExpr = union(enum) {
    Bool: bool,
    Integer: []const u8,
    Float: []const u8,
    String: []const u8,
    Array: ExprList,
};

pub const AssignExpr = struct {
    name: Token,
    value: *Expr,
};

pub const ExprType = enum {
    Assign,

    Binary,
    Unary,
    Literal,
    Variable,
    Call,
    Struct,

    Grouping,
    Get,
    Set,
};

pub const CallExpr = struct {
    callee: *Expr,
    paren: Token,
    arguments: ExprList,
};

pub const StructInit = struct {
    field: Token,
    expr: *Expr,
};

pub const StructInitList = std.ArrayList(StructInit);

pub const StructExpr = struct {
    name: Token,
    inits: StructInitList,
};

pub const GetExpr = struct {
    struc: *Expr,
    name: Token,
};

pub const SetExpr = struct {
    struc: *Expr,
    field: Token,
    value: *Expr,
};

pub const Expr = union(ExprType) {
    Assign: AssignExpr,

    Binary: BinaryExpr,
    Unary: UnaryExpr,
    Literal: LiteralExpr,
    Struct: StructExpr,

    Variable: Token,
    Grouping: *Expr,
    Call: CallExpr,

    Get: GetExpr,
    Set: SetExpr,
};

pub const IfStmt = struct {
    condition: *Expr,
    then_branch: Block,
    else_branch: ?Block,
};

pub const LoopStmt = struct {
    condition: ?*Expr,
    then_branch: Block,
};

pub const ForStmt = struct {
    index: ?Token,
    value: Token,
    array: Token,
    block: Block,
};

pub const VarDeclStmt = struct {
    name: Token,
    value: *Expr,
};

pub const Stmt = union(enum) {
    Expr: *Expr,
    Println: *Expr,

    VarDecl: VarDeclStmt,
    If: IfStmt,
    Loop: LoopStmt,
    For: ForStmt,

    Return: ReturnStmt,

    pub const ReturnStmt = struct {
        keyword: Token,
        value: *Expr,
    };

    pub fn mkPrintln(allocator: *std.mem.Allocator, expr: *Expr) !*Stmt {
        var stmt = try allocator.create(Stmt);
        stmt.* = Stmt{ .Println = expr };
        return stmt;
    }

    pub fn mkIfStmt(
        allocator: *std.mem.Allocator,
        condition: *Expr,
        then: Block,
        else_branch: ?Block,
    ) !*Stmt {
        var stmt = try allocator.create(Stmt);
        stmt.* = Stmt{
            .If = IfStmt{
                .condition = condition,
                .then_branch = then,
                .else_branch = else_branch,
            },
        };

        return stmt;
    }

    pub fn mkLoop(
        allocator: *std.mem.Allocator,
        condition: ?*Expr,
        then: Block,
    ) !*Stmt {
        var stmt = try allocator.create(Stmt);
        stmt.* = Stmt{
            .Loop = LoopStmt{
                .condition = condition,
                .then_branch = then,
            },
        };

        return stmt;
    }

    pub fn mkFor(allocator: *std.mem.Allocator, index: ?Token, value: Token, array: Token, block: Block) !*Stmt {
        var stmt = try allocator.create(Stmt);
        stmt.* = Stmt{
            .For = ForStmt{
                .index = index,
                .value = value,
                .array = array,
                .block = block,
            },
        };

        return stmt;
    }

    pub fn mkReturn(allocator: *std.mem.Allocator, tok: Token, value: *Expr) !*Stmt {
        var stmt = try allocator.create(Stmt);
        stmt.* = Stmt{
            .Return = ReturnStmt{
                .keyword = tok,
                .value = value,
            },
        };

        return stmt;
    }

    pub fn mkVarDecl(allocator: *std.mem.Allocator, name: Token, value: *Expr) !*Stmt {
        var stmt = try allocator.create(Stmt);
        stmt.* = Stmt{
            .VarDecl = VarDeclStmt{
                .name = name,
                .value = value,
            },
        };

        return stmt;
    }
};

pub const FieldList = std.ArrayList(StructField);

pub const StructField = struct {
    name: Token,
    typ: Token,

    mutable: bool = false,
    public: bool = false,
    mutable_outside: bool = false,
};

pub const Struct = struct {
    name: Token,
    fields: FieldList,
};

pub const Enum = struct {
    name: Token,

    // TODO allow custom values for the enum fields
    // just like c enums can skip values
    fields: TokenList,
};

pub const Node = union(NodeType) {
    Root: NodeList,
    FnDecl: FnDecl,
    ConstDecl: ConstList,
    Struct: Struct,
    Enum: Enum,

    Block: StmtList,

    Stmt: *Stmt,

    pub fn mkRoot(allocator: *std.mem.Allocator) !*Node {
        var node = try allocator.create(Node);
        node.* = Node{ .Root = NodeList.init(allocator) };
        return node;
    }

    pub fn mkFnDecl(
        allocator: *std.mem.Allocator,
        name: Token,
        params: ParamList,
        return_type: Token,
        block: StmtList,
        method: ?*MethodData,
    ) !*Node {
        var node = try allocator.create(Node);

        node.* = Node{
            .FnDecl = FnDecl{
                .func_name = name,
                .params = params,
                .return_type = return_type,
                .body = block,
                .method = method,
            },
        };
        return node;
    }

    pub fn mkStructDecl(allocator: *std.mem.Allocator, name: Token, fields: FieldList) !*Node {
        var node = try allocator.create(Node);
        node.* = Node{
            .Struct = Struct{
                .name = name,
                .fields = fields,
            },
        };

        return node;
    }

    pub fn mkEnumDecl(allocator: *std.mem.Allocator, name: Token, fields: TokenList) !*Node {
        var node = try allocator.create(Node);
        node.* = Node{
            .Enum = Enum{
                .name = name,
                .fields = fields,
            },
        };

        return node;
    }
};