import { Builtins } from "./arch.ts";
import { BinaryType, Expr, Stmt } from "./ast.ts";
import { LocalLeaf, Locals, LocalsFnRoot } from "./lowerer_locals.ts";
import { Ops } from "./mod.ts";
import { Assembler, Label } from "./assembler.ts";
import { VType, vtypeToString } from "./vtype.ts";

export class Lowerer {
    private program = new Assembler();
    private locals: Locals = new LocalsFnRoot();
    private fnStmtIdAddrMap: { [key: number]: number } = {};
    private breakStack: Label[] = [];

    public lower(stmts: Stmt[]) {
        for (const stmt of stmts) {
            this.lowerStaticStmt(stmt);
        }
    }

    public finish(): number[] {
        return this.program.assemble();
    }

    private lowerStaticStmt(stmt: Stmt) {
        switch (stmt.kind.type) {
            case "fn":
                return this.lowerFnStmt(stmt);
            case "error":
            case "break":
            case "return":
            case "let":
            case "assign":
            case "expr":
        }
        throw new Error(`unhandled static statement '${stmt.kind.type}'`);
    }

    private lowerStmt(stmt: Stmt) {
        switch (stmt.kind.type) {
            case "error":
                break;
            case "break":
                return this.lowerBreakStmt(stmt);
            case "return":
                break;
            case "fn":
                return this.lowerFnStmt(stmt);
            case "let":
                return this.lowerLetStmt(stmt);
            case "assign":
                return this.lowerAssignStmt(stmt);
            case "expr":
                this.lowerExpr(stmt.kind.expr);
                this.program.push(Ops.Pop);
                return;
        }
        throw new Error(`unhandled stmt '${stmt.kind.type}'`);
    }

    private lowerAssignStmt(stmt: Stmt) {
        if (stmt.kind.type !== "assign") {
            throw new Error();
        }
        this.lowerExpr(stmt.kind.value);
        switch (stmt.kind.subject.kind.type) {
            case "field": {
                this.lowerExpr(stmt.kind.subject.kind.subject);
                this.addPushStringOp(stmt.kind.subject.kind.value);
                this.program.addOp(Ops.Builtin);
                this.program.addLit(Builtins.StructSet);
                return;
            }
            case "index": {
                this.lowerExpr(stmt.kind.subject.kind.subject);
                this.lowerExpr(stmt.kind.subject.kind.value);
                this.program.addOp(Ops.Builtin);
                this.program.addLit(Builtins.ArraySet);
                return;
            }
            case "sym": {
                this.program.addOp(Ops.StoreLocal);
                this.program.addLit(
                    this.locals.symId(stmt.kind.subject.kind.sym.ident),
                );
                return;
            }
            default:
                throw new Error();
        }
    }

    private lowerBreakStmt(stmt: Stmt) {
        if (stmt.kind.type !== "break") {
            throw new Error();
        }
        if (stmt.kind.expr) {
            this.lowerExpr(stmt.kind.expr);
        }
        this.program.addOp(Ops.Jump);
        this.program.addRef(this.breakStack.at(-1)!);
    }

    private lowerFnStmt(stmt: Stmt) {
        if (stmt.kind.type !== "fn") {
            throw new Error();
        }
        const outerLocals = this.locals;
        this.locals = new LocalsFnRoot(outerLocals);
        const outerProgram = this.program;
        this.program = new Assembler();

        for (const { ident } of stmt.kind.params) {
            this.locals.allocSym(ident);
            this.program.push(
                Ops.StoreLocal,
                this.locals.symId(ident),
            );
        }
        this.lowerExpr(stmt.kind.body);
        this.program.push(Ops.Return);

        this.locals = outerLocals;
        outerProgram.concat(this.program);
        this.program = outerProgram;
    }

    private lowerLetStmt(stmt: Stmt) {
        if (stmt.kind.type !== "let") {
            throw new Error();
        }
        this.lowerExpr(stmt.kind.value);
        this.locals.allocSym(stmt.kind.param.ident),
            this.program.push(Ops.StoreLocal);
        this.program.push(this.locals.symId(stmt.kind.param.ident));
    }

    private lowerExpr(expr: Expr) {
        switch (expr.kind.type) {
            case "error":
                break;
            case "sym":
                return this.lowerSymExpr(expr);
            case "null":
                break;
            case "int":
                return this.lowerIntExpr(expr);
            case "bool":
                break;
            case "string":
                return this.lowerStringExpr(expr);
            case "ident":
                break;
            case "group":
                break;
            case "field":
                break;
            case "index":
                break;
            case "call":
                return this.lowerCallExpr(expr);
            case "unary":
                break;
            case "binary":
                return this.lowerBinaryExpr(expr);
            case "if":
                return this.lowerIfExpr(expr);
            case "loop":
                return this.lowerLoopExpr(expr);
            case "block":
                return this.lowerBlockExpr(expr);
        }
        throw new Error(`unhandled expr '${expr.kind.type}'`);
    }

    private lowerSymExpr(expr: Expr) {
        if (expr.kind.type !== "sym") {
            throw new Error();
        }
        if (expr.kind.sym.type === "let") {
            this.program.push(
                Ops.LoadLocal,
                this.locals.symId(expr.kind.ident),
            );
            return;
        }
        if (expr.kind.sym.type === "fn_param") {
            this.program.push(
                Ops.LoadLocal,
                this.locals.symId(expr.kind.ident),
            );
            return;
        }
        if (expr.kind.sym.type === "fn") {
            const addr = this.fnStmtIdAddrMap[expr.kind.sym.stmt.id];
            this.program.push(Ops.PushPtr, addr);
            return;
        }
        throw new Error(`unhandled sym type '${expr.kind.sym.type}'`);
    }

    private lowerIntExpr(expr: Expr) {
        if (expr.kind.type !== "int") {
            throw new Error();
        }
        this.program.push(Ops.PushInt, expr.kind.value);
    }

    private lowerStringExpr(expr: Expr) {
        if (expr.kind.type !== "string") {
            throw new Error();
        }
        this.addPushStringOp(expr.kind.value);
    }

    private addPushStringOp(value: string) {
        this.program.addOp(Ops.PushString);
        for (let i = 0; i < value.length; ++i) {
            this.program.addLit(value.charCodeAt(i));
        }
    }

    private lowerBinaryExpr(expr: Expr) {
        if (expr.kind.type !== "binary") {
            throw new Error();
        }
        this.lowerExpr(expr.kind.left);
        this.lowerExpr(expr.kind.right);
        const vtype = expr.kind.left.vtype!;
        if (vtype.type === "int") {
            switch (expr.kind.binaryType) {
                case "+":
                    this.program.addOp(Ops.Add);
                    return;
                case "*":
                    this.program.addOp(Ops.Multiply);
                    return;
                case "==":
                    this.program.addOp(Ops.Equal);
                    return;
                case ">=":
                    this.program.addOp(Ops.LessThan);
                    this.program.addOp(Ops.Not);
                    return;
                default:
            }
        }
        if (vtype.type === "string") {
            if (expr.kind.binaryType === "+") {
                this.program.push(Ops.Builtin, Builtins.StringConcat);
                return;
            }
            if (expr.kind.binaryType === "==") {
                this.program.push(Ops.Builtin, Builtins.StringEqual);
                return;
            }
            if (expr.kind.binaryType === "!=") {
                this.program.push(Ops.Builtin, Builtins.StringEqual, Ops.Not);
                return;
            }
        }
        throw new Error(
            `unhandled binaryType` +
                ` '${vtypeToString(expr.vtype!)}' aka. ` +
                ` '${vtypeToString(expr.kind.left.vtype!)}'` +
                ` ${expr.kind.binaryType}` +
                ` '${vtypeToString(expr.kind.left.vtype!)}'`,
        );
    }

    private lowerCallExpr(expr: Expr) {
        if (expr.kind.type !== "call") {
            throw new Error();
        }
        for (const arg of expr.kind.args) {
            this.lowerExpr(arg);
        }
        this.lowerExpr(expr.kind.subject);
    }

    private lowerIfExpr(expr: Expr) {
        if (expr.kind.type !== "if") {
            throw new Error();
        }

        const falseLabel = this.program.makeLabel();
        const doneLabel = this.program.makeLabel();

        this.lowerExpr(expr.kind.cond);

        this.program.push(Ops.Not, Ops.JumpIfTrue);
        this.program.addRef(falseLabel);

        this.lowerExpr(expr.kind.truthy);

        this.program.push(Ops.Jump);
        this.program.addRef(doneLabel);

        this.program.setLabel(falseLabel);

        if (expr.kind.falsy) {
            this.lowerExpr(expr.kind.falsy!);
        }

        this.program.setLabel(doneLabel);
    }

    private lowerLoopExpr(expr: Expr) {
        if (expr.kind.type !== "loop") {
            throw new Error();
        }
        const contineLabel = this.program.makeLabel();
        const breakLabel = this.program.makeLabel();

        this.breakStack.push(breakLabel);

        this.program.setLabel(contineLabel);
        this.lowerExpr(expr.kind.body);
        this.program.addOp(Ops.Jump);
        this.program.addRef(breakLabel);
        this.program.setLabel(breakLabel);
        if (expr.vtype!.type === "null") {
            this.program.addOp(Ops.PushNull);
        }

        this.breakStack.pop();
    }

    private lowerBlockExpr(expr: Expr) {
        if (expr.kind.type !== "block") {
            throw new Error();
        }
        const outerLocals = this.locals;
        this.locals = new LocalLeaf(this.locals);
        for (const stmt of expr.kind.stmts) {
            this.lowerStmt(stmt);
        }
        if (expr.kind.expr) {
            this.lowerExpr(expr.kind.expr);
        } else {
            this.program.addOp(Ops.PushNull);
        }
        this.locals = outerLocals;
    }

    public printProgram() {
        this.program.printProgram();
    }
}