export class Node {
    private static idCounter = 0;

    static create<Tag extends NodeKind["tag"]>(
        line: number,
        tag: Tag,
        kind: Omit<NodeKind & { tag: Tag }, "tag">,
    ): Node {
        return new Node(
            Node.idCounter++,
            line,
            { tag, ...kind } as NodeKind & { tag: Tag },
        );
    }

    private constructor(
        public id: number,
        public line: number,
        public kind: NodeKind,
    ) {}

    visit(v: Visitor) {
        if (v.visit(this) === "break") {
            return;
        }
        this.visitBelow(v);
    }

    visitBelow(v: Visitor) {
        const visit = (...nodes: (Node | null)[]) => {
            for (const node of nodes) {
                node?.visit(v);
            }
        };

        const k = this.kind;
        switch (k.tag) {
            case "Error":
                return visit();
            case "File":
                return visit(...k.stmts);
            case "Block":
                return visit(...k.stmts);
            case "ExprStmt":
                return visit(k.expr);
            case "AssignStmt":
                return visit(k.place, k.expr);
            case "FnStmt":
                return visit(...k.params, k.retTy, k.body);
            case "ReturnStmt":
                return visit(k.expr);
            case "LetStmt":
                return visit(k.param, k.expr);
            case "Param":
                return visit(k.ty);
            case "IdentExpr":
                return visit();
            case "IntExpr":
                return visit();
            case "CallExpr":
                return visit(k.expr, ...k.args);
            case "IdentTy":
                return visit();
        }
        const _: never = k;
    }
}

export type NodeKind =
    | { tag: "Error" }
    | { tag: "File"; stmts: Node[] }
    | { tag: "Block"; stmts: Node[] }
    | { tag: "ExprStmt"; expr: Node }
    | { tag: "AssignStmt"; place: Node; expr: Node }
    | {
        tag: "FnStmt";
        ident: string;
        params: Node[];
        retTy: Node | null;
        body: Node;
    }
    | { tag: "ReturnStmt"; expr: Node | null }
    | { tag: "LetStmt"; param: Node; expr: Node }
    | { tag: "Param"; ident: string; ty: Node | null }
    | { tag: "IdentExpr"; ident: string }
    | { tag: "IntExpr"; value: number }
    | { tag: "CallExpr"; expr: Node; args: Node[] }
    | { tag: "IdentTy"; ident: string };

export interface Visitor {
    visit(node: Node): void | "break";
}

export type NodeWithKind<
    Tag extends NodeKind["tag"],
> = Node & { kind: { tag: Tag } };

export function assertNodeWithKind<Tag extends NodeKind["tag"]>(
    node: Node,
    tag: Tag,
): asserts node is NodeWithKind<Tag> {
    if (node.kind.tag !== tag) {
        throw new Error();
    }
}