import { AstCreator, Expr, ExprKind, Stmt, StmtKind } from "../ast.ts";
import { AstVisitor, visitExpr, VisitRes, visitStmts } from "../ast_visitor.ts";
import { Pos } from "../token.ts";

export class SpecialLoopDesugarer implements AstVisitor {
    public constructor(
        private astCreator: AstCreator,
    ) {}

    public desugar(stmts: Stmt[]) {
        visitStmts(stmts, this);
    }

    visitWhileExpr(expr: Expr): VisitRes {
        if (expr.kind.type !== "while") {
            throw new Error();
        }
        const npos: Pos = { index: 0, line: 1, col: 1 };
        const Expr = (kind: ExprKind, pos = npos) =>
            this.astCreator.expr(kind, pos);
        const Stmt = (kind: StmtKind, pos = npos) =>
            this.astCreator.stmt(kind, pos);

        expr.kind = {
            type: "loop",
            body: Expr({
                type: "block",
                stmts: [
                    Stmt({
                        type: "expr",
                        expr: Expr({
                            type: "if",
                            cond: Expr({
                                type: "unary",
                                unaryType: "not",
                                subject: expr.kind.cond,
                            }),
                            truthy: Expr({
                                type: "block",
                                stmts: [
                                    Stmt({ type: "break" }),
                                ],
                            }),
                        }),
                    }),
                    Stmt({
                        type: "expr",
                        expr: expr.kind.body,
                    }),
                ],
            }),
        };

        visitExpr(expr, this);
        return "stop";
    }

    visitForInExpr(expr: Expr): VisitRes {
        if (expr.kind.type !== "for_in") {
            throw new Error();
        }
        const npos: Pos = { index: 0, line: 1, col: 1 };
        const Expr = (kind: ExprKind, pos = npos) =>
            this.astCreator.expr(kind, pos);
        const Stmt = (kind: StmtKind, pos = npos) =>
            this.astCreator.stmt(kind, pos);

        expr.kind = {
            type: "block",
            stmts: [
                Stmt({
                    type: "let",
                    param: this.astCreator.param({
                        ident: "::values",
                        mut: true,
                        pos: npos,
                    }),
                    value: expr.kind.value,
                }),
                Stmt({
                    type: "let",
                    param: this.astCreator.param({
                        ident: "::length",
                        mut: false,
                        pos: npos,
                    }),
                    value: Expr({
                        type: "call",
                        subject: Expr({
                            type: "path",
                            subject: Expr({
                                type: "ident",
                                ident: "std",
                            }),
                            ident: "array_length",
                        }),
                        args: [
                            Expr({
                                type: "ident",
                                ident: "::values",
                            }),
                        ],
                    }),
                }),
                Stmt({
                    type: "let",
                    param: this.astCreator.param({
                        ident: "::index",
                        mut: true,
                        pos: npos,
                    }),
                    value: Expr({ type: "int", value: 0 }),
                }, expr.pos),
                Stmt({
                    type: "expr",
                    expr: Expr({
                        type: "loop",
                        body: Expr({
                            type: "block",
                            stmts: [
                                Stmt({
                                    type: "expr",
                                    expr: Expr({
                                        type: "if",
                                        cond: Expr({
                                            type: "unary",
                                            unaryType: "not",
                                            subject: Expr({
                                                type: "binary",
                                                binaryType: "<",
                                                left: Expr({
                                                    type: "ident",
                                                    ident: "::index",
                                                }),
                                                right: Expr({
                                                    type: "ident",
                                                    ident: "::length",
                                                }),
                                            }),
                                        }),
                                        truthy: Expr({
                                            type: "block",
                                            stmts: [
                                                Stmt({
                                                    type: "break",
                                                }),
                                            ],
                                        }),
                                    }),
                                }),
                                Stmt({
                                    type: "let",
                                    param: expr.kind.param,
                                    value: Expr({
                                        type: "index",
                                        subject: Expr({
                                            type: "ident",
                                            ident: "::values",
                                        }),
                                        value: Expr({
                                            type: "ident",
                                            ident: "::index",
                                        }),
                                    }),
                                }, expr.pos),
                                Stmt({
                                    type: "expr",
                                    expr: expr.kind.body,
                                }),
                                Stmt({
                                    type: "assign",
                                    assignType: "+=",
                                    subject: Expr({
                                        type: "ident",
                                        ident: "::index",
                                    }),
                                    value: Expr({
                                        type: "int",
                                        value: 1,
                                    }),
                                }),
                            ],
                        }),
                    }),
                }),
            ],
        };

        visitExpr(expr, this);
        return "stop";
    }

    visitForExpr(expr: Expr): VisitRes {
        if (expr.kind.type !== "for") {
            throw new Error();
        }
        const Expr = (
            kind: ExprKind,
            pos: Pos = { index: 0, line: 1, col: 1 },
        ) => this.astCreator.expr(kind, pos);
        const Stmt = (
            kind: StmtKind,
            pos: Pos = { index: 0, line: 1, col: 1 },
        ) => this.astCreator.stmt(kind, pos);

        expr.kind = {
            type: "block",
            stmts: [
                ...[expr.kind.decl]
                    .filter((v) => v !== undefined),
                Stmt({
                    type: "expr",
                    expr: Expr({
                        type: "loop",
                        body: Expr({
                            type: "block",
                            stmts: [
                                ...[expr.kind.cond]
                                    .filter((v) => v !== undefined)
                                    .map(
                                        (cond) =>
                                            Stmt({
                                                type: "expr",
                                                expr: Expr({
                                                    type: "if",
                                                    cond: Expr({
                                                        type: "unary",
                                                        unaryType: "not",
                                                        subject: cond,
                                                    }),
                                                    truthy: Expr({
                                                        type: "block",
                                                        stmts: [
                                                            Stmt({
                                                                type: "break",
                                                            }),
                                                        ],
                                                    }),
                                                }),
                                            }, cond.pos),
                                    ),
                                Stmt({
                                    type: "expr",
                                    expr: expr.kind.body,
                                }),
                                ...[expr.kind.incr]
                                    .filter((v) => v !== undefined),
                            ],
                        }),
                    }),
                }),
            ],
        };

        visitExpr(expr, this);
        return "stop";
    }
}