// raw thinking notepad (based on github:lyramakesmusic/neuralese-leaker)

import { Type } from "@earendil-works/pi-ai";
import { type ExtensionAPI, getMarkdownTheme } from "@earendil-works/pi-coding-agent";
import { Container, Markdown } from "@earendil-works/pi-tui";
import { loadSettings } from "./_char_common/settings";

const TOOL_NAME = "notepad";
const PARAMETER_IDLE_TIMEOUT_MS = 10_000;

const SCRATCH_REPLY = "Continue thinking, call a tool, or respond to the user.";

const SYSTEM_PROMPT_ADDENDUM = `# notepad

- Call notepad first, before anything else, on every new task.
- Re-enter the notepad after every tool result, before your next action.`;

export default function (pi: ExtensionAPI) {
  let enabled = true;
  let watchedCallId: string | undefined;
  let idleTimer: ReturnType<typeof setTimeout> | undefined;
  let resumeTimer: ReturnType<typeof setTimeout> | undefined;
  let resumeAfterAgentEnd = false;
  const truncatedCalls = new Set<string>();

  function clearIdleTimer() {
    if (idleTimer) clearTimeout(idleTimer);
    idleTimer = undefined;
    watchedCallId = undefined;
  }

  function isActive() {
    return enabled && pi.getActiveTools().includes(TOOL_NAME);
  }

  function setActive(on: boolean) {
    enabled = on;
    if (!on) clearIdleTimer();
    const active = pi.getActiveTools();
    if (on && !active.includes(TOOL_NAME)) {
      pi.setActiveTools([...active, TOOL_NAME]);
    } else if (!on && active.includes(TOOL_NAME)) {
      pi.setActiveTools(active.filter((t) => t !== TOOL_NAME));
    }
  }

  pi.registerTool({
    name: TOOL_NAME,
    label: "Notepad",
    description:
      "Notepad: Unlike typical scratchpads, it has no budget cap. " +
      "Use it before and between actions.",
    promptSnippet: "notepad; think here before and between actions",
    promptGuidelines: [
      "Use notepad to plan before acting, and re-enter it after every tool result.",
    ],
    parameters: Type.Object({
      text: Type.String({
        description: "Working notes",
      }),
    }),
    renderShell: "self",
    async execute(toolCallId) {
      const terminate = truncatedCalls.delete(toolCallId);
      return {
        content: [{ type: "text", text: SCRATCH_REPLY }],
        details: {},
        ...(terminate ? { terminate: true } : {}),
      };
    },
    renderCall(args, theme, context) {
      const text = typeof args?.text === "string" ? args.text : "";
      const md =
        (context.lastComponent as Markdown | undefined) ??
        new Markdown("", 1, 0, getMarkdownTheme(), {
          color: (t) => theme.fg("thinkingText", t),
          italic: true,
        });
      md.setText(text);
      return md;
    },
    renderResult: () => new Container(),
  });

  pi.on("session_start", async (_event, ctx) => {
    const tools = loadSettings(ctx.cwd).tools;
    const fromSettings =
      tools && typeof tools === "object"
        ? (tools as Record<string, unknown>).notepad
        : undefined;
    enabled = fromSettings !== false;
    if (enabled) setActive(true);
  });

  pi.on("message_update", (event, ctx) => {
    if (!isActive()) return;
    const update = event.assistantMessageEvent;
    if (
      update.type !== "toolcall_start" &&
      update.type !== "toolcall_delta" &&
      update.type !== "toolcall_end"
    ) {
      return;
    }

    const call = update.partial.content[update.contentIndex];
    if (call?.type !== "toolCall" || call.name !== TOOL_NAME) return;
    if (update.type === "toolcall_end") {
      clearIdleTimer();
      return;
    }

    clearIdleTimer();
    watchedCallId = call.id;
    idleTimer = setTimeout(() => {
      if (watchedCallId !== call.id) return;
      clearIdleTimer();
      truncatedCalls.add(call.id);
      resumeAfterAgentEnd = true;
      ctx.abort();
    }, PARAMETER_IDLE_TIMEOUT_MS);
  });

  // An aborted response normally skips tool execution. Reclassify only our
  // timed-out call so its partial arguments become an ordinary tool result.
  pi.on("message_end", (event) => {
    clearIdleTimer();
    if (event.message.role !== "assistant" || event.message.stopReason !== "aborted") return;

    const truncated = event.message.content.some(
      (part) => part.type === "toolCall" && truncatedCalls.has(part.id),
    );
    if (!truncated) return;

    const { errorMessage: _errorMessage, ...message } = event.message;
    return {
      message: {
        ...message,
        stopReason: "toolUse" as const,
        content: message.content.map((part) =>
          part.type === "toolCall" && truncatedCalls.has(part.id)
            ? {
                ...part,
                arguments: {
                  ...part.arguments,
                  work: typeof part.arguments.work === "string" ? part.arguments.work : "",
                },
              }
            : part,
        ),
      },
    };
  });

  // The salvaged result terminates the aborted run; resume once its signal is
  // gone so the next provider request gets a fresh one.
  pi.on("agent_end", (_event, ctx) => {
    if (!resumeAfterAgentEnd) return;
    resumeAfterAgentEnd = false;

    const resume = () => {
      if (!ctx.isIdle()) {
        resumeTimer = setTimeout(resume, 10);
        return;
      }
      resumeTimer = undefined;
      pi.sendMessage(
        {
          customType: "notepad-timeout",
          content:
            "The previous notepad call was truncated after its arguments stopped streaming. Continue, calling notepad again if more reasoning is needed.",
          display: false,
        },
        { triggerTurn: true },
      );
    };
    resumeTimer = setTimeout(resume, 0);
  });

  pi.on("session_shutdown", () => {
    clearIdleTimer();
    if (resumeTimer) clearTimeout(resumeTimer);
    resumeTimer = undefined;
    resumeAfterAgentEnd = false;
    truncatedCalls.clear();
  });

  pi.on("before_agent_start", async (event) => {
    if (!isActive()) return;
    pi.setThinkingLevel("off");
    if (event.systemPrompt.includes("# raw thinking (notepad)")) return;
    return { systemPrompt: event.systemPrompt + "\n\n" + SYSTEM_PROMPT_ADDENDUM };
  });

  pi.registerCommand("notepad", {
    description: "Toggle raw-thinking notepad: /notepad [on|off|status]",
    getArgumentCompletions: (prefix) => {
      const normalized = (prefix || "").trim().toLowerCase();
      const items = ["on", "off", "status"].filter((v) => v.startsWith(normalized));
      return items.length > 0 ? items.map((value) => ({ value, label: value })) : null;
    },
    handler: async (args, ctx) => {
      const arg = (args || "").trim().toLowerCase();
      if (arg === "" || arg === "status") {
        ctx.ui.notify(`Notepad raw thinking: ${isActive() ? "on" : "off"}`, "info");
        return;
      }
      if (arg !== "on" && arg !== "off") {
        ctx.ui.notify("Usage: /notepad [on|off|status]", "warning");
        return;
      }
      setActive(arg === "on");
      ctx.ui.notify(`Notepad raw thinking: ${arg} (this session)`, "info");
    },
  });
}