import { createHash } from "node:crypto";
import type { Message } from "@earendil-works/pi-ai/compat";
import type { BtwEntry, BtwSummary, BtwSummarySourceRef, BtwThread } from "./threads.ts";

export type SummaryPolicy = { triggerTokens: number; retainTokens: number; inputMaxTokens: number };
export type SummarySnapshot = {
  source: BtwSummarySourceRef[];
  sourceHash: string;
  throughEntryId: string;
  input: string;
};

/** Deliberately simple, deterministic estimate used for BTW-only context policy. */
export function estimateTokens(text: string): number {
  return text === "" ? 0 : Math.ceil(Buffer.byteLength(text, "utf8") / 3);
}

function stable(value: unknown): string {
  return JSON.stringify(value);
}

/** Message overhead plus a stable serialization of its provider-visible value. */
export function estimateMessages(messages: readonly Message[]): number {
  return messages.reduce((total, message) => total + 8 + estimateTokens(stable(message)), 0);
}

export function latestSourceRefs(entries: readonly BtwEntry[]): BtwSummarySourceRef[] {
  return entries.map((entry) => ({ entryId: entry.id, attemptId: entry.attempts.at(-1)!.id }));
}

export function canonicalTuple(entry: BtwEntry): [string, string, string, string, string | null] {
  const attempt = entry.attempts.at(-1)!;
  return [entry.id, attempt.id, entry.question, attempt.answer, attempt.error ?? null];
}

export function sourceHash(entries: readonly BtwEntry[]): string {
  return createHash("sha256").update(stable(entries.map(canonicalTuple))).digest("hex");
}

export function canonicalEntryBlock(entry: BtwEntry): string {
  const attempt = entry.attempts.at(-1)!;
  const answer = attempt.answer || "(no answer)";
  return `Q: ${entry.question}\nA: ${answer}${attempt.error === undefined ? "" : `\nError: ${attempt.error}`}`;
}

/** A persisted summary is valid only for an exact latest-attempt entry prefix. */
export function validateSummary(thread: Pick<BtwThread, "entries" | "summary">): boolean {
  const summary = thread.summary;
  if (!summary || !summary.text || !summary.source.length) return false;
  const prefix = thread.entries.slice(0, summary.source.length);
  if (prefix.length !== summary.source.length || prefix.at(-1)?.id !== summary.throughEntryId) return false;
  if (!summary.source.every((ref, index) => ref.entryId === prefix[index]!.id && ref.attemptId === prefix[index]!.attempts.at(-1)!.id)) return false;
  return sourceHash(prefix) === summary.sourceHash;
}

export function summaryCoveredCount(thread: Pick<BtwThread, "entries" | "summary">): number {
  return validateSummary(thread) ? thread.summary!.source.length : 0;
}

export function contextBudget(
  triggerTokens: number,
  modelContextWindow: number | undefined,
  answerMaxTokens: number,
  systemPrompt: string | undefined,
  prefix: readonly Message[],
  question: string,
): number {
  const available = (modelContextWindow ?? Number.MAX_SAFE_INTEGER) - answerMaxTokens - estimateTokens(systemPrompt ?? "") - estimateMessages(prefix) - estimateTokens(question) - 1024;
  return Math.min(triggerTokens, Math.max(0, available));
}

function truncateCodePoints(text: string, budget: number): string {
  if (budget <= 0) return "";
  if (estimateTokens(text) <= budget) return text;
  const marker = "…[truncated]";
  if (estimateTokens(marker) > budget) {
    let output = "";
    for (const point of text) {
      if (estimateTokens(output + point) > budget) break;
      output += point;
    }
    return output;
  }
  let output = "";
  for (const point of text) {
    if (estimateTokens(output + point + marker) > budget) break;
    output += point;
  }
  return output + marker;
}

/** Truncate a newest oversized raw exchange without splitting Unicode code points. */
export function truncateEntryBlock(entry: BtwEntry, budget: number): string {
  const attempt = entry.attempts.at(-1)!;
  const question = `Q: ${entry.question}\nA: `;
  const answer = attempt.answer || "(no answer)";
  const error = attempt.error === undefined ? "" : `\nError: ${attempt.error}`;
  if (estimateTokens(question) >= budget) return truncateCodePoints(question, budget);
  return truncateCodePoints(question + answer + error, budget);
}

function chooseRaw(entries: readonly BtwEntry[], budget: number): string[] {
  if (budget <= 0 || !entries.length) return [];
  const selected: string[] = [];
  let used = 0;
  for (let index = entries.length - 1; index >= 0; index--) {
    const block = canonicalEntryBlock(entries[index]!);
    const extra = selected.length ? estimateTokens("\n\n") : 0;
    if (used + extra + estimateTokens(block) <= budget) {
      selected.unshift(block);
      used += extra + estimateTokens(block);
      continue;
    }
    if (!selected.length) {
      const clipped = truncateEntryBlock(entries[index]!, budget);
      if (clipped) selected.unshift(clipped);
    }
    break;
  }
  return selected;
}

const SUMMARY_HEADER = "Earlier side Q/A summary (model-generated reference; may be incomplete):\n";
const RAW_HEADER = "Recent raw side Q/A (authoritative; takes precedence):\n";

/** Assemble bounded summary+recent context; recent raw blocks always receive the budget first. */
export function assembleBoundedContext(thread: BtwThread, budget: number, excludeEntryId?: string): string {
  if (budget <= 0) return "";
  const valid = validateSummary(thread);
  const summary = valid && !thread.summary!.source.some((ref) => ref.entryId === excludeEntryId) ? thread.summary : undefined;
  const covered = summary ? summary.source.length : 0;
  const entries = thread.entries.slice(covered).filter((entry) => entry.id !== excludeEntryId);
  const rawAllowance = Math.max(0, budget - estimateTokens(RAW_HEADER));
  const rawBlocks = chooseRaw(entries, rawAllowance);
  let raw = rawBlocks.length ? RAW_HEADER + rawBlocks.join("\n\n") : "";
  if (estimateTokens(raw) > budget) raw = "";
  const separator = raw ? estimateTokens("\n\n") : 0;
  const summaryAllowance = budget - estimateTokens(raw) - separator - estimateTokens(SUMMARY_HEADER);
  const summaryText = summary && summaryAllowance > 0
    ? SUMMARY_HEADER + truncateCodePoints(summary.text, summaryAllowance)
    : "";
  // Raw context is selected first and must remain byte-for-byte intact. The
  // separator is part of the summary allocation, so this composition is bounded
  // without a final combined truncation that could clip the authoritative tail.
  const result = summaryText && raw ? `${summaryText}\n\n${raw}` : summaryText || raw;
  return estimateTokens(result) <= budget ? result : raw;
}

function summaryInput(previous: BtwSummary | undefined, blocks: readonly string[]): string {
  const prior = previous ? `Previous side Q/A summary (model-generated reference; may be incomplete):\n${previous.text}` : "";
  const raw = blocks.length ? `Raw side Q/A to incorporate (authoritative):\n${blocks.join("\n\n")}` : "";
  return prior && raw ? `${prior}\n\n${raw}` : prior || raw;
}

/**
 * Select an immutable, bounded summary job. It never includes main-session context or tools.
 * A candidate extends only a valid existing prefix and leaves a newest raw suffix unsummarized.
 */
export function selectSummarySnapshot(thread: BtwThread, policy: SummaryPolicy): SummarySnapshot | null {
  const existing = validateSummary(thread) ? thread.summary : undefined;
  const start = existing?.source.length ?? 0;
  const unsummarized = thread.entries.slice(start);
  const total = estimateTokens(existing?.text ?? "") + unsummarized.reduce((sum, entry) => sum + estimateTokens(canonicalEntryBlock(entry)), 0);
  if (total <= policy.triggerTokens || !unsummarized.length || policy.inputMaxTokens <= 0) return null;

  let retained = 0;
  let retainedCount = 0;
  for (let index = unsummarized.length - 1; index >= 0; index--) {
    const size = estimateTokens(canonicalEntryBlock(unsummarized[index]!));
    const separator = retainedCount > 0 ? estimateTokens("\n\n") : 0;
    if (retainedCount > 0 && retained + separator + size > policy.retainTokens) break;
    retained += separator + size;
    retainedCount++;
  }
  const coverable = unsummarized.slice(0, Math.max(0, unsummarized.length - retainedCount));
  if (!coverable.length) return null;

  const blocks: string[] = [];
  for (const entry of coverable) {
    const next = [...blocks, canonicalEntryBlock(entry)];
    if (estimateTokens(summaryInput(existing, next)) > policy.inputMaxTokens) break;
    blocks.push(canonicalEntryBlock(entry));
  }
  if (!blocks.length) return null;
  const coveredEntries = thread.entries.slice(0, start + blocks.length);
  const source = latestSourceRefs(coveredEntries);
  return {
    source,
    sourceHash: sourceHash(coveredEntries),
    throughEntryId: source.at(-1)!.entryId,
    input: summaryInput(existing, blocks),
  };
}