hydro / src / registry / commandRegistry.ts

import type { Message } from "@fluxerjs/core";
import { COMMAND_META_KEY, type CommandMeta } from "@/decorators/command";
import { GUARDS_KEY, type Guard } from "@/decorators/guards";
import { parseArgs, type ArgSchema } from "@/args";
import type { BotContext, ICommand } from "@/types";
import type { CommandCtor } from "@/types";
import { Logger } from "tslog";
import { MODULE_META_KEY, MODULE_COMMANDS_KEY, type ModuleCommand } from "@/base/ModuleBase";

interface FileRecord {
  filePath: string;
  triggers: string[];
}

interface ModuleRecord {
  name: string;
  triggers: string[];
  instance: object;
  onUnload?: () => void | Promise<void>;
}

type AnyCommand = ICommand<any>;

interface ModuleEntry {
  type: "module";
  instance: object;
  methodName: string;
  args?: ArgSchema;
}

interface ClassEntry {
  type: "class";
  instance: AnyCommand;
  args?: ArgSchema;
}

type Entry = ModuleEntry | ClassEntry;

export class CommandRegistry {
  private map = new Map<string, Entry>();
  private guards = new Map<string, Guard[]>();
  private fileSources = new Map<string, FileRecord>();
  private modules = new Map<string, ModuleRecord>();
  private logger: Logger<CommandRegistry>;

  constructor() {
    this.logger = new Logger({
      type: "pretty",
      name: "CommandRegistry",
    });
  }

  register(prefix: string, filePath: string | null, ...ctors: CommandCtor[]): void {
    for (const ctor of ctors) {
      try {
        const meta: CommandMeta = Reflect.getMetadata(COMMAND_META_KEY, ctor);
        if (!meta) throw new Error(`${ctor.name} has missing @Command decorators`);

        const instance = new ctor();
        const instanceGuards: Guard[] = Reflect.getMetadata(GUARDS_KEY, ctor) ?? [];
        const triggers = [meta.name, ...(meta.aliases ?? [])].map((n) => `${prefix}${n}`);

        for (const trigger of triggers) {
          this.map.set(trigger, { type: "class", instance, args: instance.args });
          this.guards.set(trigger, instanceGuards);
        }

        if (filePath) {
          this.fileSources.set(meta.name, { filePath, triggers });
        }
      } catch (e) {
        if (e instanceof Error)
          this.logger.error(`Something went wrong while registring: ${e.message}`);
      }
    }
  }

  registerModule(
    prefix: string,
    instance: object,
    options?: { namespace?: string; guards?: Guard[] },
  ): void {
    const meta = Reflect.getMetadata(MODULE_META_KEY, instance.constructor);
    if (!meta) {
      this.logger.error(`${instance.constructor.name} has missing @Module decorator`);
      return;
    }

    const commands: ModuleCommand[] =
      Reflect.getMetadata(MODULE_COMMANDS_KEY, instance.constructor) ?? [];
    if (commands.length === 0) {
      this.logger.warn(`Module ${meta.name} has no commands`);
      return;
    }

    const ns = options?.namespace ?? meta.name;
    const moduleGuards: Guard[] = options?.guards ?? [];
    const triggers: string[] = [];

    for (const cmd of commands) {
      const commandTriggers = [cmd.meta.name, ...(cmd.meta.aliases ?? [])].map(
        (n) => `${prefix}${ns}:${n}`,
      );
      triggers.push(...commandTriggers);

      for (const trigger of commandTriggers) {
        this.map.set(trigger, {
          type: "module",
          instance,
          methodName: cmd.methodName,
          args: cmd.meta.args,
        });
        this.guards.set(trigger, [...moduleGuards]);
      }
    }

    this.modules.set(meta.name, {
      name: meta.name,
      triggers,
      instance,
      onUnload: (instance as { onUnload?: () => void | Promise<void> }).onUnload?.bind(instance),
    });

    (instance as { onLoad?: (registry: CommandRegistry) => void | Promise<void> }).onLoad?.(this);

    this.logger.info(`Registered module "${meta.name}" with ${commands.length} commands`);
  }

  unregisterModule(name: string): void {
    const record = this.modules.get(name);
    if (!record) return;

    for (const trigger of record.triggers) {
      this.map.delete(trigger);
      this.guards.delete(trigger);
    }

    record.onUnload?.();
    this.modules.delete(name);
    this.logger.info(`Unregistered module "${name}"`);
  }

  unregisterFile(filePath: string): void {
    for (const [, record] of this.fileSources) {
      if (record.filePath === filePath) {
        for (const trigger of record.triggers) {
          this.map.delete(trigger);
          this.guards.delete(trigger);
        }
        this.fileSources.delete(record.triggers[0]!);
      }
    }
  }

  async handle(ctx: BotContext, prefix: string, message: Message): Promise<void> {
    if (!message.content.startsWith(prefix)) return;

    try {
      const [rawCmd, ...args] = message.content.slice(prefix.length).trim().split(/\s+/);
      const key = `${prefix}${rawCmd?.toLowerCase()}`;

      const entry = this.map.get(key);
      if (!entry) return;

      const guards = this.guards.get(key) ?? [];
      for (const guard of guards) {
        if (!(await guard(message))) return;
      }

      const parsedArgs = entry.args ? await parseArgs(entry.args, args) : ({} as never);

      this.logger.info(
        `Executing ${rawCmd} ${args.length ? `with arguments: ${args.join(", ")}` : "without arguments"}`,
      );

      if (entry.type === "class") {
        await entry.instance.execute(ctx, message, parsedArgs);
      } else {
        const method = (entry.instance as Record<string, unknown>)[entry.methodName];
        if (typeof method === "function") {
          await (
            method as (ctx: BotContext, message: Message, args: unknown) => Promise<void>
          ).call(entry.instance, ctx, message, parsedArgs);
        }
      }
    } catch (err) {
      if (err instanceof Error) {
        this.logger.error(`Something went wrong while executing this command: ${err.message}`);
        message.reply({
          content: `Something went wrong while executing this command:\n\`\`\`\n${err.message}\n\`\`\``,
        });
      }
    }
  }
}