import type { SessionEntry } from "@earendil-works/pi-coding-agent";
import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
import { estimateTokens, stableStringify, toBoundedJson, truncateMiddle } from "./json.js";

const MESSAGE_BUDGET = 10_000;
const TOOL_BUDGET = 10_000;
const MESSAGE_ENTRY_CAP = 2_000;
const TOOL_ENTRY_CAP = 1_000;
const MAX_RECENT_NON_USER = 40;

export type TranscriptKind =
  | "user"
  | "assistant"
  | "tool-call"
  | "tool-result"
  | "compaction-summary"
  | "branch-summary";

interface Evidence {
  number: number;
  kind: TranscriptKind;
  budget: "message" | "tool";
  text: string;
  tokens: number;
}

export interface CompactTranscript {
  text: string;
  omitted: boolean;
  retainedEntries: number;
}

function contentText(content: string | (TextContent | ImageContent)[]): string {
  if (typeof content === "string") return content;
  return content
    .map((item) => (item.type === "text" ? item.text : `<image mime_type="${item.mimeType}" omitted="true" />`))
    .join("\n");
}

function collect(entries: readonly SessionEntry[]): Evidence[] {
  const evidence: Evidence[] = [];
  let number = 0;

  const add = (kind: TranscriptKind, budget: "message" | "tool", text: string): void => {
    number += 1;
    const cap = budget === "message" ? MESSAGE_ENTRY_CAP : TOOL_ENTRY_CAP;
    const bounded = truncateMiddle(text, cap).text;
    evidence.push({ number, kind, budget, text: bounded, tokens: estimateTokens(bounded) });
  };

  for (const entry of entries) {
    if (entry.type === "compaction") {
      add("compaction-summary", "message", entry.summary);
      continue;
    }
    if (entry.type === "branch_summary") {
      add("branch-summary", "message", entry.summary);
      continue;
    }
    if (entry.type !== "message") continue;
    const message = entry.message;
    if (message.role === "user") {
      add("user", "message", contentText(message.content));
      continue;
    }
    if (message.role === "assistant") {
      const assistantText = message.content
        .filter((item): item is TextContent => item.type === "text")
        .map((item) => item.text)
        .join("\n");
      if (assistantText.length > 0) add("assistant", "message", assistantText);
      for (const item of message.content) {
        if (item.type !== "toolCall") continue;
        const bounded = toBoundedJson({ name: item.name, arguments: item.arguments }, TOOL_ENTRY_CAP);
        add("tool-call", "tool", stableStringify(bounded.value));
      }
      continue;
    }
    if (message.role === "toolResult") {
      add("tool-result", "tool", `${message.toolName} (${message.isError ? "error" : "success"}):\n${contentText(message.content)}`);
    }
    // Custom/extension and hidden-thinking context is intentionally excluded.
  }
  return evidence;
}

function addIfFits(selected: Set<number>, entry: Evidence, budget: { message: number; tool: number }): boolean {
  if (selected.has(entry.number) || entry.tokens > budget[entry.budget]) return false;
  selected.add(entry.number);
  budget[entry.budget] -= entry.tokens;
  return true;
}

export function compactTranscript(
  contextEntries: readonly SessionEntry[],
  branchEntries: readonly SessionEntry[],
): CompactTranscript {
  const contextIds = new Set(contextEntries.map((entry) => entry.id));
  const historicalUsers = branchEntries.filter(
    (entry) => entry.type === "message" && entry.message.role === "user" && !contextIds.has(entry.id),
  );
  const evidence = collect([...historicalUsers, ...contextEntries]);
  const selected = new Set<number>();
  const budget = { message: MESSAGE_BUDGET, tool: TOOL_BUDGET };
  const users = evidence.filter((entry) => entry.kind === "user");
  const latestCompactionSummary = evidence.filter((entry) => entry.kind === "compaction-summary").at(-1);
  const latestBranchSummary = evidence.filter((entry) => entry.kind === "branch-summary").at(-1);

  if (latestCompactionSummary !== undefined) addIfFits(selected, latestCompactionSummary, budget);
  if (latestBranchSummary !== undefined) addIfFits(selected, latestBranchSummary, budget);

  const firstUser = users[0];
  const latestUser = users.at(-1);
  if (firstUser !== undefined) addIfFits(selected, firstUser, budget);
  if (latestUser !== undefined) addIfFits(selected, latestUser, budget);
  for (let index = users.length - 2; index >= 1; index -= 1) {
    addIfFits(selected, users[index]!, budget);
  }

  let nonUserCount = 0;
  for (let index = evidence.length - 1; index >= 0 && nonUserCount < MAX_RECENT_NON_USER; index -= 1) {
    const entry = evidence[index]!;
    if (entry.kind === "user") continue;
    if (addIfFits(selected, entry, budget)) nonUserCount += 1;
  }

  const retained = evidence.filter((entry) => selected.has(entry.number));
  const omitted = retained.length !== evidence.length;
  const lines = retained.map(
    (entry) => `<entry number="${entry.number}" kind="${entry.kind}">\n${entry.text}\n</entry>`,
  );
  if (omitted) lines.push("<omission>Some conversation entries were omitted.</omission>");
  return { text: lines.join("\n\n"), omitted, retainedEntries: retained.length };
}