/**
 * pi-memory — Persistent memory extension for pi.
 *
 * Learns corrections, preferences, and patterns from sessions.
 * Injects relevant memory into future conversations.
 *
 * Lifecycle:
 * - session_start: open store, inject memory as a one-shot custom message
 * - (memory context is no longer injected per-turn — see v1.2.0 changelog)
 * - agent_end: queue messages for consolidation
 * - session_shutdown: consolidate and close store
 *
 * Tools:
 * - memory_search: search semantic memory
 * - memory_remember: manually add a memory
 * - memory_forget: delete a memory
 * - memory_lessons: list learned corrections
 * - memory_stats: show memory statistics
 */
import type { ExtensionAPI, AgentToolResult, SessionEntry } from "@mariozechner/pi-coding-agent";
import { Type, type TSchema } from "@sinclair/typebox";
import { join } from "node:path";
import { homedir } from "node:os";
import { readFileSync } from "node:fs";
import { MemoryStore } from "./store.js";
import { buildContextBlock, type InjectorConfig } from "./injector.js";

type ToolResult = AgentToolResult<unknown>;
function ok(text: string): ToolResult { return { content: [{ type: "text", text }], details: {} }; }

/**
 * Strip one layer of surrounding quotes from a string value.
 * Some local models (e.g. Qwen on certain runners) double-JSON-encode tool
 * arguments, emitting `"\"fact\""` instead of `"fact"`. We defensively
 * unwrap so these calls don't fail schema validation / equality checks.
 */
function stripQuotes<T>(v: T): T {
  if (typeof v !== "string") return v;
  const s = v.trim();
  if (s.length >= 2) {
    const first = s[0];
    const last = s[s.length - 1];
    if ((first === '"' && last === '"') || (first === "'" && last === "'")) {
      try {
        // Prefer JSON.parse for double-quoted (handles escapes)
        if (first === '"') return JSON.parse(s) as unknown as T;
      } catch { /* fall through */ }
      return s.slice(1, -1) as unknown as T;
    }
  }
  return v;
}
import {
  buildConsolidationPrompt,
  parseConsolidationResponse,
  applyExtracted,
  type ConsolidationInput,
} from "./consolidator.js";

const DEFAULT_MEMORY_DIR = join(homedir(), ".pi", "memory");
const DEFAULT_DB_PATH = join(DEFAULT_MEMORY_DIR, "memory.db");
const GLOBAL_SETTINGS_PATH = join(homedir(), ".pi", "agent", "settings.json");

/**
 * Default model used for session-end consolidation when no user override is
 * present in settings.json.  Preserves the historical behavior for existing
 * users — overridable via `memory.consolidationModel` (global or project).
 */
export const DEFAULT_CONSOLIDATION_MODEL = "claude-sonnet-4-20250514";

/**
 * Resolve the memory DB path for a given working directory.
 *
 * Priority (highest first):
 *   1. "pi-memory".localPath from {cwd}/.pi/settings.json → {localPath}/memory.db
 *   2. "pi-total-recall".localPath cascade → {localPath}/memory/memory.db
 *   3. Global default: ~/.pi/memory/memory.db  (preserves existing behavior)
 */
/**
 * Emit a warning when a settings block contains keys outside a known
 * schema. Catches silent typos like `LocalPath` vs `localPath` — an unknown
 * key is usually a misspelled known key that got silently ignored, leaving
 * the user wondering why their config didn't take effect.
 *
 * Logs to stderr (console.error) since this runs inside module-level code
 * at session_start; ctx.ui isn't reliably available yet.
 */
function warnUnknownKeys(block: unknown, blockName: string, knownKeys: readonly string[]): void {
  if (!block || typeof block !== "object") return;
  const unknown = Object.keys(block as Record<string, unknown>).filter((k) => !knownKeys.includes(k));
  if (unknown.length === 0) return;
  console.error(
    `pi-memory: ignoring unknown key(s) in settings.json "${blockName}" block: ${unknown.join(", ")} (expected: ${knownKeys.join(", ")})`,
  );
}

const PI_MEMORY_KNOWN_KEYS = ["localPath", "lessonInjection", "consolidationModel", "perTurnInjection"] as const;
const PI_TOTAL_RECALL_KNOWN_KEYS = ["localPath"] as const;

export function resolveDbPath(cwd: string): string {
  // Try reading the local project settings for an explicit localPath override
  try {
    const localSettingsPath = join(cwd, ".pi", "settings.json");
    const raw = readFileSync(localSettingsPath, "utf-8");
    const settings = JSON.parse(raw);

    // Package-specific override wins.
    const piMemory = settings?.["pi-memory"];
    warnUnknownKeys(piMemory, "pi-memory", PI_MEMORY_KNOWN_KEYS);
    if (piMemory && typeof piMemory === "object" && typeof piMemory.localPath === "string" && piMemory.localPath) {
      return join(piMemory.localPath, "memory.db");
    }

    // pi-total-recall cascade.
    const piTotalRecall = settings?.["pi-total-recall"];
    warnUnknownKeys(piTotalRecall, "pi-total-recall", PI_TOTAL_RECALL_KNOWN_KEYS);
    if (piTotalRecall && typeof piTotalRecall === "object" && typeof piTotalRecall.localPath === "string" && piTotalRecall.localPath) {
      return join(piTotalRecall.localPath, "memory", "memory.db");
    }
  } catch {
    // No local settings or parse error — use global default
  }
  // Default: global shared memory (preserves existing behavior)
  return DEFAULT_DB_PATH;
}

/**
 * Apply a single settings-block (the object under `memory` / `pi-memory`) to
 * `config`.  Invalid fields are ignored so a malformed value for one key
 * cannot clobber a valid value already set by a higher-priority source.
 */
function mergeMemorySettings(config: InjectorConfig, memorySettings: unknown): void {
  if (!memorySettings || typeof memorySettings !== "object") return;
  const m = memorySettings as Record<string, unknown>;

  if (m.lessonInjection === "all" || m.lessonInjection === "selective") {
    config.lessonInjection = m.lessonInjection;
  }
  if (typeof m.perTurnInjection === "boolean") {
    config.perTurnInjection = m.perTurnInjection;
  }
  if (typeof m.consolidationModel === "string" && m.consolidationModel.trim()) {
    config.consolidationModel = m.consolidationModel.trim();
  }
}

/**
 * Read pi-memory config from settings.json.
 * Looks for a "memory" (or project-local "pi-memory") key with
 * extension-specific settings.
 *
 * Example settings.json:
 * {
 *   "memory": {
 *     "perTurnInjection": true,
 *     "lessonInjection": "selective",
 *     "consolidationModel": "openai/gpt-4.1-mini"
 *   }
 * }
 *
 * Exported for tests.
 */
export function readSettingsConfig(cwd?: string): InjectorConfig {
  const config: InjectorConfig = {};

  // Read global settings
  try {
    const raw = readFileSync(GLOBAL_SETTINGS_PATH, "utf-8");
    const settings = JSON.parse(raw);
    mergeMemorySettings(config, settings?.memory);
  } catch {
    // no global settings
  }

  // Override with local project settings if available
  if (cwd) {
    try {
      const raw = readFileSync(join(cwd, ".pi", "settings.json"), "utf-8");
      const settings = JSON.parse(raw);
      // Accept either `memory` (preferred) or `pi-memory` (package-scoped).
      mergeMemorySettings(config, settings?.memory ?? settings?.["pi-memory"]);
    } catch {
      // no local settings
    }
  }

  return config;
}

export default function (pi: ExtensionAPI) {
  let store: MemoryStore | null = null;
  let pendingUserMessages: string[] = [];
  let pendingAssistantMessages: string[] = [];
  let sessionCwd: string = "";
  let sessionId: string | undefined;
  let cachedCtx: any = null;
  let resolvedDbPath: string = DEFAULT_DB_PATH;
  let injectorConfig: InjectorConfig = readSettingsConfig();

  // ─── Lifecycle ───────────────────────────────────────────────────

  pi.on("session_start", async (_event, ctx) => {
    try {
      sessionCwd = ctx.cwd;
      cachedCtx = ctx;
      sessionId = (ctx as any).sessionId ?? (ctx as any).session?.id;

      // Resolve per-agent DB path from local settings or cwd
      resolvedDbPath = resolveDbPath(sessionCwd);
      injectorConfig = readSettingsConfig(sessionCwd);

      store = new MemoryStore(resolvedDbPath);

      // Seed pending messages from existing session history so that
      // /memory-consolidate works even when resuming a session (the
      // historical messages never fire agent_end).  See #5.
      pendingUserMessages = [];
      pendingAssistantMessages = [];
      try {
        const branch = ctx.sessionManager.getBranch();
        for (const entry of branch) {
          if (entry.type !== "message") continue;
          const msg = (entry as any).message;
          if (!msg) continue;
          if (msg.role === "user") {
            const text = extractText(msg.content);
            if (text) pendingUserMessages.push(text);
          } else if (msg.role === "assistant") {
            const text = extractText(msg.content);
            if (text) pendingAssistantMessages.push(text);
          }
        }
      } catch {
        // Session may not have entries yet (brand-new session)
      }

      const stats = store.stats();
      if (stats.semantic + stats.lessons > 0) {
        ctx.ui.setStatus("pi-memory", `Memory: ${stats.semantic} facts, ${stats.lessons} lessons`);
        // The captured ctx may be stale by the time this fires (resume,
        // /new, /fork, /reload). Stale-ctx access throws synchronously;
        // swallow it — by then the new session has set its own status.
        setTimeout(() => {
          try { ctx.ui.setStatus("pi-memory", ""); } catch { /* ctx stale: harmless */ }
        }, 5000);
      }

      // Inject stored memory as a one-shot custom message BEFORE any user
      // message arrives. Matches pi-knowledge-search's pattern.
      //
      // Skipped when `perTurnInjection: true` — in that mode the
      // before_agent_start handler below takes over with per-turn semantic
      // matching via systemPrompt mutation.
      //
      // Historical note: v1.0.x mutated event.systemPrompt in before_agent_start.
      // That broke provider prefix caches on every turn boundary (any drift in
      // the system block re-writes the conversation suffix at cacheWrite rates).
      //
      // v1.1.x returned { message } from before_agent_start. That was worse: the
      // custom message landed AFTER the user's question in history, so the model
      // responded to the memory block instead of the user.
      //
      // v1.2.0 injects once at session_start using fallback mode (all facts +
      // lessons, 8KB cap). Correct ordering, stable cache, simpler model.
      //
      // v1.3.x adds `perTurnInjection: true` as an opt-in to restore v1.0.x
      // per-turn selective behavior (mutates systemPrompt, breaks cache on
      // every turn boundary — users opt in knowing the tradeoff).
      if (!injectorConfig.perTurnInjection) {
        try {
          const alreadyInjected = ctx.sessionManager
            .getEntries()
            .some(
              (e: SessionEntry) =>
                e.type === "custom_message" && e.customType === "pi-memory-context",
            );
          if (!alreadyInjected) {
            const { text, stats: injStats } = buildContextBlock(
              store,
              sessionCwd,
              undefined, // no prompt → fallback: dump all relevant memory
              injectorConfig,
            );
            if (text) {
              pi.sendMessage({
                customType: "pi-memory-context",
                content: text,
                display: false,
                details: injStats,
              });
            }
          }
        } catch {
          // Injection is nice-to-have; never break startup over it.
        }
      }
    } catch (err: any) {
      ctx.ui.notify(`pi-memory: failed to open store: ${err.message}`, "warning");
    }
  });

  // ----------------------------------------------------------------
  // Opt-in per-turn selective injection (v1.3.0).
  //
  // When `perTurnInjection: true` is set, run a semantic search against the
  // current user prompt and append matching memory to event.systemPrompt.
  // MUST use systemPrompt (not { message }) — returning { message } puts the
  // content AFTER the user message and causes the model to respond to the
  // injected memory instead of the user. See v1.1.x postmortem.
  //
  // This breaks provider prefix caches on every turn boundary — an accepted
  // cost for users who want per-query relevance from large memory stores.
  // ----------------------------------------------------------------
  pi.on("before_agent_start", async (event, ctx) => {
    if (!store) return;
    if (!injectorConfig.perTurnInjection) return;

    const { text } = buildContextBlock(store, ctx.cwd, event.prompt, injectorConfig);
    if (!text) return;

    return {
      systemPrompt: `${event.systemPrompt}\n\n${text}`,
    };
  });


  pi.on("agent_end", async (event, _ctx) => {
    // Collect messages for consolidation at shutdown
    for (const msg of event.messages) {
      if (msg.role === "user" && "content" in msg) {
        const text = extractText(msg.content);
        if (text) {
          pendingUserMessages.push(text);
          if (pendingUserMessages.length > 60) pendingUserMessages.shift();
        }
      } else if (msg.role === "assistant" && "content" in msg) {
        const text = extractText(msg.content);
        if (text) {
          pendingAssistantMessages.push(text);
          if (pendingAssistantMessages.length > 60) pendingAssistantMessages.shift();
        }
      }
    }
  });

  // Consolidate memory when switching sessions (/new, /resume)
  pi.on("session_before_switch", async (_event, ctx) => {
    if (!store) return;

    if (pendingUserMessages.length >= 3) {
      ctx.ui.setStatus("pi-memory", "🧠 Consolidating memory...");
      try {
        await consolidateSession();
      } catch {
        // Best-effort
      } finally {
        // Always clear — even if consolidateSession threw synchronously or
        // ctx went stale. Stuck indicator otherwise pins for the rest of
        // the session. See samfoy/pi-total-recall#5.
        try { ctx.ui.setStatus("pi-memory", ""); } catch { /* ctx stale: harmless */ }
      }
    }

    // Reset for the next session
    pendingUserMessages = [];
    pendingAssistantMessages = [];
  });

  pi.on("session_shutdown", async () => {
    if (!store) return;

    // Immediate visual feedback — user sees this as soon as C-c C-c fires
    if (cachedCtx) {
      cachedCtx.ui.setStatus("pi-memory", "🧠 Consolidating memory...");
    }

    // Consolidate if we have enough conversation
    if (pendingUserMessages.length >= 3) {
      try {
        await consolidateSession();
      } catch {
        // Best-effort — don't crash on shutdown
      }
    }

    store.close();
    store = null;
  });

  // ─── Consolidation ──────────────────────────────────────────────

  async function consolidateSession(): Promise<void> {
    if (!store) return;

    const input: ConsolidationInput = {
      userMessages: pendingUserMessages,
      assistantMessages: pendingAssistantMessages,
      cwd: sessionCwd,
      sessionId,
    };

    const currentFacts = store.listSemantic(undefined, 200).map(f => ({ key: f.key, value: f.value }));
    const currentLessons = store.listLessons(undefined, 100).map(l => ({ rule: l.rule, category: l.category }));
    const prompt = buildConsolidationPrompt(input, currentFacts, currentLessons);

    // Use pi's exec to call the LLM via a lightweight pi session.
    // Use a fast model to avoid blocking shutdown for too long.
    //
    // Defence in depth: pi.exec has a 45s timeout, but we also wrap the
    // whole call in a hard 60s backstop. If pi.exec's timeout ever fails
    // to kill the child (e.g. stuck in syscall), the Promise.race below
    // still rejects and lets the caller clear its status indicator.
    const EXEC_TIMEOUT_MS = 45_000;
    const HARD_TIMEOUT_MS = 60_000;
    let backstopHandle: ReturnType<typeof setTimeout> | undefined;
    try {
      const execPromise = pi.exec("pi", [
        "-p", prompt,
        "--print",
        "--no-extensions",
        "--model", injectorConfig.consolidationModel ?? DEFAULT_CONSOLIDATION_MODEL,
      ], {
        timeout: EXEC_TIMEOUT_MS,
        cwd: sessionCwd,
      });

      const result = await Promise.race([
        execPromise,
        new Promise<never>((_, reject) => {
          backstopHandle = setTimeout(
            () => reject(new Error("consolidation backstop timeout")),
            HARD_TIMEOUT_MS,
          );
        }),
      ]);

      if (result.code === 0 && result.stdout) {
        const extracted = parseConsolidationResponse(result.stdout);
        const applied = applyExtracted(store!, extracted, `session:${sessionId ?? "unknown"}`);
        if (applied.semantic + applied.lessons > 0) {
          // Log but don't notify — we're shutting down
          console.error(`pi-memory: consolidated ${applied.semantic} facts, ${applied.lessons} lessons`);
        }
      }
    } catch {
      // Timeout or exec failure — skip consolidation this session
    } finally {
      if (backstopHandle) clearTimeout(backstopHandle);
    }
  }

  // ─── Tools ──────────────────────────────────────────────────────

  pi.registerTool({
    name: "memory_search",
    label: "Memory Search",
    description: "Search persistent memory for facts, preferences, and project patterns the user has established across sessions.",
    parameters: Type.Object({
      query: Type.String({ description: "Search query" }),
      limit: Type.Optional(Type.Number({ description: "Max results (default 10)" })),
    }) as any,
    async execute(_id, params, _signal, _update, _ctx) {
      if (!store) return ok("Memory store not initialized");

      const results = store.searchSemantic(params.query, params.limit ?? 10);
      if (results.length === 0) {
        return ok("No matching memories found.");
      }

      const text = results.map(r =>
        `${r.key}: ${r.value} (confidence: ${r.confidence}, source: ${r.source})`
      ).join("\n");

      return ok(text);
    },
  });

  pi.registerTool({
    name: "memory_remember",
    label: "Memory Remember",
    description: "Store a fact, preference, or lesson in persistent memory. Use dotted keys like pref.editor, project.rosie.lang, tool.sed.usage. For corrections, use type='lesson'.",
    parameters: Type.Object({
      type: Type.String({ description: "'fact' for key-value, 'lesson' for a correction" }),
      key: Type.Optional(Type.String({ description: "Dotted key for facts (e.g. pref.commit_style)" })),
      value: Type.Optional(Type.String({ description: "Value for facts" })),
      rule: Type.Optional(Type.String({ description: "Rule text for lessons" })),
      category: Type.Optional(Type.String({ description: "Category for lessons (default: general)" })),
      negative: Type.Optional(Type.Boolean({ description: "True if this is something to AVOID" })),
    }) as any,
    async execute(_id, params, _signal, _update, _ctx) {
      if (!store) return ok("Memory store not initialized");

      // Defensively unwrap double-quoted string args from misbehaving model runners.
      params = {
        ...params,
        type: stripQuotes(params.type),
        key: stripQuotes(params.key),
        value: stripQuotes(params.value),
        rule: stripQuotes(params.rule),
        category: stripQuotes(params.category),
      };

      if (params.type !== "fact" && params.type !== "lesson") {
        return ok(`Invalid type: ${params.type}. Must be 'fact' or 'lesson'.`);
      }

      if (params.type === "fact") {
        if (!params.key || !params.value) {
          return ok("Both key and value required for facts");
        }
        store.setSemantic(params.key, params.value, 0.95, "user");
        return ok(`Remembered: ${params.key} = ${params.value}`);
      }

      if (params.type === "lesson") {
        if (!params.rule) {
          return ok("Rule text required for lessons");
        }
        const result = store.addLesson(params.rule, params.category ?? "general", "user", params.negative ?? false);
        if (result.success) {
          return ok(`Lesson learned: ${params.rule}`);
        }
        return ok(`Already known (${result.reason}): ${params.rule}`);
      }

      return ok("Unknown type");
    },
  });

  pi.registerTool({
    name: "memory_forget",
    label: "Memory Forget",
    description: "Remove a fact or lesson from persistent memory.",
    parameters: Type.Object({
      type: Type.String(),
      key: Type.Optional(Type.String({ description: "Key for facts" })),
      id: Type.Optional(Type.String({ description: "ID for lessons" })),
    }) as any,
    async execute(_id, params, _signal, _update, _ctx) {
      if (!store) return ok("Memory store not initialized");

      params = {
        ...params,
        type: stripQuotes(params.type),
        key: stripQuotes(params.key),
        id: stripQuotes(params.id),
      };

      if (params.type !== "fact" && params.type !== "lesson") {
        return ok(`Invalid type: ${params.type}. Must be 'fact' or 'lesson'.`);
      }

      if (params.type === "fact" && params.key) {
        const deleted = store.deleteSemantic(params.key);
        return ok(deleted ? `Forgot: ${params.key}` : `Not found: ${params.key}`);
      }

      if (params.type === "lesson" && params.id) {
        const deleted = store.deleteLesson(params.id);
        return ok(deleted ? `Forgot lesson ${params.id}` : `Not found: ${params.id}`);
      }

      return ok("Provide key (for facts) or id (for lessons)");
    },
  });

  pi.registerTool({
    name: "memory_lessons",
    label: "Memory Lessons",
    description: "List learned corrections and lessons from past sessions.",
    parameters: Type.Object({
      category: Type.Optional(Type.String({ description: "Filter by category" })),
      limit: Type.Optional(Type.Number({ description: "Max results (default 50)" })),
    }) as any,
    async execute(_id, params, _signal, _update, _ctx) {
      if (!store) return ok("Memory store not initialized");

      const lessons = store.listLessons(params.category, params.limit ?? 50);
      if (lessons.length === 0) {
        return ok("No lessons learned yet.");
      }

      const text = lessons.map(l =>
        `${l.negative ? "❌" : "✅"} [${l.category}] ${l.rule} (id: ${l.id.slice(0, 8)})`
      ).join("\n");

      return ok(text);
    },
  });

  pi.registerTool({
    name: "memory_stats",
    label: "Memory Stats",
    description: "Show memory statistics — how many facts, lessons, and events are stored.",
    parameters: Type.Object({}) as any,
    async execute(_id, _params, _signal, _update, _ctx) {
      if (!store) return ok("Memory store not initialized");

      const stats = store.stats();
      const text = `Memory: ${stats.semantic} semantic facts, ${stats.lessons} active lessons, ${stats.events} events logged\nDB: ${resolvedDbPath}`;
      return ok(text);
    },
  });

  // ─── Commands ──────────────────────────────────────────────────

  pi.registerCommand("memory-consolidate", {
    description: "Manually trigger memory consolidation for the current session",
    async handler(_args, ctx) {
      if (!store) {
        ctx.ui.notify("Memory store not initialized", "warning");
        return;
      }

      if (pendingUserMessages.length < 2) {
        ctx.ui.notify("Not enough conversation to consolidate (need at least 2 user messages)", "warning");
        return;
      }

      ctx.ui.notify("Consolidating session memory...", "info");
      try {
        await consolidateSession();
        const stats = store.stats();
        ctx.ui.notify(`Memory updated: ${stats.semantic} facts, ${stats.lessons} lessons`, "info");
      } catch (err: any) {
        ctx.ui.notify(`Consolidation failed: ${err.message}`, "error");
      }
    },
  });
}

// ─── Helpers ─────────────────────────────────────────────────────────

function extractText(content: unknown): string {
  if (typeof content === "string") return content;
  if (Array.isArray(content)) {
    return content
      .filter((c: any) => c.type === "text" && typeof c.text === "string")
      .map((c: any) => c.text)
      .join("\n");
  }
  return "";
}