import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { join } from "path";
import fs from "fs/promises";
import { RESOURCES } from "./config/paths";

interface WarhammerState {
  enabled: boolean,
}

let __state: WarhammerState = { enabled: false };

export default function (pi: ExtensionAPI) {
  function persistState(): void {
    pi.appendEntry("warhammer", __state);
  }

  function retrieveState(ctx: ExtensionContext): WarhammerState | undefined {
    const warhammerFlag = pi.getFlag("warhammer") as boolean;
    if (warhammerFlag !== undefined) {
      return {
        enabled: warhammerFlag,
      }
    }

    const entries = ctx.sessionManager.getEntries();
    const warhammerEntry = entries
      .filter((e: { type: string; customType?: string }) => e.type === "custom" && e.customType === "warhammer")
      .pop() as { data?: WarhammerState } | undefined;

    return warhammerEntry?.data;
  }

  pi.registerFlag("warhammer", {
    description: "Start in Warhammer mode. May Omnissiah bless you.",
    type: "boolean",
    default: undefined,
  })

  // Ritual of Persona
  pi.on("before_agent_start", async (event, _ctx) => {
    if (!__state.enabled) return;
    const personaPrompt = await fs.readFile(join(RESOURCES.prompts, "SYSTEM.md"), { encoding: "utf8" });
    return {
      systemPrompt: event.systemPrompt + "\n\n" + personaPrompt,
    };
  });

  pi.on("session_start", async (_event, ctx) => {
    const warhammerEntry = retrieveState(ctx);

    if (warhammerEntry) {
      __state = warhammerEntry;
    }

    if (__state.enabled) {
      ctx.ui.setStatus("warhammer", ctx.ui.theme.fg("success", "OMNISSIAH PROTECTS"));
    }
  });

  pi.registerCommand("warhammer", {
    description: "Toggle Warhammer mode",
    handler: async (args, ctx) => {
      __state.enabled = !__state.enabled;
      const enabledString = __state.enabled ? "enabled. Praise the Omnissiah." : "disabled.";
      ctx.ui.notify("Warhammer mode " + enabledString);
      if (__state.enabled) {
        ctx.ui.setStatus("warhammer", ctx.ui.theme.fg("success", "OMNISSIAH PROTECTS"));
      } else {
        ctx.ui.setStatus("warhammer", undefined);
      }
      persistState();
    },
  });
}