import {
  copyToClipboard,
  type ExtensionAPI,
  type ExtensionContext,
} from "@earendil-works/pi-coding-agent";

export function extractLastCodeBlock(text: string): string | undefined {
  const blocks = text.matchAll(/^```[^\r\n]*\r?\n([\s\S]*?)^```[ \t]*$/gm);
  let last: string | undefined;

  for (const block of blocks) last = block[1]?.replace(/\r?\n$/, "");

  return last;
}

function getLastAssistantText(ctx: ExtensionContext): string | undefined {
  const branch = ctx.sessionManager.getBranch();

  for (let index = branch.length - 1; index >= 0; index--) {
    const entry = branch[index];
    if (entry?.type !== "message" || entry.message.role !== "assistant") continue;
    if (entry.message.stopReason === "aborted" && entry.message.content.length === 0) continue;

    return (
      entry.message.content
        .filter((part) => part.type === "text")
        .map((part) => part.text)
        .join("")
        .trim() || undefined
    );
  }
}

export default function copyblock(pi: ExtensionAPI) {
  pi.registerCommand("copyblock", {
    description: "Copy the last code block from the latest assistant response",
    handler: async (_args, ctx) => {
      const block = extractLastCodeBlock(getLastAssistantText(ctx) ?? "");
      if (block === undefined) {
        ctx.ui.notify("No code block found in the latest message. Nothing was copied.", "warning");
        return;
      }

      try {
        await copyToClipboard(block);
        ctx.ui.notify("Copied last code block to clipboard.", "info");
      } catch (error) {
        ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
      }
    },
  });
}