import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { randomUUID } from "node:crypto";
import { readFile, writeFile } from "node:fs/promises";
import { join } from "node:path";
import {
  createContextViewPolicy,
  fencedHistoryBlocks,
} from "../src/context-view-policy.js";
import { createDelegateRunner } from "../src/delegate-runner.js";
import { createHistoryCompactor } from "../src/history-compactor.js";
import { createPiDelegateProcessAdapter } from "../src/pi-delegate-process.js";
import {
  createContextCostManager,
  type ContextCostManager,
  type ContextCostMode,
  type TextContent,
} from "../src/context-cost-manager.js";

const STATUS_KEY = "context-cost";

interface ContextCostSettings {
  mode: ContextCostMode;
  cacheTtl: "5m" | "1h" | number;
  initialResultTokenThreshold: number;
  previewTokenLimit: number;
}

const DEFAULT_SETTINGS: ContextCostSettings = {
  mode: "observe",
  cacheTtl: "5m",
  initialResultTokenThreshold: 12_000,
  previewTokenLimit: 2_000,
};

async function loadSettings(path: string): Promise<ContextCostSettings> {
  try {
    const candidate = JSON.parse(await readFile(path, "utf8")) as Partial<ContextCostSettings>;
    const mode = candidate.mode;
    return {
      mode: mode === "safe" || mode === "aggressive" ? mode : "observe",
      cacheTtl:
        candidate.cacheTtl === "1h" ||
        candidate.cacheTtl === "5m" ||
        (typeof candidate.cacheTtl === "number" && candidate.cacheTtl > 0)
          ? candidate.cacheTtl
          : DEFAULT_SETTINGS.cacheTtl,
      initialResultTokenThreshold:
        typeof candidate.initialResultTokenThreshold === "number" &&
        candidate.initialResultTokenThreshold > 0
          ? candidate.initialResultTokenThreshold
          : DEFAULT_SETTINGS.initialResultTokenThreshold,
      previewTokenLimit:
        typeof candidate.previewTokenLimit === "number" &&
        candidate.previewTokenLimit > 0
          ? candidate.previewTokenLimit
          : DEFAULT_SETTINGS.previewTokenLimit,
    };
  } catch {
    return { ...DEFAULT_SETTINGS };
  }
}

function cacheTtlMilliseconds(value: ContextCostSettings["cacheTtl"]): number {
  if (typeof value === "number") return value;
  return value === "1h" ? 60 * 60_000 : 5 * 60_000;
}

export default function contextCostExtension(pi: ExtensionAPI): void {
  let manager: ContextCostManager | undefined;
  let mode: ContextCostMode = "observe";
  let lastProviderActivityAt = Date.now();
  let maintenanceRequested = false;
  let manualShake: "tools" | "all" | undefined;
  let settings = { ...DEFAULT_SETTINGS };
  let settingsPath: string | undefined;
  const contextViewPolicy = createContextViewPolicy();
  const delegateRunner = createDelegateRunner(createPiDelegateProcessAdapter());

  pi.on("session_start", async (_event, ctx) => {
    settingsPath = join(ctx.sessionManager.getSessionDir(), "context-cost.json");
    settings = await loadSettings(settingsPath);
    mode = settings.mode;
    lastProviderActivityAt = Date.now();
    maintenanceRequested = false;
    manualShake = undefined;
    manager = createContextCostManager({
      artifactRoot: join(
        ctx.sessionManager.getSessionDir(),
        "context-cost-artifacts",
      ),
      mode,
      initialResultTokenThreshold: settings.initialResultTokenThreshold,
      previewTokenLimit: settings.previewTokenLimit,
    });
    await manager.purgeExpired({
      olderThan: new Date(Date.now() - 30 * 24 * 60 * 60 * 1_000),
    });
    ctx.ui.setStatus(STATUS_KEY, `cost: ${mode}`);
  });

  pi.on("after_provider_response", async () => {
    lastProviderActivityAt = Date.now();
  });

  pi.on("context", async (event, ctx) => {
    const usage = ctx.getContextUsage();
    const contextWindow = ctx.model?.contextWindow;
    if (!usage || usage.tokens === null || !contextWindow) return;

    const requestedManualShake = manualShake;
    manualShake = undefined;
    const result = contextViewPolicy.apply({
      messages: event.messages,
      mode: requestedManualShake && mode === "observe" ? "safe" : mode,
      contextTokens: usage.tokens,
      contextWindow,
      cacheAgeMs: Date.now() - lastProviderActivityAt,
      cacheTtlMs: cacheTtlMilliseconds(settings.cacheTtl),
      manual: requestedManualShake,
    });
    if (
      result.telemetry.supersededResults === 0 &&
      result.telemetry.shakenResults === 0 &&
      result.telemetry.fencedBlocks === 0
    ) {
      return;
    }
    const label = mode === "observe" ? "potential" : "context saved";
    ctx.ui.setStatus(
      STATUS_KEY,
      `cost: ${label} ~${result.telemetry.estimatedTokensSaved.toLocaleString()} tokens`,
    );
    return mode === "observe" && !requestedManualShake
      ? undefined
      : { messages: result.messages };
  });

  pi.on("turn_end", async (_event, ctx) => {
    const usage = ctx.getContextUsage();
    const contextWindow = ctx.model?.contextWindow;
    if (
      mode === "observe" ||
      maintenanceRequested ||
      !usage ||
      usage.tokens === null ||
      !contextWindow
    ) {
      return;
    }
    const trigger = mode === "aggressive" ? 0.55 : 0.65;
    if (usage.tokens / contextWindow < trigger) return;
    maintenanceRequested = true;
    ctx.compact({
      customInstructions: "context-cost automatic threshold",
      onComplete: () => { maintenanceRequested = false; },
      onError: () => { maintenanceRequested = false; },
    });
  });

  pi.on("session_compact", async () => {
    maintenanceRequested = false;
  });

  pi.on("session_compact_failed", async () => {
    maintenanceRequested = false;
  });

  pi.on("session_before_compact", async (event, ctx) => {
    if (mode === "observe") return;
    const configured = process.env.PI_CONTEXT_COST_SUMMARIZER_MODEL;
    const available = ctx.modelRegistry.getAvailable();
    const summarizerModel = configured
      ? available.find(
          (model) => `${model.provider}/${model.id}` === configured,
        )
      : available.find((model) => model.id.toLowerCase().includes("haiku"));
    if (!summarizerModel) {
      ctx.ui.notify(
        "Context cost: no Haiku summarizer model is available; pi compaction will run.",
        "warning",
      );
      return;
    }

    const compactor = createHistoryCompactor({
      summarize: async ({ prompt, signal }) => {
        const response = await ctx.modelRegistry.complete(
          summarizerModel,
          {
            messages: [
              {
                role: "user",
                content: [{ type: "text", text: prompt }],
                timestamp: Date.now(),
              },
            ],
          },
          {
            maxTokens: 8_192,
            signal,
            cacheRetention: "none",
            sessionId: randomUUID(),
          },
        );
        manager?.recordUsage(response.usage);
        return {
          text: response.content
            .filter((block): block is { type: "text"; text: string } =>
              block.type === "text",
            )
            .map((block) => block.text)
            .join("\n"),
          usage: response.usage,
        };
      },
    });
    const result = await compactor.compact({
      messages: [
        ...event.preparation.messagesToSummarize,
        ...event.preparation.turnPrefixMessages,
      ],
      firstKeptEntryId: event.preparation.firstKeptEntryId,
      tokensBefore: event.preparation.tokensBefore,
      signal: event.signal,
    });
    return result ? { compaction: result as any } : undefined;
  });

  pi.on("message_end", async (event) => {
    if (manager && event.message.role === "assistant") {
      manager.recordUsage(event.message.usage);
    }
  });

  pi.on("tool_result", async (event, ctx) => {
    const textContent = event.content.filter(
      (block): block is TextContent => block.type === "text",
    );
    if (!manager || textContent.length !== event.content.length) {
      return;
    }

    const processed = await manager.processToolResult({
      sessionId: ctx.sessionManager.getSessionId(),
      toolCallId: event.toolCallId,
      toolName: event.toolName,
      input: event.input as Record<string, unknown>,
      content: textContent,
      details: event.details,
      isError: event.isError,
    });

    if (processed.telemetry.estimatedTokensSaved > 0) {
      const label = processed.changed ? "saved" : "potential";
      ctx.ui.setStatus(
        STATUS_KEY,
        `cost: ${label} ~${processed.telemetry.estimatedTokensSaved.toLocaleString()} tokens`,
      );
    }

    if (processed.changed) {
      return { content: processed.content };
    }
  });

  if (typeof pi.registerTool === "function") {
    pi.registerTool({
      name: "delegate_context_work",
      label: "Delegate Context Work",
      description: "Run a bounded exploration task in an isolated cheap-model pi session. Returns concise findings and usage.",
      parameters: Type.Object({
        task: Type.String({ description: "Self-contained exploration task" }),
        model: Type.Optional(Type.String({ description: "Worker model as provider/model" })),
        tokenCeiling: Type.Optional(Type.Number({ minimum: 1, default: 50_000 })),
        costCeiling: Type.Optional(Type.Number({ minimum: 0, default: 0.25 })),
        tools: Type.Optional(Type.Array(Type.String())),
      }),
      async execute(_toolCallId, params, signal, _onUpdate, ctx) {
        const discovered = ctx.modelRegistry
          .getAvailable()
          .find((model) => model.id.toLowerCase().includes("haiku"));
        const model =
          params.model ??
          process.env.PI_CONTEXT_COST_WORKER_MODEL ??
          (discovered ? `${discovered.provider}/${discovered.id}` : undefined);
        if (!model) {
          return {
            content: [{ type: "text", text: "No worker model is configured. Set PI_CONTEXT_COST_WORKER_MODEL or pass model." }],
            details: { configured: false },
          };
        }
        const result = await delegateRunner.run({
          task: params.task,
          model,
          cwd: ctx.cwd,
          tools: params.tools ?? ["read", "grep", "find", "ls", "bash"],
          tokenCeiling: params.tokenCeiling ?? 50_000,
          costCeiling: params.costCeiling ?? 0.25,
          signal,
        });
        manager?.recordUsage({
          input: result.usage.input,
          output: result.usage.output,
          cacheRead: result.usage.cacheRead,
          cacheWrite: result.usage.cacheWrite,
          cost: { total: result.usage.cost },
        });
        const limitNotice = result.aborted
          ? "\n\n[Delegate aborted.]"
          : result.limitExceeded
            ? `\n\n[Delegate stopped: ${result.limitExceeded} ceiling exceeded.]`
            : "";
        const output = `${result.output || "(no output)"}${limitNotice}`;
        return {
          content: [{ type: "text", text: output.slice(0, 50 * 1024) }],
          details: { model, ...result },
          usage: {
            input: result.usage.input,
            output: result.usage.output,
            cacheRead: result.usage.cacheRead,
            cacheWrite: result.usage.cacheWrite,
            totalTokens:
              result.usage.input +
              result.usage.output +
              result.usage.cacheRead +
              result.usage.cacheWrite,
            cost: {
              input: 0,
              output: 0,
              cacheRead: 0,
              cacheWrite: 0,
              total: result.usage.cost,
            },
          },
        };
      },
    });
  }

  pi.registerCommand("cost-mode", {
    description: "Set context-cost mode (observe, safe, or aggressive)",
    handler: async (args, ctx) => {
      const requestedMode = args.trim() as ContextCostMode;
      if (
        requestedMode !== "observe" &&
        requestedMode !== "safe" &&
        requestedMode !== "aggressive"
      ) {
        ctx.ui.notify("Usage: /cost-mode <observe|safe|aggressive>", "warning");
        return;
      }
      if (!manager) {
        ctx.ui.notify("Context cost: session has not started", "warning");
        return;
      }

      mode = requestedMode;
      settings.mode = requestedMode;
      manager.setMode(requestedMode);
      if (settingsPath) {
        await writeFile(settingsPath, `${JSON.stringify(settings, null, 2)}\n`, {
          encoding: "utf8",
          mode: 0o600,
        });
      }
      ctx.ui.setStatus(STATUS_KEY, `cost: ${requestedMode}`);
      ctx.ui.notify(`Context cost mode: ${requestedMode}`, "info");
    },
  });

  pi.registerCommand("cost-shake", {
    description: "Slim recoverable history on the next model request",
    handler: async (args, ctx) => {
      const option = args.trim();
      if (option !== "" && option !== "--all") {
        ctx.ui.notify("Usage: /cost-shake [--all]", "warning");
        return;
      }
      manualShake = option === "--all" ? "all" : "tools";
      ctx.ui.notify(
        `Context cost: manual ${manualShake} shake armed for the next model request`,
        "info",
      );
    },
  });

  pi.registerCommand("cost-compact", {
    description: "Compact context with the configured cheap model",
    handler: async (_args, ctx) => {
      ctx.compact({
        customInstructions: "context-cost manual compaction",
        onComplete: () => { maintenanceRequested = false; },
        onError: () => { maintenanceRequested = false; },
      });
    },
  });

  pi.registerCommand("cost-purge", {
    description: "Delete context-cost artifacts for the active session",
    handler: async (_args, ctx) => {
      if (!manager) {
        ctx.ui.notify("Context cost: session has not started", "warning");
        return;
      }
      await manager.purge({ sessionId: ctx.sessionManager.getSessionId() });
      ctx.ui.notify("Context-cost artifacts purged for this session", "info");
    },
  });

  pi.registerCommand("cost-restore", {
    description: "Restore archived tool output into future context",
    handler: async (args, ctx) => {
      const artifactId = args.trim();
      if (!manager || !artifactId) {
        ctx.ui.notify("Usage: /cost-restore <artifact-id>", "warning");
        return;
      }

      try {
        const historyPrefix = "cost-history://";
        const blockPrefix = "cost-history-block://";
        const historyToolCallId = artifactId.startsWith(historyPrefix)
          ? artifactId.slice(historyPrefix.length)
          : undefined;
        const requestedBlockId = artifactId.startsWith(blockPrefix)
          ? artifactId.slice(blockPrefix.length)
          : undefined;
        let content: string;
        let details: Record<string, string>;
        if (requestedBlockId) {
          const block = ctx.sessionManager
            .getBranch()
            .flatMap((entry) =>
              entry.type === "message" &&
              (entry.message.role === "user" || entry.message.role === "assistant")
                ? fencedHistoryBlocks(entry.message.content)
                : [],
            )
            .find((candidate) => candidate.id === requestedBlockId);
          if (!block) {
            throw new Error(`History block not found: ${requestedBlockId}`);
          }
          content = block.text;
          details = { historyBlockId: requestedBlockId };
        } else if (historyToolCallId) {
          const entry = [...ctx.sessionManager.getBranch()]
            .reverse()
            .find(
              (candidate) =>
                candidate.type === "message" &&
                candidate.message.role === "toolResult" &&
                candidate.message.toolCallId === historyToolCallId,
            );
          if (
            !entry ||
            entry.type !== "message" ||
            entry.message.role !== "toolResult"
          ) {
            throw new Error(`History result not found: ${historyToolCallId}`);
          }
          content = entry.message.content
            .filter((block): block is TextContent => block.type === "text")
            .map((block) => block.text)
            .join("\n");
          details = { historyToolCallId };
        } else {
          content = await manager.restore({
            sessionId: ctx.sessionManager.getSessionId(),
            artifactId,
          });
          details = { artifactId };
        }
        pi.sendMessage(
          {
            customType: "context-cost-restored",
            content,
            display: true,
            details,
          },
          { deliverAs: "nextTurn" },
        );
        ctx.ui.notify(`Restored context-cost reference ${artifactId}`, "info");
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        ctx.ui.notify(`Could not restore artifact: ${message}`, "error");
      }
    },
  });

  pi.registerCommand("cost-status", {
    description: "Show context-cost savings for this session",
    handler: async (_args, ctx) => {
      const status = manager?.getStatus();
      const provider = status?.provider;
      const usage = provider
        ? `${provider.requests.toLocaleString()} ${provider.requests === 1 ? "request" : "requests"} | input ${provider.inputTokens.toLocaleString()} | cache read ${provider.cacheReadTokens.toLocaleString()} | cache write ${provider.cacheWriteTokens.toLocaleString()} | output ${provider.outputTokens.toLocaleString()} | cost $${provider.cost.toFixed(4)}`
        : "0 requests";
      const message = status
        ? `Context cost: ${usage} | ${status.archivedResults} archived | ~${status.estimatedTokensSaved.toLocaleString()} tokens saveable/saved`
        : "Context cost: session has not started";
      ctx.ui.notify(message, "info");
    },
  });
}