import { StringEnum } from "@earendil-works/pi-ai";
import type {
  ExtensionAPI,
  ExtensionCommandContext,
  ExtensionContext,
  SessionEntry,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { EpochController } from "./controller.ts";
import { loadConfig } from "./config.ts";
import { DEFAULT_CONFIG } from "./constants.ts";
import {
  collectHistoryItems,
  formatHistoryRead,
  readHistory,
  searchHistory,
} from "./history.ts";
import { listNotes, replayNotes, validateNoteWrite } from "./notes.ts";
import { getEpochState } from "./epoch.ts";
import { evaluatePressure } from "./pressure.ts";

const NEW_CONTEXT_DESCRIPTION =
  "Start a new context epoch: end the current working context and continue the SAME Pi session in a fresh one. " +
  "The old epoch's messages leave the active context but stay searchable via context_history. " +
  "No summary is generated. Before calling this, persist any non-reconstructable state (user decisions, constraints, blockers, external side effects) with context_notes. " +
  "Do not save things that are cheap to re-derive from source files, git, tests, or session history.";

const NOTES_DESCRIPTION =
  "Session-scoped, branch-aware notes that survive context epoch rollovers. " +
  "Store ONLY semantic state that cannot be cheaply reconstructed: explicit user decisions and rejected approaches (with reasons), non-obvious user constraints, unresolved blockers, acceptance criteria, external side effects, facts not represented in workspace state. " +
  "Do NOT store things re-derivable from source files, git, tests, package metadata, or session history (e.g. what a source file contains - just read it again later). " +
  "Notes are not long-term memory: they end when the session ends.";

const HISTORY_DESCRIPTION =
  "Search and read the durable session transcript, including messages retired by previous context epochs. " +
  "search finds old messages by keyword (returns stable entry refs); read retrieves one entry plus surrounding context by ref. " +
  "Use this to re-ground in prior work without keeping old epochs in the active context. Reasoning/thinking content is never returned.";

const TOOLS_NOTE_GUIDELINE =
  "context_notes is for non-reconstructable semantic state only (user decisions, constraints, blockers, acceptance criteria). Do not copy source code or file contents into notes; re-read files instead.";

/**
 * pi-context-epochs extension entry point.
 *
 * Model-facing surface: exactly three tools (new_context, context_notes,
 * context_history) plus one /epoch command. Everything else is runtime
 * machinery: the durable epoch boundary rides on pi's own compaction path.
 */
export default function createEpochsExtension(pi: ExtensionAPI): void {
  const { config, warnings } = loadConfig();
  const controller = new EpochController(pi, config);

  // -------------------------------------------------------------------------
  // Lifecycle
  // -------------------------------------------------------------------------
  pi.on("session_start", (_event, ctx) => {
    controller.onSessionStart(ctx);
    if (warnings.length > 0) {
      for (const warning of warnings) {
        if (ctx.hasUI) ctx.ui.notify(warning, "warning");
      }
    }
  });

  pi.on("session_before_compact", (event, ctx) =>
    controller.onBeforeCompact(
      {
        reason: event.reason,
        customInstructions: event.customInstructions,
        preparation: event.preparation,
        branchEntries: event.branchEntries,
      },
      ctx,
    ),
  );

  pi.on("session_compact", (event, ctx) => controller.onSessionCompact(event, ctx));
  pi.on("session_compact_failed", (_event, ctx) => controller.onSessionCompactFailed(ctx));
  pi.on("agent_end", (event) => controller.onAgentEnd(event as unknown as { messages: Array<{ stopReason?: string }> }));
  pi.on("agent_settled", (_event, ctx) => controller.onAgentSettled(ctx));
  pi.on("turn_end", (_event, ctx) => controller.onTurnEnd(ctx));
  pi.on("session_tree", (_event, ctx) => controller.onSessionTree(ctx));
  pi.on("model_select", (_event, ctx) => controller.onModelSelect(ctx));
  pi.on("session_shutdown", () => controller.onSessionShutdown());

  // -------------------------------------------------------------------------
  // Tools
  // -------------------------------------------------------------------------
  pi.registerTool({
    name: "new_context",
    label: "New context epoch",
    description: NEW_CONTEXT_DESCRIPTION,
    promptSnippet: "Start a fresh working context epoch in the same session (no summary; prior history stays retrievable)",
    promptGuidelines: [
      "Call new_context when the current working context is exhausted or no longer useful. First save non-reconstructable state with context_notes; after the rollover, re-ground from the workspace, notes, and context_history.",
    ],
    parameters: Type.Object({}),
    async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
      const arm = controller.armModelRollover(ctx);
      const state = getEpochState(ctx.sessionManager.getBranch());
      if (!arm.armed) {
        return {
          content: [{ type: "text", text: `No rollover performed. ${arm.reason ?? ""}`.trim() }],
          details: { armed: false, epochIndex: state.epochIndex },
        };
      }
      return {
        content: [
          {
            type: "text",
            text:
              `Context epoch boundary armed (next epoch: ${arm.epochIndex}). ` +
              "It takes effect when your current response finishes. The next turn starts in a fresh working context " +
              "containing only an epoch marker; old messages remain available via context_history. " +
              "Finish your current response now; do not call new_context again.",
          },
        ],
        details: { armed: true, nextEpochIndex: arm.epochIndex },
      };
    },
  });

  pi.registerTool({
    name: "context_notes",
    label: "Context notes",
    description: NOTES_DESCRIPTION,
    promptSnippet: "Session-scoped notes for non-reconstructable state (decisions, constraints, blockers); survives epoch rollovers",
    promptGuidelines: [TOOLS_NOTE_GUIDELINE],
    parameters: Type.Object({
      action: StringEnum(["list", "read", "write", "delete"] as const, {
        description: "list: show note keys; read: fetch one note; write: create/update; delete: remove",
      }),
      key: Type.Optional(Type.String({ description: "Note key (required for read/write/delete)" })),
      content: Type.Optional(Type.String({ description: "Note content (required for write)" })),
    }),
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const branch = ctx.sessionManager.getBranch();
      const notes = replayNotes(branch);

      if (params.action === "list") {
        const items = listNotes(notes);
        if (items.length === 0) {
          return { content: [{ type: "text", text: "No context notes in this session branch." }], details: { keys: [] } };
        }
        const text = ["Available context notes:", ...items.map((n) => `- ${n.key} (${n.content.length} chars)`)].join("\n");
        return { content: [{ type: "text", text }], details: { keys: items.map((n) => n.key) } };
      }

      if (params.action === "read") {
        if (!params.key) return err("context_notes read requires 'key'.");
        const note = notes.get(params.key.trim());
        if (!note) {
          const available = listNotes(notes).map((n) => n.key);
          return err(`No note named "${params.key}". Available: ${available.length ? available.join(", ") : "(none)"}`);
        }
        return {
          content: [{ type: "text", text: `[context note: ${note.key}]\n${note.content}` }],
          details: { key: note.key },
        };
      }

      if (params.action === "write") {
        const validation = validateNoteWrite(params.key, params.content, controller.config.noteMaxChars);
        if (!validation.ok) return err(validation.error ?? "invalid note");
        const key = validation.key as string;
        const exists = notes.has(key);
        pi.appendEntry("pi-context-epochs/note", {
          action: "write",
          key,
          content: validation.content,
          timestamp: new Date().toISOString(),
        });
        return {
          content: [
            {
              type: "text",
              text: `Context note "${key}" ${exists ? "updated" : "created"} (${(validation.content as string).length} chars). It persists on the session branch and survives epoch rollovers.`,
            },
          ],
          details: { key, action: "write" },
        };
      }

      // delete
      if (!params.key) return err("context_notes delete requires 'key'.");
      const key = params.key.trim();
      if (!notes.has(key)) return err(`No note named "${key}" to delete.`);
      pi.appendEntry("pi-context-epochs/note", { action: "delete", key, timestamp: new Date().toISOString() });
      return { content: [{ type: "text", text: `Context note "${key}" deleted.` }], details: { key, action: "delete" } };
    },
  });

  pi.registerTool({
    name: "context_history",
    label: "Context history",
    description: HISTORY_DESCRIPTION,
    promptSnippet: "Search/read the durable session transcript including retired epochs (lexical search, stable entry refs)",
    parameters: Type.Object({
      action: StringEnum(["search", "read"] as const, {
        description: "search: find old messages by keyword; read: retrieve one entry (plus neighbors) by ref",
      }),
      query: Type.Optional(Type.String({ description: "Search keywords (required for search)" })),
      ref: Type.Optional(Type.String({ description: "Entry ref from search results (required for read)" })),
      limit: Type.Optional(Type.Number({ description: `Max search results (default ${DEFAULT_CONFIG.historySearchLimit})` })),
      around: Type.Optional(Type.Number({ description: `Neighbor entries per side for read (default ${DEFAULT_CONFIG.historyReadAround})` })),
      scope: Type.Optional(
        StringEnum(["branch", "session"] as const, {
          description: "branch (default): current branch only; session: all entries in the session file",
        }),
      ),
    }),
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const branch = ctx.sessionManager.getBranch();
      const all = ctx.sessionManager.getEntries();
      const branchIds = new Set(branch.map((entry) => entry.id));

      if (params.action === "search") {
        const query = (params.query ?? "").trim();
        if (!query) return err("context_history search requires 'query'.");
        const scope = params.scope === "session" ? "session" : "branch";
        const entries = scope === "session" ? all : branch;
        const matches = searchHistory(entries, {
          query,
          limit: clampInt(params.limit, 1, 50, DEFAULT_CONFIG.historySearchLimit),
        });
        if (matches.length === 0) {
          return { content: [{ type: "text", text: `No matches for "${query}" in ${scope} history.` }], details: { matches: 0 } };
        }
        const lines = matches.map(
          (m, i) =>
            `${i + 1}. [epoch ${m.epochIndex} · ref ${m.entryId} · ${m.kind}${m.toolName ? ` ${m.toolName}` : ""}] ${m.preview.replace(/\s+/g, " ").slice(0, 260)}`,
        );
        const text = [
          `Found ${matches.length} match(es) for "${query}" (scope: ${scope}). Use context_history read with a ref to see full entries.`,
          ...lines,
        ].join("\n");
        return { content: [{ type: "text", text }], details: { matches: matches.length } };
      }

      // read
      const ref = (params.ref ?? "").trim();
      if (!ref) return err("context_history read requires 'ref' (an entry id from search results).");
      const scope = params.scope === "session" ? "session" : "branch";
      const entries = scope === "session" ? all : branch;
      const { blocks, refFound } = readHistory(entries, {
        ref,
        around: clampInt(params.around, 0, 10, DEFAULT_CONFIG.historyReadAround),
        maxChars: controller.config.historyMaxChars,
      });
      if (!refFound || blocks.length === 0) {
        return err(
          `Entry "${ref}" not found in ${scope} history. Run a search to get valid refs; refs are entry ids like "a1b2c3d4".`,
        );
      }
      const text = formatHistoryRead(blocks);
      return {
        content: [{ type: "text", text: `Context history around ${ref} (scope: ${scope}):\n\n${text}` }],
        details: { ref, blocks: blocks.length },
      };
    },
  });

  // -------------------------------------------------------------------------
  // /epoch command
  // -------------------------------------------------------------------------
  pi.registerCommand("epoch", {
    description: "Context epochs: status, new epoch, notes, history, debug",
    getArgumentCompletions: (prefix: string) => {
      const options = ["new", "notes", "history", "debug"].filter((o) => o.startsWith(prefix));
      return options.length > 0 ? options.map((value) => ({ value, label: value })) : null;
    },
    handler: async (args: string, ctx: ExtensionCommandContext) => {
      const sub = args.trim().split(/\s+/)[0] ?? "";
      if (sub === "new") {
        await handleEpochNew(ctx);
        return;
      }
      if (sub === "notes") {
        await handleEpochNotes(ctx);
        return;
      }
      if (sub === "debug") {
        handleEpochDebug(ctx);
        return;
      }
      if (sub === "history") {
        handleEpochHistory(ctx);
        return;
      }
      handleEpochStatus(ctx);
    },
  });

  async function handleEpochNew(ctx: ExtensionCommandContext): Promise<void> {
    await ctx.waitForIdle();
    const outcome = await controller.manualRollover(ctx);
    if (outcome.performed) {
      notify(ctx, `New context epoch ${outcome.details?.epochIndex} started (no summary generated).`);
    } else if (outcome.reason?.includes("already at epoch")) {
      notify(ctx, `No rollover: ${outcome.reason}`);
    } else {
      notify(ctx, `Rollover failed: ${outcome.reason ?? "unknown error"}`, "warning");
    }
  }

  async function handleEpochNotes(ctx: ExtensionCommandContext): Promise<void> {
    const notes = replayNotes(ctx.sessionManager.getBranch());
    const items = listNotes(notes);
    if (items.length === 0) {
      notify(ctx, "No context notes in this session branch.");
      return;
    }
    const lines = ["Context notes:", ...items.map((n) => `- ${n.key} (${n.content.length} chars, entry ${n.entryId})`)];
    notify(ctx, lines.join("\n"));
  }

  function handleEpochStatus(ctx: ExtensionContext): void {
    const branch = ctx.sessionManager.getBranch();
    const state = getEpochState(branch);
    const usage = ctx.getContextUsage();
    const evaluation = evaluatePressure(usage, controller.config);
    const notes = replayNotes(branch);
    const last = controller.getLastRollover();

    const lines: string[] = ["Context Epochs"];
    lines.push(`Current epoch: ${state.epochIndex}${state.epochId ? ` (${state.epochId})` : ""}`);
    if (usage) {
      const tokens = usage.tokens !== null ? formatTokens(usage.tokens) : "unknown";
      lines.push(`Context: ${tokens} / ${formatTokens(usage.contextWindow)}`);
      if (evaluation.remaining !== null) lines.push(`Remaining: ${formatTokens(evaluation.remaining)}`);
      lines.push(`Pressure: ${evaluation.band}`);
    }
    lines.push(`History epochs: ${state.epochIndex}`);
    lines.push(`Notes: ${notes.size}`);
    if (last) {
      const ago = formatAge(Date.now() - last.at);
      lines.push(`Last rollover: ${ago} · ${last.details.trigger} · epoch ${last.details.epochIndex} (${last.details.epochId})`);
    } else if (state.boundary) {
      const ago = formatAge(Date.now() - new Date(state.boundary.timestamp).getTime());
      const details = state.boundary.details as { trigger?: string };
      lines.push(`Last rollover: ${ago} · ${details?.trigger ?? "?"} · epoch ${state.epochIndex} (${state.epochId})`);
    } else {
      lines.push("Last rollover: (none)");
    }
    notify(ctx, lines.join("\n"));
  }

  function handleEpochHistory(ctx: ExtensionContext): void {
    const branch = ctx.sessionManager.getBranch();
    const state = getEpochState(branch);
    if (state.boundaries.length === 0) {
      ctx.ui.notify("No context epoch boundaries yet (epoch 0).", "info");
      return;
    }
    const lines = ["Epoch boundaries (oldest first):"];
    state.boundaries.forEach((boundary, i) => {
      const details = boundary.details as { epochId?: string; trigger?: string } | undefined;
      lines.push(
        `- epoch ${i + 1}: entry ${boundary.id} · id ${details?.epochId ?? "?"} · trigger ${details?.trigger ?? "?"} · ${boundary.timestamp}`,
      );
    });
    notify(ctx, lines.join("\n"));
  }

  function handleEpochDebug(ctx: ExtensionContext): void {
    const branch = ctx.sessionManager.getBranch();
    const all = ctx.sessionManager.getEntries();
    const state = getEpochState(branch);
    const contextEntries = ctx.sessionManager.buildContextEntries();
    const activeTokens = controller.estimateActiveTokens(ctx);
    const usage = ctx.getContextUsage();
    const notes = replayNotes(branch);
    const historyItems = collectHistoryItems(branch);

    const lines: string[] = ["Context Epochs debug"];
    lines.push(`Session file: ${ctx.sessionManager.getSessionFile() ?? "(in-memory)"}`);
    lines.push(`Session id: ${ctx.sessionManager.getSessionId() ?? "?"}`);
    lines.push(`Branch entries (raw, durable): ${branch.length}`);
    lines.push(`Total entries in file: ${all.length}`);
    lines.push(`Active model context entries (post-boundary projection): ${contextEntries.length}`);
    lines.push(`Estimated active context tokens: ${activeTokens}${usage?.tokens !== null && usage ? ` (reported usage: ${usage.tokens})` : ""}`);
    lines.push(`Estimated full-branch tokens (incl. retired epochs): ${branch.reduce((sum, entry) => sum + entryChars(entry), 0) / 4 | 0}`);
    lines.push(`Epoch boundaries on branch: ${state.epochIndex}`);
    lines.push(`History index items (branch): ${historyItems.length}`);
    lines.push(`Notes: ${notes.size}`);
    lines.push(`Pending rollover: ${controller.isPending() ? "armed" : "no"}`);
    const conflicts = controller.getConflicts();
    lines.push(`Detected conflicting extensions: ${conflicts.length ? conflicts.join(", ") : "none"}`);
    if (state.boundary) {
      const details = state.boundary.details as { epochId?: string; trigger?: string; createdAt?: string } | undefined;
      lines.push(
        `Latest boundary: ${details?.epochId} · trigger ${details?.trigger} · created ${details?.createdAt} · firstKeptEntryId ${state.boundary.firstKeptEntryId}`,
      );
    }
    notify(ctx, lines.join("\n"));
  }

  function formatTokens(tokens: number): string {
    if (tokens >= 1000) return `${(tokens / 1000).toFixed(1)}k`;
    return String(tokens);
  }

  function formatAge(ms: number): string {
    const minutes = Math.floor(ms / 60_000);
    if (minutes < 1) return "just now";
    if (minutes < 60) return `${minutes}m ago`;
    const hours = Math.floor(minutes / 60);
    if (hours < 24) return `${hours}h ago`;
    return `${Math.floor(hours / 24)}d ago`;
  }
}

function entryChars(entry: SessionEntry): number {
  if (entry.type === "compaction" || entry.type === "branch_summary") return entry.summary?.length ?? 0;
  const content =
    entry.type === "message"
      ? (entry.message as { content?: unknown }).content
      : entry.type === "custom_message"
        ? entry.content
        : undefined;
  if (typeof content === "string") return content.length;
  if (Array.isArray(content)) {
    let chars = 0;
    for (const block of content) {
      if (block && typeof block === "object") {
        const b = block as { text?: string; thinking?: string; arguments?: unknown; name?: string };
        chars += (b.text?.length ?? 0) + (b.thinking?.length ?? 0) + (b.name?.length ?? 0) + JSON.stringify(b.arguments ?? {}).length;
      }
    }
    return chars;
  }
  return 0;
}

/** Notify via the TUI when available; fall back to stderr in headless modes. */
function notify(ctx: { hasUI: boolean; mode: string; ui: { notify(text: string, level?: string): void } }, text: string, level: "info" | "warning" = "info"): void {
  if (ctx.hasUI) {
    ctx.ui.notify(text, level);
  } else if (ctx.mode === "print" || ctx.mode === "json") {
    console.error(text);
  }
}

function clampInt(value: number | undefined, min: number, max: number, fallback: number): number {
  if (typeof value !== "number" || !Number.isFinite(value)) return fallback;
  return Math.min(max, Math.max(min, Math.round(value)));
}

function err(message: string): { content: Array<{ type: "text"; text: string }>; isError: true; details: Record<string, never> } {
  return { content: [{ type: "text", text: message }], isError: true, details: {} };
}