export type JsonValue = null | boolean | number | string | JsonValue[] | { [key: string]: JsonValue };

const encoder = new TextEncoder();

export function estimateTokens(text: string): number {
  return Math.ceil(encoder.encode(text).byteLength / 4);
}

function prefixWithinBytes(points: string[], byteLimit: number): string {
  let bytes = 0;
  let end = 0;
  while (end < points.length) {
    const size = encoder.encode(points[end]).byteLength;
    if (bytes + size > byteLimit) break;
    bytes += size;
    end += 1;
  }
  return points.slice(0, end).join("");
}

function suffixWithinBytes(points: string[], byteLimit: number): string {
  let bytes = 0;
  let start = points.length;
  while (start > 0) {
    const size = encoder.encode(points[start - 1]).byteLength;
    if (bytes + size > byteLimit) break;
    bytes += size;
    start -= 1;
  }
  return points.slice(start).join("");
}

/** Bounded middle truncation that never splits a Unicode code point. */
export function truncateMiddle(text: string, maxApproxTokens: number): { text: string; truncated: boolean } {
  if (estimateTokens(text) <= maxApproxTokens) return { text, truncated: false };

  const points = Array.from(text);
  const maxBytes = maxApproxTokens * 4;
  const markerReserve = Math.min(120, Math.max(40, Math.floor(maxBytes / 3)));
  const contentBytes = Math.max(0, maxBytes - markerReserve);
  const prefix = prefixWithinBytes(points, Math.ceil(contentBytes / 2));
  const suffix = suffixWithinBytes(points.slice(Array.from(prefix).length), Math.floor(contentBytes / 2));
  const omittedBytes = Math.max(0, encoder.encode(text).byteLength - encoder.encode(prefix).byteLength - encoder.encode(suffix).byteLength);
  const marker = `<truncated omitted_approx_tokens="${Math.ceil(omittedBytes / 4)}" />`;
  const result = `${prefix}${marker}${suffix}`;

  // Tiny test budgets can be smaller than the marker. Production budgets are large.
  if (estimateTokens(result) <= maxApproxTokens) return { text: result, truncated: true };
  return { text: prefixWithinBytes(Array.from(marker), maxBytes), truncated: true };
}

export function toBoundedJson(value: unknown, stringTokenCap = 16_000): { value: JsonValue; truncated: boolean } {
  let truncated = false;

  const visit = (candidate: unknown, path: string): JsonValue => {
    if (candidate === null || typeof candidate === "boolean") return candidate;
    if (typeof candidate === "number") {
      if (!Number.isFinite(candidate)) throw new Error(`Non-finite number in planned action at ${path}`);
      return candidate;
    }
    if (typeof candidate === "string") {
      const bounded = truncateMiddle(candidate, stringTokenCap);
      truncated ||= bounded.truncated;
      return bounded.text;
    }
    if (Array.isArray(candidate)) return candidate.map((item, index) => visit(item, `${path}[${index}]`));
    if (typeof candidate === "object") {
      const object = candidate as Record<string, unknown>;
      const result: { [key: string]: JsonValue } = Object.create(null);
      for (const key of Object.keys(object).sort()) {
        const item = object[key];
        if (item === undefined) continue;
        result[key] = visit(item, `${path}.${key}`);
      }
      return result;
    }
    throw new Error(`Non-JSON value in planned action at ${path}`);
  };

  return { value: visit(value, "$"), truncated };
}

export function stableStringify(value: JsonValue): string {
  return JSON.stringify(value, null, 2);
}