#!/usr/bin/env node
/**
 * OpenClaw Compression — Byte-Range Surgery
 *
 * Cache-safe byte-level replacement for transparent proxy compression.
 * Migrated from Headroom `crates/headroom-core/src/transforms/live_zone.rs`
 *
 * Core idea:
 *   When compressing tool outputs in a proxy, only replace the bytes of
 *   the tool_result block. All other bytes (system prompt, conversation
 *   prefix, tool_use blocks) are copied VERBATIM from the original.
 *   This guarantees the SHA-256 of the prefix is byte-identical to the
 *   original, so the provider's prompt cache still hits.
 *
 *   Bytes outside rewritten ranges are *literally copied* from the input,
 *   never re-serialized.
 *
 * Usage:
 *   const { ByteSurgeon, computeLiveZone } = require('./byte-surgery');
 *
 *   const requestBody = '...'; // raw HTTP body bytes
 *   const replacements = surgeon.planSurgery(requestBody, blocks);
 *   const newBody = surgeon.operate(requestBody, replacements);
 */

const crypto = require('crypto');

// ─── Types ───────────────────────────────────────────────────────

/**
 * Byte-range replacement.
 * ECC/Headroom pattern: (start_byte, end_byte) → replacement_bytes
 */
class ByteRange {
  constructor(start, end, replacement) {
    this.start = start;          // inclusive
    this.end = end;              // exclusive
    this.replacement = replacement; // Buffer or string
  }

  get length() { return this.end - this.start; }
}

/**
 * A block in the request body — corresponds to a message content block.
 */
class ContentBlock {
  constructor(type, start, end, content) {
    this.type = type;            // 'system' | 'user_text' | 'tool_result' | 'tool_use' | 'assistant'
    this.start = start;          // byte offset in body
    this.end = end;              // byte offset
    this.content = content;      // decoded string content
  }
}

// ─── Live Zone Detection ─────────────────────────────────────────

/**
 * Identify the "live zone" — the portion of the request that can be
 * safely modified without breaking provider prompt cache.
 *
 * Headroom logic:
 *   - Frozen zone: from byte 0 to the last tool_use/assistant block that
 *     is before the latest user message. NEVER modified.
 *   - Hot zone: tool_use, thinking blocks. NEVER modified (they affect caching).
 *   - Live zone: tool_result blocks in the latest user message. ONLY these
 *     are compressed.
 *
 * @param {Buffer} body — full request body
 * @param {string} format — 'anthropic' | 'openai' | 'generic'
 * @returns {object} { blocks, frozenEnd, liveBlocks }
 */
function computeLiveZone(body, format = 'generic') {
  const blocks = parseBlocks(body, format);
  if (blocks.length === 0) return { blocks, frozenEnd: 0, liveBlocks: [] };

  // Find the latest user message
  // All tool_result blocks are live; frozen = before first tool_result

  // Frozen zone: everything BEFORE the first tool_result block
  // Live blocks: ALL tool_result blocks
  const liveBlocks = [];
  let frozenEnd = 0;

  for (const b of blocks) {
    if (b.type === 'tool_result') {
      liveBlocks.push(b);
      if (frozenEnd === 0) frozenEnd = b.start;
    }
  }

  if (frozenEnd === 0) frozenEnd = body.length;

  return { blocks, frozenEnd, liveBlocks };
}

/**
 * Parse HTTP request body into content blocks.
 * Handles Anthropic Messages API and OpenAI Chat Completions API formats.
 */
function parseBlocks(body, format = 'generic') {
  const blocks = [];
  let text;

  try {
    text = Buffer.isBuffer(body) ? body.toString('utf8') : String(body);
  } catch {
    return blocks;
  }

  if (format === 'anthropic' || format === 'generic') {
    // Anthropic: body.messages[].content[].type = "tool_result"
    try {
      const parsed = JSON.parse(text);
      const messages = parsed.messages || [];
      let offset = 0;

      for (const msg of messages) {
        const content = msg.content;
        if (typeof content === 'string') {
          const start = text.indexOf(JSON.stringify(content), offset);
          if (start < 0) continue;
          const end = start + JSON.stringify(content).length;
          blocks.push(new ContentBlock(
            msg.role === 'user' ? 'user_text' : msg.role,
            start, end, content
          ));
          offset = end;
        } else if (Array.isArray(content)) {
          for (const block of content) {
            const blockJson = JSON.stringify(block);
            const start = text.indexOf(blockJson, offset);
            if (start < 0) continue;
            const end = start + blockJson.length;
            blocks.push(new ContentBlock(block.type || 'unknown', start, end, blockJson));
            offset = end;
          }
        }
      }
    } catch {
      // Fallback: find tool_result blocks by byte position in text
      const re = /"type"\s*:\s*"tool_result"/g;
      while ((match = re.exec(text)) !== null) {
        let blockStart = match.index;
        // Walk back to start of this JSON object
        let depth = 0, startFound = false;
        for (let i = blockStart; i >= 0; i--) {
          if (text[i] === '}') depth++;
          if (text[i] === '{') {
            if (depth === 0) { blockStart = i; startFound = true; break; }
            depth--;
          }
        }
        if (!startFound) continue;
        // Walk forward to end of this object
        depth = 0;
        let blockEnd = blockStart;
        for (let i = blockStart; i < text.length; i++) {
          if (text[i] === '{') depth++;
          if (text[i] === '}') {
            depth--;
            if (depth === 0) { blockEnd = i + 1; break; }
          }
        }
        const content = text.substring(blockStart, blockEnd);
        blocks.push(new ContentBlock('tool_result', blockStart, blockEnd, content));
      }
    }
  }

  if (format === 'openai') {
    // OpenAI: body.messages[].role = "tool", content = "<output>"
    // Similar parsing...
  }

  return blocks;
}

// ─── Byte Surgeon ────────────────────────────────────────────────

class ByteSurgeon {
  /**
   * Plan surgery: compute byte-range replacements for compressed blocks.
   *
   * @param {Buffer|string} body — Original request body
   * @param {Array} compressedBlocks — [{ blockIndex, compressedContent }]
   * @param {string} format — Provider format
   * @returns {ByteRange[]} Sorted replacement ranges
   */
  planSurgery(body, compressedBlocks, format = 'generic') {
    const { blocks, liveBlocks, frozenEnd } = computeLiveZone(body, format);
    const replacements = [];

    for (const compressed of compressedBlocks) {
      if (compressed.blockIndex < liveBlocks.length) {
        const block = liveBlocks[compressed.blockIndex];
        const newContent = JSON.stringify({ type: 'tool_result', content: compressed.compressedContent });

        replacements.push(new ByteRange(block.start, block.end, Buffer.from(newContent)));
      }
    }

    // Sort by start offset
    replacements.sort((a, b) => a.start - b.start);

    // Validate no overlapping ranges
    for (let i = 1; i < replacements.length; i++) {
      if (replacements[i].start < replacements[i - 1].end) {
        throw new Error(`Overlapping byte ranges: [${replacements[i-1].start},${replacements[i-1].end}) and [${replacements[i].start},${replacements[i].end})`);
      }
    }

    return { replacements, frozenEnd, liveBlockCount: liveBlocks.length };
  }

  /**
   * Execute surgery: splice replacements into the original body.
   *
   * CRITICAL: bytes OUTSIDE replaced ranges are COPIED VERBATIM.
   * Never re-serialize, never re-encode. This is what keeps the
   * provider prompt cache prefix SHA-256 byte-identical.
   *
   * Headroom: live_zone.rs:apply_replacements()
   */
  operate(body, replacements) {
    const original = Buffer.isBuffer(body) ? body : Buffer.from(body);
    const parts = [];
    let cursor = 0;

    for (const r of replacements) {
      // Copy untouched bytes before this replacement
      if (cursor < r.start) {
        parts.push(original.subarray(cursor, r.start));
      }

      // Insert replacement
      if (Buffer.isBuffer(r.replacement)) {
        parts.push(r.replacement);
      } else {
        parts.push(Buffer.from(r.replacement));
      }

      cursor = r.end;
    }

    // Copy remaining bytes after last replacement
    if (cursor < original.length) {
      parts.push(original.subarray(cursor));
    }

    return Buffer.concat(parts);
  }

  /**
   * Validate that surgery didn't corrupt the structure.
   * Checks: valid JSON, same message count, frozen prefix byte-identical.
   */
  validate(original, result, format = 'generic') {
    const checks = { valid: true, warnings: [] };

    // 1. Valid JSON
    try {
      JSON.parse(result.toString('utf8'));
    } catch (e) {
      checks.valid = false;
      checks.warnings.push(`Invalid JSON after surgery: ${e.message}`);
    }

    // 2. Frozen prefix unchanged
    const { frozenEnd } = computeLiveZone(original, format);
    if (frozenEnd > 0) {
      const originalPrefix = Buffer.isBuffer(original)
        ? original.subarray(0, frozenEnd)
        : Buffer.from(original).subarray(0, frozenEnd);
      const resultPrefix = result.subarray(0, Math.min(frozenEnd, result.length));

      if (!originalPrefix.equals(resultPrefix)) {
        checks.warnings.push('Frozen prefix changed — provider cache will MISS');
      }
    }

    // 3. Size check
    checks.originalSize = original.length;
    checks.resultSize = result.length;
    checks.savingsPercent = Math.round((1 - result.length / original.length) * 10000) / 100;

    return checks;
  }
}

// ─── SHA-256 Prefix Verification ─────────────────────────────────

/**
 * Verify that the SHA-256 of the frozen prefix is unchanged.
 * This is the definitive check for provider prompt cache safety.
 */
function verifyCachePrefix(original, result, frozenEnd) {
  const origPrefix = Buffer.isBuffer(original)
    ? original.subarray(0, frozenEnd)
    : Buffer.from(original).subarray(0, frozenEnd);
  const resPrefix = result.subarray(0, Math.min(frozenEnd, result.length));

  const origHash = crypto.createHash('sha256').update(origPrefix).digest('hex');
  const resHash = crypto.createHash('sha256').update(resPrefix).digest('hex');

  return {
    match: origHash === resHash,
    originalHash: origHash,
    resultHash: resHash,
    prefixBytes: frozenEnd,
  };
}

// ─── Module exports ──────────────────────────────────────────────

module.exports = {
  ByteSurgeon,
  ByteRange,
  ContentBlock,
  computeLiveZone,
  parseBlocks,
  verifyCachePrefix,
};
