import { MatrixClient, MsgType } from "matrix-js-sdk";
import type { ICallbackStore } from "../types.js";
import { config } from "../../config.js";
import { getTextGemini, getImageGemini } from "../../services/ai/ai.js";
import { alts } from "./alts.js";
import type { IAdminInstructions } from "./types.js";
import { getUserById, getUserName } from "../../helpers.js";
import { prices } from "./prices.js";

let client: MatrixClient;

const registerModuleAI = (
    matrixClient: MatrixClient,
    callbackStore: ICallbackStore,
) => {
    client = matrixClient;

    callbackStore.messageCallbacks.push({
        startConditions: [`${config.app.triggerPrefix}ai `],
        callbackFunc: onAI,
        allowedRoles: ["USER", "MODERATOR", "ADMIN"],
    });
    callbackStore.messageCallbacks.push({
        startConditions: [`${config.app.triggerPrefix}img `],
        callbackFunc: onImageGen,
        allowedRoles: ["USER", "MODERATOR", "ADMIN"],
    });
};

const onAI = async (text: string, roomId: string, sender: string) => {
    const user = getUserById(sender);

    if (text.trim().length < 5) {
        return;
    }

    let textMod = text.replace("!ai", "").trim().toLowerCase();
    let instructions = {
        prefferedLanguages: ["english", "slovak"],
        adminText:
            "Your name is Aslobot. Be concise, try to keep text as short as possible. Be helpful, but you are encouraged to be sarcastic and make fun of people",
        alts: alts,
    } as IAdminInstructions;

    const username = getUserName(user);
    textMod = `Admin Instructions:\n${JSON.stringify(instructions)}\nPrompt by ${username}:\n${textMod}`;

    const responseAI = await getTextGemini(textMod);

    user.aiCost += responseAI.tokens * prices.text;

    client.sendTextMessage(roomId, responseAI.text);
};

const onImageGen = async (text: string, roomId: string, sender: string) => {
    const user = getUserById(sender);

    let textMod = text.replace("!img", "").trim().toLowerCase();
    alts.forEach((alt) => {
        alt.keys.forEach((key) => {
            textMod = textMod.replaceAll(key, alt.alt);
        });
    });

    const responseAI = await getImageGemini(textMod);

    user.aiCost += responseAI.tokens * prices.image;
    if (!responseAI.image || responseAI.image.length < 10) {
        return;
    }

    const imageName = `photo-img-gen.png`;

    const uploadResult = await client.uploadContent(responseAI.image, {
        type: "image/png",
        name: imageName,
    });

    await client.sendMessage(roomId, {
        msgtype: MsgType.Image,
        body: imageName,
        url: uploadResult.content_uri,
        info: {
            mimetype: "image/png",
            size: responseAI.image.length,
        },
    });
};

export { registerModuleAI, onImageGen };