/**
 * pi-memory v1.1.0 — Three-Layer Memory System
 * 
 * Inspired by AEF's three-layer architecture (Grand Synthesis pattern #1):
 *   Layer 1: Working Memory — current session (handled by pi's compaction)
 *   Layer 2: Project Memory — MEMORY.md per project (learnings, patterns, decisions)
 *   Layer 3: Episodic Memory — per-task summaries that persist across sessions
 * 
 * /memory status              → show memory stats
 * /memory learn <insight>     → record a learning for current project
 * /memory recall [query]      → search project + episodic memory
 * /memory handover            → generate handover doc (focus, blockers, next actions)
 * /memory episode <summary>   → save episode (auto-called on session end)
 * 
 * Tools: memory_learn, memory_recall, memory_handover
 */
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import { existsSync, readFileSync, writeFileSync, appendFileSync, mkdirSync } from "node:fs";
import { join, basename } from "node:path";
import { homedir } from "node:os";

const MEMORY_DIR = join(homedir(), ".pi", "memory");
const EPISODES_DIR = join(MEMORY_DIR, "episodes");
const RST = "\x1b[0m", B = "\x1b[1m", D = "\x1b[2m";
const GREEN = "\x1b[32m", YELLOW = "\x1b[33m", CYAN = "\x1b[36m";

function ensureDirs() {
  if (!existsSync(MEMORY_DIR)) mkdirSync(MEMORY_DIR, { recursive: true });
  if (!existsSync(EPISODES_DIR)) mkdirSync(EPISODES_DIR, { recursive: true });
}

// ── Layer 2: Project Memory (MEMORY.md) ─────────────────────────
function getProjectMemoryPath(): string {
  const cwd = process.cwd();
  return join(cwd, "MEMORY.md");
}

function loadProjectMemory(): string {
  const p = getProjectMemoryPath();
  return existsSync(p) ? readFileSync(p, "utf-8") : "";
}

function appendProjectLearning(insight: string, category?: string) {
  ensureDirs();
  const p = getProjectMemoryPath();
  const ts = new Date().toISOString().slice(0, 10);
  const cat = category ? `[${category}]` : "";
  
  if (!existsSync(p)) {
    const projectName = basename(process.cwd());
    writeFileSync(p, `# ${projectName} — Project Memory\n\nLearnings, patterns, and decisions tracked across sessions.\n\n## Learnings\n\n`);
  }
  
  appendFileSync(p, `- ${ts} ${cat} ${insight}\n`);
  return { path: p, insight };
}

// ── Layer 3: Episodic Memory (per-task summaries) ───────────────
interface Episode {
  ts: string;
  project: string;
  summary: string;
  focus?: string;
  blockers?: string[];
  nextActions?: string[];
  filesChanged?: string[];
}

function getEpisodesFile(): string {
  const projectName = basename(process.cwd()).replace(/[^a-zA-Z0-9-_]/g, "_");
  return join(EPISODES_DIR, `${projectName}.jsonl`);
}

function saveEpisode(episode: Omit<Episode, "ts" | "project">) {
  ensureDirs();
  const record: Episode = {
    ts: new Date().toISOString(),
    project: basename(process.cwd()),
    ...episode
  };
  appendFileSync(getEpisodesFile(), JSON.stringify(record) + "\n");
  return record;
}

function loadEpisodes(limit = 10): Episode[] {
  const f = getEpisodesFile();
  if (!existsSync(f)) return [];
  const lines = readFileSync(f, "utf-8").trim().split("\n").filter(Boolean);
  return lines.slice(-limit).map(l => { try { return JSON.parse(l); } catch { return null; } }).filter(Boolean) as Episode[];
}

function searchMemory(query: string): { projectHits: string[]; episodeHits: Episode[] } {
  const q = query.toLowerCase();
  
  // Search project memory
  const projectMem = loadProjectMemory();
  const projectHits = projectMem.split("\n")
    .filter(line => line.toLowerCase().includes(q))
    .slice(0, 10);
  
  // Search episodes
  const f = getEpisodesFile();
  let episodeHits: Episode[] = [];
  if (existsSync(f)) {
    const lines = readFileSync(f, "utf-8").trim().split("\n").filter(Boolean);
    episodeHits = lines
      .map(l => { try { return JSON.parse(l) as Episode; } catch { return null; } })
      .filter((e): e is Episode => e !== null && JSON.stringify(e).toLowerCase().includes(q))
      .slice(-5);
  }
  
  return { projectHits, episodeHits };
}

// ── Handover Protocol (AEF pattern #10) ─────────────────────────
function generateHandover(): string {
  const episodes = loadEpisodes(3);
  const projectMem = loadProjectMemory();
  const lastEp = episodes[episodes.length - 1];
  
  let handover = `# Handover — ${basename(process.cwd())}\n`;
  handover += `Generated: ${new Date().toISOString().slice(0, 19)}\n\n`;
  
  if (lastEp) {
    handover += `## Last Session\n`;
    handover += `- **Focus:** ${lastEp.focus || lastEp.summary}\n`;
    if (lastEp.blockers?.length) handover += `- **Blockers:** ${lastEp.blockers.join(", ")}\n`;
    if (lastEp.nextActions?.length) {
      handover += `- **Next Actions:**\n`;
      for (const a of lastEp.nextActions) handover += `  - ${a}\n`;
    }
    if (lastEp.filesChanged?.length) handover += `- **Files Changed:** ${lastEp.filesChanged.join(", ")}\n`;
    handover += `\n`;
  }
  
  if (episodes.length > 1) {
    handover += `## Recent Episodes (${episodes.length})\n`;
    for (const ep of episodes) {
      handover += `- ${ep.ts.slice(0, 10)} — ${ep.summary}\n`;
    }
    handover += `\n`;
  }
  
  // Extract recent learnings from MEMORY.md
  const learnings = projectMem.split("\n").filter(l => l.startsWith("- 202")).slice(-5);
  if (learnings.length) {
    handover += `## Recent Learnings\n`;
    for (const l of learnings) handover += `${l}\n`;
  }
  
  return handover;
}

export default function init(pi: ExtensionAPI) {
  ensureDirs();

  // ── Commands ──────────────────────────────────────────────────
  pi.registerCommand("memory", {
    description: "Three-layer memory: /memory status|learn|recall|handover|episode",
    handler: async (args: string) => {
      const parts = (args || "").trim().split(/\s+/);
      const cmd = parts[0] || "status";
      const rest = parts.slice(1).join(" ");

      if (cmd === "status") {
        const projectMem = loadProjectMemory();
        const episodes = loadEpisodes(100);
        const learningCount = projectMem.split("\n").filter(l => l.startsWith("- 202")).length;
        
        let out = `${B}${CYAN}🧠 Three-Layer Memory${RST}\n\n`;
        out += `${B}Layer 1 — Working Memory${RST}\n`;
        out += `  ${D}Handled by pi's compaction system${RST}\n\n`;
        out += `${B}Layer 2 — Project Memory${RST} ${existsSync(getProjectMemoryPath()) ? GREEN + "✓" + RST : YELLOW + "(none)" + RST}\n`;
        out += `  File: ${D}${getProjectMemoryPath()}${RST}\n`;
        out += `  Learnings: ${learningCount}\n\n`;
        out += `${B}Layer 3 — Episodic Memory${RST}\n`;
        out += `  Episodes: ${episodes.length}\n`;
        out += `  File: ${D}${getEpisodesFile()}${RST}\n`;
        
        if (episodes.length) {
          const last = episodes[episodes.length - 1];
          out += `  Last: ${D}${last.ts.slice(0, 10)} — ${last.summary.slice(0, 80)}${RST}\n`;
        }
        
        return out;
      }

      if (cmd === "learn") {
        if (!rest) return `${YELLOW}Usage: /memory learn <insight>${RST}`;
        const result = appendProjectLearning(rest);
        return `${GREEN}✅ Recorded in MEMORY.md:${RST} ${rest}`;
      }

      if (cmd === "recall") {
        if (!rest) {
          const episodes = loadEpisodes(5);
          const projectMem = loadProjectMemory();
          const learnings = projectMem.split("\n").filter(l => l.startsWith("- 202")).slice(-5);
          
          let out = `${B}${CYAN}🔍 Recent Memory${RST}\n\n`;
          if (learnings.length) {
            out += `${B}Learnings:${RST}\n`;
            for (const l of learnings) out += `  ${l}\n`;
            out += `\n`;
          }
          if (episodes.length) {
            out += `${B}Episodes:${RST}\n`;
            for (const ep of episodes) out += `  ${D}${ep.ts.slice(0, 10)}${RST} ${ep.summary}\n`;
          }
          return out || `${YELLOW}No memory yet. Use /memory learn <insight> to start.${RST}`;
        }
        
        const results = searchMemory(rest);
        let out = `${B}${CYAN}🔍 Memory Search: "${rest}"${RST}\n\n`;
        if (results.projectHits.length) {
          out += `${B}Project Memory:${RST}\n`;
          for (const h of results.projectHits) out += `  ${h}\n`;
          out += `\n`;
        }
        if (results.episodeHits.length) {
          out += `${B}Episodes:${RST}\n`;
          for (const ep of results.episodeHits) out += `  ${D}${ep.ts.slice(0, 10)}${RST} ${ep.summary}\n`;
        }
        if (!results.projectHits.length && !results.episodeHits.length) {
          out += `${YELLOW}No matches found.${RST}`;
        }
        return out;
      }

      if (cmd === "handover") {
        const handover = generateHandover();
        const handoverPath = join(process.cwd(), "HANDOVER.md");
        writeFileSync(handoverPath, handover);
        return `${GREEN}✅ Handover generated:${RST} ${handoverPath}\n\n${handover}`;
      }

      if (cmd === "episode") {
        if (!rest) return `${YELLOW}Usage: /memory episode <summary>${RST}`;
        const ep = saveEpisode({ summary: rest });
        return `${GREEN}✅ Episode saved:${RST} ${ep.summary}`;
      }

      return `${YELLOW}Usage: /memory status|learn|recall|handover|episode${RST}`;
    },
  });

  // ── Tools ─────────────────────────────────────────────────────
  pi.registerTool({
    name: "memory_learn",
    label: "Memory Learn",
    description: "Record a learning/insight into the project's MEMORY.md. Persists across sessions.",
    parameters: Type.Object({
      insight: Type.String({ description: "The learning or insight to record" }),
      category: Type.Optional(Type.String({ description: "Category: pattern, decision, gotcha, architecture, process" })),
    }),
    async execute(toolCallId, params, signal, onUpdate, ctx) {
      const result = appendProjectLearning(params.insight, params.category);
      return {
        content: [{ type: "text", text: `Recorded in ${result.path}: ${result.insight}` }],
        details: { path: result.path },
      };
    },
  });

  pi.registerTool({
    name: "memory_recall",
    label: "Memory Recall",
    description: "Search project memory (MEMORY.md) and episodic memory for relevant context.",
    parameters: Type.Object({
      query: Type.Optional(Type.String({ description: "Search query. Omit to get recent memory." })),
      limit: Type.Optional(Type.Number({ description: "Max episodes to return (default 5)" })),
    }),
    async execute(toolCallId, params, signal, onUpdate, ctx) {
      if (!params.query) {
        const episodes = loadEpisodes(params.limit || 5);
        const projectMem = loadProjectMemory();
        return {
          content: [{ type: "text", text: JSON.stringify({ projectMemory: projectMem.slice(-2000), recentEpisodes: episodes }, null, 2) }],
          details: {},
        };
      }
      return {
        content: [{ type: "text", text: JSON.stringify(searchMemory(params.query), null, 2) }],
        details: {},
      };
    },
  });

  pi.registerTool({
    name: "memory_handover",
    label: "Memory Handover",
    description: "Generate a handover document for the current project.",
    parameters: Type.Object({
      summary: Type.Optional(Type.String({ description: "Summary of what was accomplished" })),
      focus: Type.Optional(Type.String({ description: "Current focus area" })),
      blockers: Type.Optional(Type.Array(Type.String(), { description: "Current blockers" })),
      nextActions: Type.Optional(Type.Array(Type.String(), { description: "Next actions" })),
    }),
    async execute(toolCallId, params, signal, onUpdate, ctx) {
      if (params.summary) {
        saveEpisode({ summary: params.summary, focus: params.focus, blockers: params.blockers, nextActions: params.nextActions });
      }
      const handover = generateHandover();
      const handoverPath = join(process.cwd(), "HANDOVER.md");
      writeFileSync(handoverPath, handover);
      return {
        content: [{ type: "text", text: `Handover saved to ${handoverPath}:\n\n${handover}` }],
        details: { path: handoverPath },
      };
    },
  });
}