/**
 * Memory Blocks Extension
 *
 * Gives the LLM persistent, self-editable memory blocks (inspired by MemGPT/Letta).
 * Blocks are stored as .pi/memory/<key>.md files with frontmatter metadata.
 *
 * - Blocks are injected into the system prompt every turn
 * - The LLM can update blocks via the `updateMemory` tool
 * - Only existing block keys can be written to (no creation by LLM)
 * - Two default blocks are provided: `user` and `agent`
 */

import * as fs from "node:fs";
import * as path from "node:path";
import matter from "gray-matter";
import yaml from "js-yaml";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { type Static, Type } from "@sinclair/typebox";
import { Value } from "@sinclair/typebox/value";
import { Text } from "@mariozechner/pi-tui";
import { isSlug } from "./lib/slug.ts";
import { stem } from "./lib/path.ts";

// --- Types ---

const BlockFrontmatter = Type.Object({
  description: Type.String({ default: "" }),
  limit: Type.Number({ default: 2000 }),
});

type BlockFrontmatter = Static<typeof BlockFrontmatter>;

export type Block = BlockFrontmatter & {
  content: string;
};

export type BlockEntry = {
  key: string;
  block: Block;
};

// --- Block I/O ---

export const parseBlock = (raw: string): Block => {
  const { data, content } = matter(raw);
  const frontmatter = Value.Decode(
    BlockFrontmatter,
    Value.Default(BlockFrontmatter, data),
  );
  return { ...frontmatter, content: content.trim() };
};

export const serializeBlock = (block: Block): string => {
  const frontmatter = yaml
    .dump(
      { description: block.description, limit: block.limit },
      { lineWidth: -1 },
    )
    .trim();
  return `---\n${frontmatter}\n---\n${block.content}\n`;
};

export const getMemoryDir = (cwd: string): string =>
  path.join(cwd, ".pi", "memory-blocks");

/**
 * List block keys in a memory directory, validating each key.
 * Throws if any .md file has an invalid key (not in slug form).
 */
export function* listBlockKeys(memoryDir: string): Generator<string> {
  for (const path of fs.readdirSync(memoryDir)) {
    if (!path.endsWith(".md")) {
      continue;
    }
    const pathStem = stem(path);
    if (!isSlug(pathStem)) {
      throw new TypeError(
        `Invalid memory block filename "${path}". Block filenames must be valid XML tag names (lowercase, hyphen, or underscore-separated slugs e.g. "user.md", "project-notes.md")`,
      );
    }
    yield pathStem;
  }
}

export const readBlock = (memoryDir: string, key: string): Block => {
  const filePath = path.join(memoryDir, `${key}.md`);
  const raw = fs.readFileSync(filePath, "utf-8");
  return parseBlock(raw);
};

export const writeBlock = (
  memoryDir: string,
  key: string,
  block: Block,
): void => {
  const filePath = path.join(memoryDir, `${key}.md`);
  fs.writeFileSync(filePath, serializeBlock(block), "utf-8");
};

export const readAllBlocks = (memoryDir: string): BlockEntry[] =>
  Array.from(listBlockKeys(memoryDir)).map((key) => ({
    key,
    block: readBlock(memoryDir, key),
  }));

export const DEFAULT_MEMORY_BLOCKS: Record<string, Block> = {
  user: {
    description:
      "Information about the user, their preferences, facts about them, and relevant context",
    limit: 2000,
    content: "",
  },
  agent: {
    description:
      "Your role, own self-concept, personality traits, and behavioral guidelines",
    limit: 2000,
    content: "",
  },
};

export const ensureDefaults = (memoryDir: string): void => {
  if (!fs.existsSync(memoryDir)) {
    fs.mkdirSync(memoryDir, { recursive: true });
  }

  for (const [key, block] of Object.entries(DEFAULT_MEMORY_BLOCKS)) {
    const filePath = path.join(memoryDir, `${key}.md`);
    if (!fs.existsSync(filePath)) {
      writeBlock(memoryDir, key, block);
    }
  }
};

// --- Format helpers ---

export const renderMemoryBlock = ({ key, block }: BlockEntry): string => {
  return `<${key}>
  <description>${block.description}</description>
  <metadata>
  - chars_current: ${block.content.length}
  - chars_limit: ${block.limit}
  </metadata>
  <value>${block.content}</value>
</${key}>`;
};

export const renderMemorySystemPrompt = (blocks: BlockEntry[]): string => {
  const blockMarkup = blocks.map((b) => renderMemoryBlock(b)).join("\n");
  return `You have persistent memory blocks that survive across sessions. Use the updateMemory tool to store important information you learn. Review your memory blocks below and keep them up to date.

<memory_blocks>
  ${blockMarkup}
</memory_blocks>
`;
};

// --- Extension ---

export default function memoryExtension(pi: ExtensionAPI) {
  let memoryDir = "";

  // Ensure default blocks exist on session start
  pi.on("session_start", async (_event, ctx) => {
    memoryDir = getMemoryDir(ctx.cwd);
    ensureDefaults(memoryDir);

    const blockDesc = readAllBlocks(memoryDir)
      .map(
        ({ key, block }) => `${key} (${block.content.length}/${block.limit})`,
      )
      .join(", ");
    ctx.ui.notify(`Memory blocks: ${blockDesc}`, "info");
  });

  // Inject memory blocks into system prompt every turn
  pi.on("before_agent_start", async (event) => {
    if (!memoryDir) return;

    const blocks = readAllBlocks(memoryDir);
    if (blocks.length === 0) return;

    return {
      systemPrompt: [event.systemPrompt, renderMemorySystemPrompt(blocks)].join(
        "\n\n",
      ),
    };
  });

  // Register the updateMemory tool
  pi.registerTool({
    name: "updateMemory",
    label: "Update Memory",
    description:
      "Update a persistent memory block. If oldText is provided, it is replaced with newText. " +
      "If oldText is omitted, newText is appended. Content is truncated to the block's character limit.",
    promptSnippet: "Update a persistent memory block (replace or append text)",
    promptGuidelines: [
      "Use updateMemory to persist important facts you learn about the user, project, or your own role.",
      "Prefer replacing outdated information over appending duplicates.",
      "Keep memory blocks concise and well-organized.",
    ],
    parameters: Type.Object({
      blockKey: Type.String({
        description:
          "Block to update (e.g. 'user', 'agent'). Must be an existing block.",
      }),
      oldText: Type.Optional(
        Type.String({
          description: "Text to find and replace. Omit to append instead.",
        }),
      ),
      newText: Type.String({
        description:
          "Replacement text, or text to append if oldText is omitted.",
      }),
    }),

    async execute(_toolCallId, params, _signal, _onUpdate, _ctx) {
      const { blockKey, oldText, newText } = params;

      // Validate block exists
      const keys = Array.from(listBlockKeys(memoryDir));
      if (!keys.includes(blockKey)) {
        throw new Error(
          `Block "${blockKey}" does not exist. Available blocks: ${keys.join(", ")}`,
        );
      }

      const block = readBlock(memoryDir, blockKey);
      let truncated = false;

      if (oldText !== undefined) {
        // Replace mode
        if (!block.content.includes(oldText)) {
          throw new Error(
            `oldText not found in block "${blockKey}". Use the memory shown in the system prompt as reference.`,
          );
        }
        block.content = block.content.replace(oldText, newText);
      } else {
        // Append mode
        block.content = block.content
          ? block.content + "\n" + newText
          : newText;
      }

      // Truncate to limit
      if (block.content.length > block.limit) {
        block.content = block.content.slice(0, block.limit);
        truncated = true;
      }

      writeBlock(memoryDir, blockKey, block);

      const usage = `${block.content.length}/${block.limit}`;
      const truncMsg = truncated ? " (truncated to limit)" : "";

      return {
        content: [
          {
            type: "text" as const,
            text: `Updated block "${blockKey}" [${usage}]${truncMsg}\n\nCurrent content:\n${block.content}`,
          },
        ],
        details: { blockKey, usage, truncated, content: block.content },
      };
    },

    renderCall(args, theme) {
      const op = args.oldText !== undefined ? "replace" : "append";
      let text = theme.fg("toolTitle", theme.bold("updateMemory "));
      text += theme.fg("accent", args.blockKey);
      text += " " + theme.fg("muted", `(${op})`);
      return new Text(text, 0, 0);
    },

    renderResult(result, { expanded }, theme) {
      const details = result.details as
        | {
            blockKey: string;
            usage: string;
            truncated: boolean;
            content: string;
          }
        | undefined;

      if (!details) {
        const text = result.content[0];
        return new Text(text?.type === "text" ? text.text : "", 0, 0);
      }

      let text =
        theme.fg("success", "✓ ") +
        theme.fg("accent", details.blockKey) +
        " " +
        theme.fg("muted", `[${details.usage}]`);

      if (details.truncated) {
        text += " " + theme.fg("warning", "(truncated)");
      }

      if (expanded && details.content) {
        text += "\n" + theme.fg("dim", details.content);
      }

      return new Text(text, 0, 0);
    },
  });

  // Register /memory command to view blocks
  pi.registerCommand("memory", {
    description: "Show all memory blocks and their usage",
    handler: async (_args, ctx) => {
      const blocks = readAllBlocks(memoryDir);

      if (blocks.length === 0) {
        ctx.ui.notify("No memory blocks found.", "info");
        return;
      }

      const lines = blocks.map(({ key, block }) => {
        const usage = `${block.content.length}/${block.limit}`;
        const bar = "█"
          .repeat(Math.round((block.content.length / block.limit) * 20))
          .padEnd(20, "░");
        const preview = block.content
          ? block.content.slice(0, 80).replace(/\n/g, " ") +
            (block.content.length > 80 ? "…" : "")
          : "(empty)";
        return `${key} [${usage}] ${bar}\n  ${block.description}\n  ${preview}`;
      });

      ctx.ui.notify(lines.join("\n\n"), "info");
    },
  });
}