nyx/bot/mod.ts

import { Api, Bot, Context, RawApi, session, SessionFlavor } from "grammy";
import { FileFlavor, hydrateFiles } from "grammy_files";
import { hydrateReply, ParseModeFlavor } from "grammy_parse_mode";
import { sequentialize } from "grammy_runner";
import { error, info, warning } from "std/log/mod.ts";
import { sessions } from "../api/sessionsRoute.ts";
import { formatUserChat } from "../utils/formatUserChat.ts";
import { omitUndef } from "../utils/omitUndef.ts";
import { broadcastCommand } from "./broadcastCommand.ts";
import { cancelCommand } from "./cancelCommand.ts";
import { img2imgCommand, img2imgQuestion } from "./img2imgCommand.ts";
import { pnginfoCommand, pnginfoQuestion } from "./pnginfoCommand.ts";
import { queueCommand } from "./queueCommand.ts";
import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts";
import { setConfig, getConfig } from "../app/config.ts";

// Set the new configuration
await setConfig({ maxUserJobs: 1, maxJobs: 500 });

// Fetch the updated configuration
const updatedConfig = await getConfig();

// Log the updated configuration to the console
console.log("Updated Configuration:", updatedConfig);

interface SessionData {
  chat: ErisChatData;
  user: ErisUserData;
}

interface ErisChatData {
  language?: string | undefined;
}

interface ErisUserData {
  params?: Record<string, string> | undefined;
}

export type ErisContext =
  & FileFlavor<ParseModeFlavor<Context>>
  & SessionFlavor<SessionData>;

type WithRetryApi<T extends RawApi> = {
  [M in keyof T]: T[M] extends (args: infer P, ...rest: infer A) => infer R
    ? (args: P extends object ? P & { maxAttempts?: number; maxWait?: number } : P, ...rest: A) => R
    : T[M];
};

type ErisApi = Api<WithRetryApi<RawApi>>;

export const bot = new Bot<ErisContext, ErisApi>(
  Deno.env.get("TG_BOT_TOKEN")!,
  {
    client: { timeoutSeconds: 20 },
  },
);

bot.use(hydrateReply);

bot.use(sequentialize((ctx) => ctx.chat?.id.toString()));

bot.use(session<SessionData, ErisContext>({
  type: "multi",
  chat: {
    initial: () => ({}),
  },
  user: {
    getSessionKey: (ctx) => ctx.from?.id.toFixed(),
    initial: () => ({}),
  },
}));

bot.api.config.use(hydrateFiles(bot.token));

// Automatically retry bot requests if we get a "too many requests" or telegram internal error
bot.api.config.use(async (prev, method, payload, signal) => {
  const maxAttempts = payload && ("maxAttempts" in payload) ? payload.maxAttempts ?? 3 : 3;
  const maxWait = payload && ("maxWait" in payload) ? payload.maxWait ?? 10 : 10;
  let attempt = 0;
  while (true) {
    attempt++;
    const result = await prev(method, payload, signal);
    if (result.ok) return result;
    if (result.error_code !== 429) return result;
    if (attempt >= maxAttempts) return result;
    const retryAfter = result.parameters?.retry_after ?? (attempt * 5);
    if (retryAfter > maxWait) return result;
    warning(
      `${method} (attempt ${attempt}) failed: ${result.error_code} ${result.description}`,
    );
    await new Promise((resolve) => setTimeout(resolve, retryAfter * 1000));
  }
});

bot.catch((err) => {
  error(
    `Handling update from ${formatUserChat(err.ctx)} failed: ${err.name} ${err.message}`,
  );
});

// if error happened, try to reply to the user with the error
bot.use(async (ctx, next) => {
  try {
    await next();
  } catch (err) {
    try {
      await ctx.reply(
        `Handling update failed: ${err}`,
        omitUndef({
          reply_to_message_id: ctx.message?.message_id,
          allow_sending_without_reply: true,
        }),
      );
    } catch {
      throw err;
    }
  }
});

// Wrap the calls in try-catch for error handling
async function setupBotCommands() {
  try {
    await bot.api.setMyShortDescription("Generate furry images in '704x704', '576x832', '832x576'. https://ko-fi.com/nyxthebot https://nyx.akiru.de/");
  } catch (err) {
    error(`Failed to set short description: ${err.message}`);
  }

  try {
    await bot.api.setMyDescription(
        "I can generate furry images from text. If you want a different size, use 'Size: 576x832' for example." +
        "Send /txt2img to generate an image.",
    );
  } catch (err) {
    error(`Failed to set description: ${err.message}`);
  }

  try {
    await bot.api.setMyCommands([
      { command: "txt2img", description: "Generate image from text" },
      { command: "img2img", description: "Generate image from image" },
      { command: "pnginfo", description: "Try to extract prompt from raw file" },
      { command: "queue", description: "Show the current queue" },
      { command: "cancel", description: "Cancel all your requests" },
    ]);
  } catch (err) {
    error(`Failed to set commands: ${err.message}`);
  }
}

// Call the setup function
setupBotCommands();

bot.command("start", async (ctx) => {
  if (ctx.match) {
    const id = ctx.match.trim();
    const session = sessions.get(id);
    if (session == null) {
      await ctx.reply(
        "Login failed: Invalid session ID",
        omitUndef({
          reply_to_message_id: ctx.message?.message_id,
        }),
      );
      return;
    }
    session.userId = ctx.from?.id;
    sessions.set(id, session);
    info(`User ${formatUserChat(ctx)} logged in`);
    // TODO: show link to web ui
    await ctx.reply(
      "Login successful! You can now return to the WebUI.",
      omitUndef({
        reply_to_message_id: ctx.message?.message_id,
      }),
    );
    return;
  }

  await ctx.reply(
    "Hello! Use the /txt2img command to generate an image",
    omitUndef({
      reply_to_message_id: ctx.message?.message_id,
    }),
  );
});

bot.command("txt2img", txt2imgCommand);
bot.use(txt2imgQuestion.middleware());

bot.command("img2img", img2imgCommand);
bot.use(img2imgQuestion.middleware());

bot.command("pnginfo", pnginfoCommand);
bot.use(pnginfoQuestion.middleware());

bot.command("queue", queueCommand);

bot.command("cancel", cancelCommand);

bot.command("broadcast", broadcastCommand);

bot.command("crash", () => {
  throw new Error("Crash command used");
});

// Set up the webhook in the telegram API and initialize the bot
await bot.api.setWebhook('https://nyx.akiru.de/webhook');
await bot.init();

// Function to handle incoming webhook requests
export async function handleWebhook(req: Request): Promise<Response> {
  try {
    const body = await req.json();
    // console.log("Received webhook data:", JSON.stringify(body, null, 2));

    // Log before processing update
    // console.log("Processing update through handleUpdate...");
    await bot.handleUpdate(body); // Process the update
    // Log after processing update
    // console.log("Update processed successfully.");

    return new Response(JSON.stringify({ status: 'ok' }), { headers: { 'Content-Type': 'application/json' } });
  } catch (error) {
    // Detailed error logging
    console.error("Error in handleWebhook:", error);
    return new Response("Error", { status: 500 });
  }
}