import { writeFileSync } from "node:fs";
import * as t from "@babel/types";
import type * as $ from "@/asm/generator/instructions";
import type { Opcode } from "@/asm/generator/instructions";
import {
    instructionList,
    pseudoInstructions,
} from "@/asm/generator/instructions";
import generateTs from "@babel/generator";

const RUNTIME_QUALIFIER = t.identifier("c");
const UTIL_QUALIFIER = t.identifier("$");
const AST_QUALIFIER = t.identifier("$ast");

const main = () => {
    writeFileSync(`${__dirname}/../text/convert.ts`, generate());
};

const generate = (): string => {
    const importUtil = t.importDeclaration(
        [t.importNamespaceSpecifier(UTIL_QUALIFIER)],
        t.stringLiteral("./util"),
    );

    const importConstructors = t.importDeclaration(
        [t.importNamespaceSpecifier(RUNTIME_QUALIFIER)],
        t.stringLiteral("../runtime"),
    );

    const importAst = t.importDeclaration(
        [t.importSpecifier(AST_QUALIFIER, AST_QUALIFIER)],
        t.stringLiteral("./grammar"),
    );

    const stmts = instructionList().flatMap(([name, instruction]) => {
        if (pseudoInstructions.has(name)) {
            return [];
        }
        return generateInstr(name, instruction);
    });

    const convertorFunc = generateConvertorFunction(instructionList());

    const file = t.file(
        t.program([
            importUtil,
            importConstructors,
            importAst,
            ...stmts,
            convertorFunc,
        ]),
    );
    t.addComment(
        file,
        "leading",
        ` AUTOGENERATED, DO NOT EDIT, generated by ../generator/gen-converter.ts`,
        true,
    );
    return generateTs(file).code;
};

function generateConvertorFunction(
    instructions: [string, Opcode][],
): t.ExportNamedDeclaration {
    const nameIdent = t.identifier("convertInstruction");

    const instrParam = t.identifier("instr");
    instrParam.typeAnnotation = t.tsTypeAnnotation(
        t.tsTypeReference(
            t.tsQualifiedName(AST_QUALIFIER, t.identifier("Instruction")),
        ),
    );

    const ctxParam = t.identifier("ctx");
    ctxParam.typeAnnotation = t.tsTypeAnnotation(
        t.tsTypeReference(
            t.tsQualifiedName(UTIL_QUALIFIER, t.identifier("Ctx")),
        ),
    );

    const locParam = t.identifier("loc");
    locParam.typeAnnotation = t.tsTypeAnnotation(
        t.tsTypeReference(
            t.tsQualifiedName(
                t.tsQualifiedName(RUNTIME_QUALIFIER, t.identifier("util")),
                t.identifier("Loc"),
            ),
        ),
    );

    const params = [ctxParam, instrParam, locParam];

    const errorMsg = t.templateLiteral(
        [
            t.templateElement({ raw: "Unexpected instruction: " }),
            t.templateElement({ raw: "" }),
        ],
        [t.identifier("name")],
    );

    const body = t.blockStatement([
        t.variableDeclaration("const", [
            t.variableDeclarator(
                t.identifier("name"),
                t.memberExpression(
                    t.memberExpression(
                        t.identifier("instr"),
                        t.identifier("name"),
                    ),
                    t.identifier("name"),
                ),
            ),
        ]),

        t.switchStatement(
            t.identifier("name"),
            instructions.flatMap(([name]) => {
                if (pseudoInstructions.has(name)) {
                    return [];
                }

                return [
                    t.switchCase(t.stringLiteral(name), [
                        t.returnStatement(
                            t.callExpression(t.identifier(name), [
                                t.identifier("ctx"),
                                t.identifier("instr"),
                                t.identifier("loc"),
                            ]),
                        ),
                    ]),
                ];
            }),
        ),

        t.throwStatement(t.newExpression(t.identifier("Error"), [errorMsg])),
    ]);

    return t.exportNamedDeclaration(
        t.variableDeclaration("const", [
            t.variableDeclarator(
                nameIdent,
                t.arrowFunctionExpression(params, body),
            ),
        ]),
    );
}

function generateBody(name: string, instruction: $.Opcode) {
    const len = argsLen(instruction.args);

    const argsConst = t.variableDeclaration("const", [
        t.variableDeclarator(
            t.identifier("argsLen"),
            t.memberExpression(t.identifier("instr"), t.identifier("args")),
        ),
    ]);

    const cond = t.binaryExpression(
        "!==",
        t.memberExpression(t.identifier("argsLen"), t.identifier("length")),
        t.numericLiteral(len),
    );

    const argumentsStr = len === 1 ? "argument" : "arguments";
    const messageLiteral = t.stringLiteral(`Expected ${len} ${argumentsStr}`);
    const lenCheck = t.ifStatement(
        cond,
        t.blockStatement([
            t.throwStatement(
                t.newExpression(
                    t.memberExpression(
                        UTIL_QUALIFIER,
                        t.identifier("ParseError"),
                    ),
                    [t.identifier("loc"), messageLiteral],
                ),
            ),
        ]),
    );

    const argsStatements = generateArgs(name, instruction.args);

    return t.blockStatement([argsConst, lenCheck, ...argsStatements]);
}

const argsLen = (args: $.args): number => {
    switch (args.$) {
        case "simpleArgs":
            return args.children.length;
        case "xchgArgs":
            return 2;
        case "dictpush":
            return 2;
    }

    throw new Error("Unexpected arg type");
};

function isIntegerArg(arg: $.arg) {
    return (
        arg.$ === "int" ||
        arg.$ === "uint" ||
        arg.$ === "refs" ||
        arg.$ === "plduzArg" ||
        arg.$ === "tinyInt" ||
        arg.$ === "runvmArg" ||
        arg.$ === "minusOne" ||
        arg.$ === "delta" ||
        arg.$ === "hash" || // TODO: separate type
        arg.$ === "setcpArg"
    );
}

function isCodeArg(arg: $.arg) {
    return (
        arg.$ === "codeSlice" ||
        arg.$ === "inlineCodeSlice" ||
        arg.$ === "refCodeSlice"
    );
}

function isStackArg(arg: $.arg) {
    if (arg.$ === "stack" || arg.$ === "s1") return true;
    if (arg.$ === "delta") {
        return isStackArg(arg.arg);
    }
    return false;
}

function isControlArg(arg: $.arg) {
    if (arg.$ === "control") return true;
    if (arg.$ === "delta") {
        return isStackArg(arg.arg);
    }
    return false;
}

function parseArgs(name: string, withCtx: boolean) {
    return t.variableDeclaration("const", [
        t.variableDeclarator(
            t.identifier("args"),
            t.callExpression(
                t.memberExpression(UTIL_QUALIFIER, t.identifier(name)),
                [
                    ...(withCtx ? [t.identifier("ctx")] : []),
                    t.identifier("instr"),
                ],
            ),
        ),
    ]);
}

function returnTwoArgInstr(name: string) {
    return t.returnStatement(
        t.callExpression(
            t.memberExpression(RUNTIME_QUALIFIER, t.identifier(name)),
            [
                t.memberExpression(
                    t.identifier("args"),
                    t.numericLiteral(0),
                    true,
                ),
                t.memberExpression(
                    t.identifier("args"),
                    t.numericLiteral(1),
                    true,
                ),
                t.identifier("loc"),
            ],
        ),
    );
}

function returnSingleArgInstr(name: string) {
    return t.returnStatement(
        t.callExpression(
            t.memberExpression(RUNTIME_QUALIFIER, t.identifier(name)),
            [t.identifier("args"), t.identifier("loc")],
        ),
    );
}

function returnThreeArgInstr(name: string) {
    return t.returnStatement(
        t.callExpression(
            t.memberExpression(RUNTIME_QUALIFIER, t.identifier(name)),
            [
                t.memberExpression(
                    t.identifier("args"),
                    t.numericLiteral(0),
                    true,
                ),
                t.memberExpression(
                    t.identifier("args"),
                    t.numericLiteral(1),
                    true,
                ),
                t.memberExpression(
                    t.identifier("args"),
                    t.numericLiteral(2),
                    true,
                ),
                t.identifier("loc"),
            ],
        ),
    );
}

const generateSimpleArgs = (name: string, args: $.arg[]): t.Statement[] => {
    if (args.length === 0) {
        return [
            t.returnStatement(
                t.callExpression(
                    t.memberExpression(RUNTIME_QUALIFIER, t.identifier(name)),
                    [t.identifier("loc")],
                ),
            ),
        ];
    }

    if (args.length === 1) {
        const arg = args[0];
        if (!arg) {
            throw new Error(
                `Unexpected args count, expected 1, got ${args.length}`,
            );
        }

        if (isIntegerArg(arg)) {
            return [
                parseArgs("singleIntegerArg", false),
                returnSingleArgInstr(name),
            ];
        }

        if (isStackArg(arg)) {
            return [
                parseArgs("singleStackArg", false),
                returnSingleArgInstr(name),
            ];
        }

        if (isControlArg(arg)) {
            return [
                parseArgs("singleControlArg", false),
                returnSingleArgInstr(name),
            ];
        }

        if (arg.$ === "largeInt") {
            return [
                parseArgs("singleBigIntArg", false),
                returnSingleArgInstr(name),
            ];
        }

        if (isCodeArg(arg)) {
            return [
                parseArgs("codeSliceArg", true),
                returnSingleArgInstr(name),
            ];
        }

        if (arg.$ === "slice") {
            return [parseArgs("sliceArg", false), returnSingleArgInstr(name)];
        }

        if (arg.$ === "debugstr") {
            return [
                parseArgs("debugstrArg", false),
                returnSingleArgInstr(name),
            ];
        }
    }

    if (args.length === 2) {
        const arg = args[0];
        const arg1 = args[1];

        if (!arg || !arg1) {
            throw new Error(
                `Unexpected args count, expected 2, got ${args.length}`,
            );
        }

        if (isIntegerArg(arg) && isIntegerArg(arg1)) {
            return [
                parseArgs("twoIntegerArgs", false),
                returnTwoArgInstr(name),
            ];
        }

        if (isStackArg(arg) && isStackArg(arg1)) {
            return [parseArgs("twoStackArgs", false), returnTwoArgInstr(name)];
        }

        if (isCodeArg(arg) && isCodeArg(arg1)) {
            return [
                parseArgs("twoCodeSliceArgs", true),
                returnTwoArgInstr(name),
            ];
        }

        if (isIntegerArg(arg) && isCodeArg(arg1)) {
            return [parseArgs("ifBitArgs", true), returnTwoArgInstr(name)];
        }
    }

    if (args.length === 3) {
        const arg = args[0];
        const arg1 = args[1];
        const arg2 = args[2];

        if (!arg || !arg1 || !arg2) {
            throw new Error(
                `Unexpected args count, expected 3, got ${args.length}`,
            );
        }

        if (isIntegerArg(arg) && isIntegerArg(arg1) && isIntegerArg(arg2)) {
            return [
                parseArgs("threeIntegerArgs", false),
                returnThreeArgInstr(name),
            ];
        }

        if (isStackArg(arg) && isStackArg(arg1) && isStackArg(arg2)) {
            return [
                parseArgs("threeStackArgs", false),
                returnThreeArgInstr(name),
            ];
        }

        if (isIntegerArg(arg) && isCodeArg(arg1) && isCodeArg(arg2)) {
            return [
                parseArgs("ifElseBitArgs", false),
                returnThreeArgInstr(name),
            ];
        }
    }

    return [];
};

const generateDictpush = (name: string): t.Statement[] => {
    return [
        t.variableDeclaration("const", [
            t.variableDeclarator(
                t.identifier("args"),
                t.callExpression(
                    t.memberExpression(
                        UTIL_QUALIFIER,
                        t.identifier("dictpushArg"),
                    ),
                    [t.identifier("ctx"), t.identifier("instr")],
                ),
            ),
        ]),
        t.returnStatement(
            t.callExpression(
                t.memberExpression(RUNTIME_QUALIFIER, t.identifier(name)),
                [
                    t.memberExpression(
                        t.identifier("args"),
                        t.numericLiteral(0),
                        true,
                    ),
                    t.memberExpression(
                        t.identifier("args"),
                        t.numericLiteral(1),
                        true,
                    ),
                    t.identifier("loc"),
                ],
            ),
        ),
    ];
};

const generateXchgArgs = (name: string): t.Statement[] => {
    return [
        t.variableDeclaration("const", [
            t.variableDeclarator(
                t.identifier("args"),
                t.callExpression(
                    t.memberExpression(
                        UTIL_QUALIFIER,
                        t.identifier("twoStackArgs"),
                    ),
                    [t.identifier("instr")],
                ),
            ),
        ]),
        t.returnStatement(
            t.callExpression(
                t.memberExpression(RUNTIME_QUALIFIER, t.identifier(name)),
                [
                    t.memberExpression(
                        t.identifier("args"),
                        t.numericLiteral(0),
                        true,
                    ),
                    t.memberExpression(
                        t.identifier("args"),
                        t.numericLiteral(1),
                        true,
                    ),
                    t.identifier("loc"),
                ],
            ),
        ),
    ];
};

const generateArgs = (name: string, args: $.args): t.Statement[] => {
    switch (args.$) {
        case "simpleArgs":
            return generateSimpleArgs(name, args.children);
        case "xchgArgs":
            return generateXchgArgs(name);
        case "dictpush":
            return generateDictpush(name);
    }

    throw new Error("Unexpected arg type");
};

const generateInstr = (name: string, instruction: $.Opcode): t.Statement[] => {
    const nameIdent = t.identifier(name);
    nameIdent.typeAnnotation = t.tsTypeAnnotation(
        t.tsTypeReference(
            t.tsQualifiedName(UTIL_QUALIFIER, t.identifier("Convert")),
        ),
    );

    const params = [
        t.identifier("ctx"),
        t.identifier("instr"),
        t.identifier("loc"),
    ];
    const lambda = t.arrowFunctionExpression(
        params,
        generateBody(name, instruction),
    );

    const constructor = t.exportNamedDeclaration(
        t.variableDeclaration("const", [
            t.variableDeclarator(nameIdent, lambda),
        ]),
    );

    return [constructor];
};

main();
