import {
  compact,
  createFileOps,
  geometry,
  getPreservedArchive,
  maxFramesForDataBudget,
  providerFrameBudget,
  renderabilityProbeText,
  resolveShapeForText,
  serializeConversation,
} from "@oh-my-pi/snapcompact";

const WORKER_PROTOCOL_VERSION = 2;
const PACKAGE_VERSION = "18.0.8";

interface WorkerInput {
  protocolVersion: number;
  preparation: {
    firstKeptEntryId: string;
    messagesToSummarize: unknown[];
    turnPrefixMessages: unknown[];
    tokensBefore: number;
    previousSummary?: string;
    previousPreserveData?: Record<string, unknown>;
    fileOps: {
      read: string[];
      written: string[];
      edited: string[];
    };
  };
  model: {
    api?: string;
    id: string;
    provider: string;
  };
  historicalToolResultsToKeep: number;
  maxFrameDataBytes: number;
  maxArchiveTokens: number;
}

function assertStringArray(value: unknown, name: string): asserts value is string[] {
  if (!Array.isArray(value) || value.some((item) => typeof item !== "string")) {
    throw new Error(`${name} must be a string array`);
  }
}

function parseInput(value: unknown): WorkerInput {
  if (!value || typeof value !== "object") throw new Error("worker input must be an object");
  const input = value as WorkerInput;
  if (input.protocolVersion !== WORKER_PROTOCOL_VERSION) {
    throw new Error("worker protocol version mismatch");
  }
  const preparation = input.preparation;
  if (!preparation || typeof preparation !== "object") throw new Error("missing preparation");
  if (typeof preparation.firstKeptEntryId !== "string") throw new Error("missing firstKeptEntryId");
  if (!Array.isArray(preparation.messagesToSummarize)) throw new Error("messagesToSummarize must be an array");
  if (!Array.isArray(preparation.turnPrefixMessages)) throw new Error("turnPrefixMessages must be an array");
  if (typeof preparation.tokensBefore !== "number") throw new Error("tokensBefore must be a number");
  if (!preparation.fileOps || typeof preparation.fileOps !== "object") throw new Error("missing fileOps");
  assertStringArray(preparation.fileOps.read, "fileOps.read");
  assertStringArray(preparation.fileOps.written, "fileOps.written");
  assertStringArray(preparation.fileOps.edited, "fileOps.edited");
  if (!input.model || typeof input.model.id !== "string" || typeof input.model.provider !== "string") {
    throw new Error("missing model identity");
  }
  if (
    !Number.isInteger(input.historicalToolResultsToKeep) ||
    input.historicalToolResultsToKeep < 0 ||
    input.historicalToolResultsToKeep > 3
  ) {
    throw new Error("historicalToolResultsToKeep must be between 0 and 3");
  }
  if (typeof input.maxFrameDataBytes !== "number" || input.maxFrameDataBytes < 170_000) {
    throw new Error("maxFrameDataBytes is too small");
  }
  if (typeof input.maxArchiveTokens !== "number" || input.maxArchiveTokens < 1) {
    throw new Error("maxArchiveTokens must be positive");
  }
  return input;
}

function conservativeTextTokens(text: string): number {
  let nonAscii = 0;
  for (let index = 0; index < text.length; index++) {
    if (text.charCodeAt(index) > 0x7f) nonAscii++;
  }
  return Math.ceil((text.length - nonAscii) / 3 + nonAscii * 2);
}

const TOOL_RESULT_PLACEHOLDER = "[tool result]";
const NEWLINE_GLYPH = "█";

function maskArchivedToolResults(text: string, keepRecent: number): string {
  // SnapCompact strips dim markers from raw bodies, so only its real closing delimiter can match.
  const outputBlock = /<out>(?:\n|█)\u000e[\s\S]*?\u000f(?:\n|█)<\/out>/g;
  const redact = Math.max(0, [...text.matchAll(outputBlock)].length - keepRecent);
  let index = 0;
  return text.replace(outputBlock, (block) => {
    const metadata = /\[tool result [^\]]+\]/.exec(block)?.[0];
    const replacement =
      index < redact
        ? `<out>${NEWLINE_GLYPH}${metadata ? `${metadata}${NEWLINE_GLYPH}` : ""}${TOOL_RESULT_PLACEHOLDER}${NEWLINE_GLYPH}</out>`
        : block;
    index++;
    return replacement;
  });
}

function sanitizePreviousArchive(
  preserveData: Record<string, unknown> | undefined,
  keepRecent: number,
): void {
  if (!preserveData) return;
  const archive = getPreservedArchive(preserveData);
  if (!archive) return;
  const source =
    archive.text ??
    [archive.textHead, archive.textTail]
      .filter((part): part is string => typeof part === "string" && part.length > 0)
      .join(NEWLINE_GLYPH);
  if (!source) return;
  archive.text = maskArchivedToolResults(source, keepRecent);
  archive.textHead = undefined;
  archive.textTail = undefined;
  preserveData.snapcompact = archive;
}

function annotateToolResults(messages: unknown[]): unknown[] {
  return messages.map((message) => {
    if (!message || typeof message !== "object") return message;
    const record = message as Record<string, unknown>;
    if (record.role !== "toolResult" || !Array.isArray(record.content)) return message;
    const hasText = record.content.some(
      (block) =>
        block !== null &&
        typeof block === "object" &&
        "type" in block &&
        block.type === "text" &&
        "text" in block &&
        typeof block.text === "string" &&
        block.text.length > 0,
    );
    if (!hasText && record.isError !== true) return message;
    const toolName = typeof record.toolName === "string" ? record.toolName : "unknown";
    const toolCallId = typeof record.toolCallId === "string" ? record.toolCallId : "unknown";
    return {
      ...record,
      content: [
        {
          type: "text",
          text: `[tool result ${toolName}#${toolCallId} error=${record.isError === true}]\n`,
        },
        ...record.content,
      ],
    };
  });
}

const input = parseInput(JSON.parse(await Bun.stdin.text()));
sanitizePreviousArchive(
  input.preparation.previousPreserveData,
  input.historicalToolResultsToKeep,
);
const fileOps = createFileOps();
for (const file of input.preparation.fileOps.read) fileOps.read.add(file);
for (const file of input.preparation.fileOps.written) fileOps.written.add(file);
for (const file of input.preparation.fileOps.edited) fileOps.edited.add(file);

const messages = annotateToolResults([
  ...input.preparation.messagesToSummarize,
  ...input.preparation.turnPrefixMessages,
]) as never[];
const serialized = serializeConversation(messages, {
  includeThinking: false,
  toolResultMaxChars: Number.POSITIVE_INFINITY,
  toolArgMaxChars: Number.POSITIVE_INFINITY,
  toolCallMaxChars: Number.POSITIVE_INFINITY,
});
const probe = renderabilityProbeText(
  serialized,
  input.preparation.previousPreserveData,
  input.preparation.previousSummary,
);
const model = { api: input.model.api as never, id: input.model.id };
const shape = resolveShapeForText(probe, model);
const textEdgeTokens = 2 * geometry(shape).capacity;
const imageTokenBudget = Math.max(0, input.maxArchiveTokens - textEdgeTokens);
const frameBudget = Math.min(
  providerFrameBudget(input.model.provider),
  maxFramesForDataBudget(input.maxFrameDataBytes),
  Math.floor(imageTokenBudget / shape.frameTokenEstimate),
);
if (frameBudget < 1) {
  throw new Error("no safe visual-token budget remains for a SnapCompact frame");
}

const result = await compact(
  {
    ...input.preparation,
    messagesToSummarize: messages,
    turnPrefixMessages: [],
    fileOps,
  },
  {
    shape,
    maxFrames: frameBudget,
    includeThinking: false,
    toolResultMaxChars: Number.POSITIVE_INFINITY,
    toolArgMaxChars: Number.POSITIVE_INFINITY,
    toolCallMaxChars: Number.POSITIVE_INFINITY,
  },
);
const preserveData = result.preserveData;
if (!preserveData) throw new Error("SnapCompact did not return preserve data");
const archive = getPreservedArchive(preserveData);
if (!archive) throw new Error("SnapCompact did not return an archive");
let frameDataBytes = archive.frames.reduce((sum, frame) => sum + frame.data.length, 0);
let dataBudgetOmittedFrames = 0;
while (frameDataBytes > input.maxFrameDataBytes && archive.frames.length > 0) {
  const removed = archive.frames.shift();
  if (!removed) break;
  frameDataBytes -= removed.data.length;
  dataBudgetOmittedFrames++;
}
if (frameDataBytes > input.maxFrameDataBytes) {
  throw new Error("rendered frame data cannot fit the configured byte budget");
}
if (dataBudgetOmittedFrames > 0) {
  (archive as typeof archive & { dataBudgetOmittedFrames: number }).dataBudgetOmittedFrames =
    dataBudgetOmittedFrames;
}
preserveData.snapcompact = archive;
const textChars = (archive.textHead?.length ?? 0) + (archive.textTail?.length ?? 0);
const archiveTokenEstimate =
  conservativeTextTokens(
    result.summary + (archive.textHead ?? "") + (archive.textTail ?? ""),
  ) + archive.frames.length * shape.frameTokenEstimate;
if (archiveTokenEstimate > input.maxArchiveTokens) {
  throw new Error(
    `rendered archive estimate ${archiveTokenEstimate} exceeds budget ${input.maxArchiveTokens}`,
  );
}

process.stdout.write(JSON.stringify({
  protocolVersion: WORKER_PROTOCOL_VERSION,
  packageVersion: PACKAGE_VERSION,
  frameBudget,
  frameTokenEstimate: shape.frameTokenEstimate,
  archiveTokenEstimate,
  summary: result.summary,
  shortSummary: result.shortSummary,
  firstKeptEntryId: result.firstKeptEntryId,
  tokensBefore: result.tokensBefore,
  readFiles: result.details?.readFiles ?? [],
  modifiedFiles: result.details?.modifiedFiles ?? [],
  preserveData,
}));