import { createHash, randomUUID } from "node:crypto";
import { Type } from "typebox";
import {
  buildSessionContext,
  convertToLlm,
  estimateTokens,
  generateSummaryWithUsage,
  sessionEntryToContextMessages,
  truncateHead,
  DEFAULT_MAX_BYTES,
  DEFAULT_MAX_LINES,
  type ExtensionAPI,
  type ExtensionContext,
  type SessionEntry,
} from "@earendil-works/pi-coding-agent";
import { DEFAULT_CONFIG, loadConfig } from "./src/config.js";
import {
  buildCodexHeaders,
  buildCompactionRequestBody,
  buildReplacementHistory,
  buildToolPayload,
  callRemoteCompaction,
  findNativeCheckpoint,
  isJsonObject,
  isOpenAICodexModel,
  mergeFeatureHeader,
  messagesToResponseItems,
  modelKey,
  NATIVE_COMPACTION_KIND,
  NATIVE_COMPACTION_VERSION,
  resolveCodexResponsesUrl,
  stripInputFromPayload,
  type JsonObject,
  type NativeCompactionDetails,
  type ResponseItem,
} from "./src/native-compaction.js";
import { SegmentIndex } from "./src/segment-index.js";
import { SemanticWorker } from "./src/semantic-worker.js";
import { shouldCompact } from "./src/token-accounting.js";
import {
  CompactionServiceRegistry,
  publishCompactionService,
  unpublishCompactionService,
} from "./src/service.js";
import {
  COMPACTION_CUSTOM_TYPE,
  COMPACTION_VERSION,
  QUERY_TOOL_NAME,
  type PiCompactionConfig,
  type ProjectionSnapshot,
  type ToolHistoryRecord,
} from "./src/types.js";

interface SessionRuntime {
  index: SegmentIndex;
  config: PiCompactionConfig;
  codexCache?: { fingerprint: string; modelKey: string; toolsHash: string; items: ResponseItem[] };
}

interface CachedPayload {
  modelKey: string;
  payload: JsonObject;
}

function errorMessage(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}

function hash(value: string): string {
  return createHash("sha256").update(value).digest("hex").slice(0, 20);
}

interface EnvironmentValue {
  hash: string;
  tokens: number;
  contextWindow: number;
}

interface EnvironmentCacheEntry extends EnvironmentValue {
  prompt: string;
  modelIdentity: string;
  activeNames: string[];
  tools: unknown[];
}

function environment(
  pi: ExtensionAPI,
  ctx: ExtensionContext,
  cache: Map<string, EnvironmentCacheEntry>,
): EnvironmentValue {
  const sessionId = ctx.sessionManager.getSessionId();
  const prompt = ctx.getSystemPrompt();
  const activeNames = pi.getActiveTools();
  const active = new Set(activeNames);
  const tools = pi.getAllTools().filter((tool) => active.has(tool.name));
  const contextWindow = ctx.model?.contextWindow ?? 0;
  const modelIdentity = ctx.model ? modelKey(ctx.model) : "none";
  const previous = cache.get(sessionId);
  if (
    previous
    && previous.prompt === prompt
    && previous.contextWindow === contextWindow
    && previous.modelIdentity === modelIdentity
    && previous.activeNames.length === activeNames.length
    && previous.activeNames.every((name, index) => name === activeNames[index])
    && previous.tools.length === tools.length
    && previous.tools.every((tool, index) => tool === tools[index])
  ) return previous;

  const serializedTools = JSON.stringify(tools.map((tool) => ({
    name: tool.name,
    description: tool.description,
    parameters: tool.parameters,
  })));
  const text = `${modelIdentity}\n${prompt}\n${serializedTools}`;
  const value: EnvironmentCacheEntry = {
    prompt,
    modelIdentity,
    activeNames: [...activeNames],
    tools,
    hash: hash(text),
    tokens: Math.ceil(text.length / 4),
    contextWindow,
  };
  cache.set(sessionId, value);
  return value;
}

function portableSummary(messages: readonly any[], maxChars = 24_000): string {
  const sections: string[] = ["Portable context checkpoint generated by pi-compaction."];
  for (const message of messages) {
    if (message?.role === "toolResult") {
      const text = Array.isArray(message.content)
        ? message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n")
        : "";
      if (text) sections.push(`[Tool ${message.toolName ?? "unknown"}] ${text}`);
      continue;
    }
    if (message?.role === "user") {
      const text = typeof message.content === "string"
        ? message.content
        : Array.isArray(message.content)
          ? message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n")
          : "";
      if (text) sections.push(`[User] ${text}`);
      continue;
    }
    if (message?.role === "assistant" && Array.isArray(message.content)) {
      const text = message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n");
      if (text) sections.push(`[Assistant] ${text}`);
    }
  }
  const combined = sections.join("\n\n");
  if (combined.length <= maxChars) return combined;
  const head = Math.floor(maxChars * 0.35);
  const tail = maxChars - head - 32;
  return `${combined.slice(0, head)}\n\n… checkpoint truncated …\n\n${combined.slice(-tail)}`;
}

function asPortableMessage(summary: string): any {
  return {
    role: "user",
    content: [{ type: "text", text: summary }],
    timestamp: Date.now(),
  };
}

function setFeatureHeader(headers: Record<string, string | null>): void {
  const existing = Object.entries(headers).find(([name]) => name.toLowerCase() === "x-codex-beta-features");
  if (existing) headers[existing[0]] = mergeFeatureHeader(existing[1]);
  else headers["x-codex-beta-features"] = mergeFeatureHeader(undefined);
}

function cleanHeaders(headers: Record<string, string | null> | undefined): Record<string, string> | undefined {
  if (!headers) return undefined;
  return Object.fromEntries(Object.entries(headers).filter((entry): entry is [string, string] => typeof entry[1] === "string"));
}

function activeToolsHash(pi: ExtensionAPI): string {
  return hash(pi.getActiveTools().join("\0"));
}

export default function piCompactionExtension(pi: ExtensionAPI): void {
  const runtimes = new Map<string, SessionRuntime>();
  const payloads = new Map<string, CachedPayload>();
  const environments = new Map<string, EnvironmentCacheEntry>();
  const worker = new SemanticWorker();
  const service = new CompactionServiceRegistry();
  publishCompactionService(service);

  const runtimeFor = (ctx: ExtensionContext): SessionRuntime => {
    const sessionId = ctx.sessionManager.getSessionId();
    let runtime = runtimes.get(sessionId);
    if (!runtime) {
      const index = new SegmentIndex(sessionId);
      index.hydrate(ctx.sessionManager.getBranch() as SessionEntry[], DEFAULT_CONFIG);
      runtime = { index, config: structuredClone(DEFAULT_CONFIG) };
      runtimes.set(sessionId, runtime);
      service.setSession(sessionId, (ids) => index.query(ids));
    }
    return runtime;
  };

  const persist = (record: unknown): void => {
    pi.appendEntry(COMPACTION_CUSTOM_TYPE, record);
  };

  const drainSemantic = (ctx: ExtensionContext, runtime: SessionRuntime): void => {
    const sessionId = ctx.sessionManager.getSessionId();
    for (const result of worker.drain(sessionId)) {
      const record = runtime.index.applySemantic(result.sourceKey, result.sourceHash, result.summary);
      if (record) persist(record);
    }
    runtime.index.setSemanticJobs(worker.count(sessionId));
  };

  const projectMessages = (
    ctx: ExtensionContext,
    runtime: SessionRuntime,
    messages: readonly any[],
  ): ProjectionSnapshot => {
    const env = environment(pi, ctx, environments);
    runtime.index.setSemanticJobs(worker.count(ctx.sessionManager.getSessionId()));
    const snapshot = runtime.index.project(messages, {
      environmentHash: env.hash,
      environmentTokens: env.tokens,
      contextWindow: env.contextWindow,
    });
    service.setSnapshot(snapshot);
    return snapshot;
  };

  const branchMessages = (
    ctx: ExtensionContext,
    runtime: SessionRuntime,
    branch: SessionEntry[],
    excludeLastAssistantError = false,
  ): { snapshot: ProjectionSnapshot; checkpoint: ReturnType<typeof findNativeCheckpoint> } => {
    const checkpoint = findNativeCheckpoint(branch);
    let messages: any[];
    if (checkpoint.status === "valid") {
      messages = branch
        .slice(checkpoint.checkpoint.entryIndex + 1)
        .flatMap((entry) => sessionEntryToContextMessages(entry));
      if (checkpoint.checkpoint.details.modelKey !== (ctx.model ? modelKey(ctx.model) : "")) {
        messages = [asPortableMessage(checkpoint.checkpoint.details.portableSummary), ...messages];
      }
    } else {
      messages = buildSessionContext(branch).messages;
    }
    if (excludeLastAssistantError) {
      const index = messages.findLastIndex((message) => message.role === "assistant");
      if (index >= 0 && messages[index]?.stopReason === "error") messages = messages.toSpliced(index, 1);
    }
    return { snapshot: projectMessages(ctx, runtime, messages), checkpoint };
  };

  const codexInput = (
    piModel: any,
    runtime: SessionRuntime,
    snapshot: ProjectionSnapshot,
    checkpoint?: NativeCompactionDetails,
  ): ResponseItem[] => {
    const toolsHash = activeToolsHash(pi);
    const key = modelKey(piModel);
    const cached = runtime.codexCache;
    if (cached && cached.fingerprint === snapshot.fingerprint && cached.modelKey === key && cached.toolsHash === toolsHash) {
      return cached.items;
    }
    const tail = messagesToResponseItems(piModel, convertToLlm(snapshot.messages as any), pi.getAllTools());
    const items = checkpoint && checkpoint.modelKey === key
      ? [...structuredClone(checkpoint.replacementHistory), ...tail]
      : tail;
    runtime.codexCache = { fingerprint: snapshot.fingerprint, modelKey: key, toolsHash, items };
    snapshot.codexItems = items;
    return items;
  };

  pi.on("session_start", async (_event, ctx) => {
    publishCompactionService(service);
    const sessionId = ctx.sessionManager.getSessionId();
    const config = await loadConfig();
    const index = new SegmentIndex(sessionId);
    index.hydrate(ctx.sessionManager.getBranch() as SessionEntry[], config);
    runtimes.set(sessionId, { index, config });
    service.setSession(sessionId, (ids) => index.query(ids));
  });

  pi.on("session_shutdown", (_event, ctx) => {
    const sessionId = ctx.sessionManager.getSessionId();
    worker.cancelSession(sessionId);
    runtimes.delete(sessionId);
    payloads.delete(sessionId);
    environments.delete(sessionId);
    service.deleteSession(sessionId);
    if (runtimes.size === 0) unpublishCompactionService(service);
  });

  pi.on("session_tree", (_event, ctx) => {
    const runtime = runtimeFor(ctx);
    runtime.index.hydrate(ctx.sessionManager.getBranch() as SessionEntry[], runtime.config);
    runtime.codexCache = undefined;
    service.clearSnapshot(ctx.sessionManager.getSessionId());
  });

  pi.on("session_compact", (_event, ctx) => {
    const runtime = runtimeFor(ctx);
    runtime.index.hydrate(ctx.sessionManager.getBranch() as SessionEntry[], runtime.config);
    runtime.codexCache = undefined;
    service.clearSnapshot(ctx.sessionManager.getSessionId());
    payloads.delete(ctx.sessionManager.getSessionId());
  });

  pi.on("model_select", (_event, ctx) => {
    const sessionId = ctx.sessionManager.getSessionId();
    payloads.delete(sessionId);
    environments.delete(sessionId);
    const runtime = runtimes.get(sessionId);
    if (runtime) runtime.codexCache = undefined;
  });

  pi.on("turn_end", (event, ctx) => {
    const runtime = runtimeFor(ctx);
    drainSemantic(ctx, runtime);
    const previous = runtime.index.getSnapshot();
    const captured = runtime.index.captureTurn(event.message, event.toolResults, runtime.config);
    for (const record of captured) persist(runtime.index.toPersisted(record));
    const pressure = previous && previous.contextWindow > 0
      ? previous.upperBound / previous.contextWindow >= runtime.config.softRatio
      : false;
    const semanticCandidates = pressure
      ? runtime.index.allRecords()
      : captured.filter((record) => record.rawText.length >= runtime.config.semantic.maxInputChars);
    for (const record of semanticCandidates) {
      worker.enqueue(ctx.sessionManager.getSessionId(), record, ctx, runtime.config);
    }
    runtime.index.setSemanticJobs(worker.count(ctx.sessionManager.getSessionId()));

    const usage = ctx.getContextUsage();
    if (usage?.tokens && previous) {
      const env = environment(pi, ctx, environments);
      runtime.index.recordUsage(
        usage.tokens,
        previous.messageEstimate + estimateTokens(event.message as any),
        env.hash,
      );
    }
  });

  pi.on("context", (event, ctx) => {
    const runtime = runtimeFor(ctx);
    const branch = ctx.sessionManager.getBranch() as SessionEntry[];
    const checkpoint = findNativeCheckpoint(branch);
    if (checkpoint.status === "invalid") {
      ctx.abort();
      if (ctx.hasUI) ctx.ui.notify("pi-compaction checkpoint is malformed; request blocked.", "error");
      return { messages: [] };
    }
    if (checkpoint.status === "valid") {
      const projected = branchMessages(ctx, runtime, branch);
      return { messages: projected.snapshot.messages as any };
    }
    return { messages: projectMessages(ctx, runtime, event.messages).messages as any };
  });

  pi.on("before_provider_headers", (event, ctx) => {
    if (isOpenAICodexModel(ctx.model)) setFeatureHeader(event.headers);
  });

  pi.on("before_provider_request", (event, ctx) => {
    const model = ctx.model;
    if (!model || !isOpenAICodexModel(model) || !isJsonObject(event.payload)) return undefined;
    const runtime = runtimeFor(ctx);
    const sessionId = ctx.sessionManager.getSessionId();
    payloads.set(sessionId, { modelKey: modelKey(model), payload: stripInputFromPayload(event.payload) });
    const checkpoint = findNativeCheckpoint(ctx.sessionManager.getBranch() as SessionEntry[]);
    if (checkpoint.status !== "valid" || checkpoint.checkpoint.details.modelKey !== modelKey(model)) return undefined;
    const snapshot = runtime.index.getSnapshot();
    if (!snapshot) return undefined;
    const payload: JsonObject = {
      ...event.payload,
      input: codexInput(model, runtime, snapshot, checkpoint.checkpoint.details),
    };
    delete payload.messages;
    delete payload.previous_response_id;
    return payload;
  });

  pi.on("session_before_compact", async (event, ctx) => {
    const runtime = runtimeFor(ctx);
    if (!runtime.config.enabled) return undefined;
    drainSemantic(ctx, runtime);
    const { snapshot, checkpoint } = branchMessages(
      ctx,
      runtime,
      event.branchEntries as SessionEntry[],
      event.reason === "overflow" && event.willRetry,
    );

    if (!shouldCompact(event.reason, snapshot, runtime.config)) return { cancel: true };

    const model = ctx.model;
    if (!model) return undefined;
    if (!runtime.config.native.enabled || !isOpenAICodexModel(model)) {
      try {
        const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
        if (!auth.ok) return undefined;
        const result = await generateSummaryWithUsage(
          snapshot.messages as any,
          model,
          16_384,
          auth.apiKey,
          cleanHeaders(auth.headers),
          event.signal,
          event.customInstructions,
          undefined,
          runtime.config.semantic.thinking,
        );
        return {
          compaction: {
            summary: result.text,
            firstKeptEntryId: event.preparation.firstKeptEntryId,
            tokensBefore: snapshot.effectiveTokens,
            usage: result.usage,
            details: {
              kind: COMPACTION_CUSTOM_TYPE,
              version: COMPACTION_VERSION,
              checkpoint: "portable",
              sourceFingerprint: snapshot.fingerprint,
            },
          },
        };
      } catch (error) {
        if (!event.signal.aborted) console.error("pi-compaction: portable checkpoint failed:", error);
        return undefined;
      }
    }

    try {
      const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
      if (!auth.ok || !auth.apiKey) throw new Error(auth.ok ? "Codex authentication unavailable" : auth.error);
      const currentCheckpoint = checkpoint.status === "valid" ? checkpoint.checkpoint.details : undefined;
      const input = codexInput(model, runtime, snapshot, currentCheckpoint);
      const sessionId = ctx.sessionManager.getSessionId();
      const cached = payloads.get(sessionId);
      const body = buildCompactionRequestBody({
        basePayload: cached?.modelKey === modelKey(model) ? cached.payload : undefined,
        model,
        input,
        instructions: ctx.getSystemPrompt(),
        tools: buildToolPayload(pi.getAllTools(), pi.getActiveTools()),
        sessionId,
      });
      if (ctx.hasUI) ctx.ui.notify("pi-compaction: creating Codex native checkpoint…", "info");
      const remote = await callRemoteCompaction({
        url: resolveCodexResponsesUrl(model.baseUrl),
        headers: buildCodexHeaders({ apiKey: auth.apiKey, headers: cleanHeaders(auth.headers), sessionId }),
        body,
        signal: event.signal,
      });
      const details: NativeCompactionDetails = {
        kind: NATIVE_COMPACTION_KIND,
        version: NATIVE_COMPACTION_VERSION,
        checkpoint: "native",
        modelKey: modelKey(model),
        replacementHistory: buildReplacementHistory(
          input,
          remote.compactionItem,
          Math.min(64_000, Math.max(4_096, Math.floor(snapshot.contextWindow * Math.max(0.05, runtime.config.targetRatio - 0.3)))),
        ),
        portableSummary: portableSummary(snapshot.messages),
        sourceFingerprint: snapshot.fingerprint,
      };
      return {
        compaction: {
          summary: `Pi native checkpoint ${randomUUID()}`,
          firstKeptEntryId: event.preparation.firstKeptEntryId,
          tokensBefore: snapshot.effectiveTokens,
          usage: remote.usage,
          details,
        },
      };
    } catch (error) {
      if (!event.signal.aborted && ctx.hasUI) {
        ctx.ui.notify(`pi-compaction native checkpoint failed; falling back to Pi: ${errorMessage(error)}`, "warning");
      }
      return undefined;
    }
  });

  pi.on("session_before_tree", async (event, ctx) => {
    const model = ctx.model;
    if (!event.preparation.userWantsSummary || !model) return undefined;
    const runtime = runtimeFor(ctx);
    if (!runtime.config.enabled) return undefined;
    const messages = event.preparation.entriesToSummarize.flatMap((entry) => sessionEntryToContextMessages(entry));
    const snapshot = projectMessages(ctx, runtime, messages);
    try {
      const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
      if (!auth.ok) return undefined;
      const result = await generateSummaryWithUsage(
        snapshot.messages as any,
        model,
        16_384,
        auth.apiKey,
        cleanHeaders(auth.headers),
        event.signal,
        event.preparation.customInstructions,
        undefined,
        runtime.config.semantic.thinking,
        undefined,
        undefined,
        undefined,
        undefined,
      );
      return {
        summary: {
          summary: result.text,
          usage: result.usage,
          details: {
            kind: COMPACTION_CUSTOM_TYPE,
            version: COMPACTION_VERSION,
            checkpoint: "portable",
            sourceFingerprint: snapshot.fingerprint,
          },
        },
      };
    } catch (error) {
      if (!event.signal.aborted) console.error("pi-compaction: branch summary failed:", error);
      return undefined;
    }
  });

  pi.registerTool({
    name: QUERY_TOOL_NAME,
    label: "Query Compacted Tool History",
    description: "Recover original tool outputs by source key or tool call id after pi-compaction replaced them with bounded summaries.",
    promptSnippet: "Recover original compacted tool evidence",
    promptGuidelines: [
      "Use context_tree_query only when a pi-compaction summary lacks evidence needed for the current task.",
    ],
    parameters: Type.Object({
      toolCallIds: Type.Array(Type.String(), { minItems: 1, maxItems: 20 }),
    }),
    async execute(_id, params, _signal, _update, ctx) {
      const records = service.query(ctx.sessionManager.getSessionId(), params.toolCallIds);
      const blocks = records.map((record: ToolHistoryRecord) => {
        const perResultBytes = Math.max(1_024, Math.floor(DEFAULT_MAX_BYTES / Math.max(1, records.length)));
        const args = truncateHead(JSON.stringify(record.args), { maxLines: 20, maxBytes: 2_048 }).content;
        const truncated = truncateHead(record.rawText, {
          maxLines: Math.max(20, Math.floor(DEFAULT_MAX_LINES / Math.max(1, records.length))),
          maxBytes: perResultBytes,
        });
        return [
          `## ${record.sourceKey}`,
          `Tool: ${record.toolName}`,
          `Args: ${args}`,
          `Status: ${record.isError ? "ERROR" : "OK"}`,
          "",
          truncated.content,
          ...(truncated.truncated ? [`[truncated: ${truncated.outputLines}/${truncated.totalLines} lines]`] : []),
        ].join("\n");
      });
      const missing = params.toolCallIds.filter((id: string) => !records.some((record) => record.sourceKey === id || record.toolCallId === id));
      for (const id of missing) blocks.push(`## ${id}\n(not found)`);
      return {
        content: [{ type: "text", text: blocks.join("\n\n---\n\n") }],
        details: { sourceKeys: records.map((record) => record.sourceKey), missing },
      };
    },
  });

  pi.registerCommand("compaction-status", {
    description: "Show effective/raw context usage and pi-compaction state",
    handler: async (_args, ctx) => {
      const metrics = service.getMetrics(ctx.sessionManager.getSessionId());
      const text = metrics
        ? [
            `Effective: ${metrics.effectiveTokens}/${metrics.contextWindow} (${metrics.percent.toFixed(1)}%, ${metrics.confidence})`,
            `Raw: ${metrics.rawTokens}/${metrics.contextWindow}`,
            `Reclaimed: ${metrics.reclaimedTokens}`,
            `Segments: ${metrics.segmentCount}; semantic jobs: ${metrics.semanticJobs}; generation: ${metrics.generation}`,
          ].join("\n")
        : "pi-compaction has no projection snapshot yet.";
      if (ctx.hasUI) ctx.ui.notify(text, "info");
      else process.stderr.write(`${text}\n`);
    },
  });

}