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

const PACKAGE_VERSION = "18.0.8";

interface WorkerInput {
  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;
  };
  includeThinking: boolean;
  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;
  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 (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 input = parseInput(JSON.parse(await Bun.stdin.text()));
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 = [
  ...input.preparation.messagesToSummarize,
  ...input.preparation.turnPrefixMessages,
] as never[];
const serialized = serializeConversation(messages, { includeThinking: input.includeThinking });
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,
    fileOps,
  },
  {
    shape,
    maxFrames: frameBudget,
    includeThinking: input.includeThinking,
  },
);
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({
  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,
}));