/**
 * StreamingToolCallBuffer — Accumulates raw tool call argument strings
 * during streaming, providing a fallback when LangChain's parsePartialJson
 * loses data (e.g., when max_tokens truncates a tool call mid-JSON).
 *
 * Architecture:
 * - handleToolCallChunks feeds raw arg fragments into this buffer
 * - When a tool call is finalized (or truncated), the ToolNode can access
 *   the accumulated raw string and extract field values that parsePartialJson
 *   may have dropped
 * - This enables "streaming write" semantics: content that streamed to the UI
 *   is never lost, even if the JSON wrapper was incomplete
 *
 * @example
 * ```ts
 * const buffer = new StreamingToolCallBuffer();
 *
 * // During streaming (called by handleToolCallChunks):
 * buffer.append(toolCallId, '{"action":"write","content":"func');
 * buffer.append(toolCallId, 'tion App() {');
 * buffer.setToolName(toolCallId, 'content_tool');
 *
 * // When ToolNode needs the content:
 * const content = buffer.extractFieldValue(toolCallId, 'content');
 * // => 'function App() {'
 * ```
 */

export class StreamingToolCallBuffer {
  /** Raw accumulated arg strings keyed by tool call ID */
  private argBuffers: Map<string, string> = new Map();
  /** Tool names keyed by tool call ID */
  private toolNames: Map<string, string> = new Map();
  /**
   * Maps chunk index → tool call ID for providers (e.g., Bedrock) that send
   * START chunks with {id, name, index} and DELTA chunks with {args, index} (no id).
   * Without this mapping, DELTA chunks cannot be associated with the correct buffer entry.
   */
  private indexToIdMap: Map<number, string> = new Map();
  /**
   * Tracks the last appended args fragment per tool call ID.
   * Used to skip exact duplicate fragments caused by LangChain's `streamEvents`
   * emitting the same `on_chat_model_stream` event multiple times when the model
   * is wrapped in nested runnables (e.g., `.bind()`, `.withConfig()`).
   */
  private lastChunks: Map<string, string> = new Map();

  /**
   * Append a raw argument string chunk for a tool call.
   * Called during streaming as each ToolCallChunk arrives.
   *
   * Skips exact duplicate fragments that arrive when LangChain emits the same
   * streaming event multiple times from nested runnable wrappers. Without this,
   * each token's args appear twice, producing malformed JSON like:
   * `{"action": "write"action": "write"...`
   *
   * @returns true if the chunk was appended, false if it was a duplicate
   */
  append(toolCallId: string, argsChunk: string): boolean {
    if (!argsChunk) return false;

    // Skip exact duplicate of the last appended fragment
    const lastChunk = this.lastChunks.get(toolCallId);
    if (lastChunk === argsChunk) {
      return false;
    }
    this.lastChunks.set(toolCallId, argsChunk);

    const existing = this.argBuffers.get(toolCallId) ?? '';
    this.argBuffers.set(toolCallId, existing + argsChunk);
    return true;
  }

  /**
   * Record the tool name for a tool call ID.
   * Usually available from the first chunk.
   */
  setToolName(toolCallId: string, name: string): void {
    if (name && !this.toolNames.has(toolCallId)) {
      this.toolNames.set(toolCallId, name);
    }
  }

  /**
   * Get the tool name for a tool call ID.
   */
  getToolName(toolCallId: string): string | undefined {
    return this.toolNames.get(toolCallId);
  }

  /**
   * Get the raw accumulated argument string for a tool call.
   */
  getRawArgs(toolCallId: string): string | undefined {
    return this.argBuffers.get(toolCallId);
  }

  /**
   * Extract a string field value from the raw accumulated JSON args.
   *
   * Uses simple string parsing (not JSON.parse) to handle truncated JSON.
   * Finds `"fieldName":"` and extracts everything after it until the next
   * unescaped quote or end of buffer.
   *
   * @param toolCallId - The tool call ID
   * @param fieldName - The JSON field name to extract (e.g., 'content')
   * @returns The extracted value, or undefined if the field wasn't found
   */
  extractFieldValue(toolCallId: string, fieldName: string): string | undefined {
    const raw = this.argBuffers.get(toolCallId);
    if (!raw) return undefined;

    // Find the field key in the JSON string: "fieldName":"
    // Handle both `"fieldName":"value"` and `"fieldName": "value"` (with space)
    const patterns = [
      `"${fieldName}":"`,
      `"${fieldName}": "`,
      `"${fieldName}" : "`,
    ];

    let startIdx = -1;
    for (const pattern of patterns) {
      const idx = raw.indexOf(pattern);
      if (idx !== -1) {
        startIdx = idx + pattern.length;
        break;
      }
    }

    if (startIdx === -1) return undefined;

    // Extract the value — handle escaped characters
    let value = '';
    let escaped = false;
    for (let i = startIdx; i < raw.length; i++) {
      const ch = raw[i];
      if (escaped) {
        // Handle common escape sequences
        switch (ch) {
          case 'n':
            value += '\n';
            break;
          case 't':
            value += '\t';
            break;
          case 'r':
            value += '\r';
            break;
          case '"':
            value += '"';
            break;
          case '\\':
            value += '\\';
            break;
          default:
            value += ch;
            break;
        }
        escaped = false;
        continue;
      }
      if (ch === '\\') {
        escaped = true;
        continue;
      }
      if (ch === '"') {
        // End of the string value
        break;
      }
      value += ch;
    }

    return value || undefined;
  }

  /**
   * Store an index → tool call ID mapping.
   * Called when a START chunk arrives with both `id` and `index` (e.g., Bedrock).
   * Subsequent DELTA chunks that only have `index` can resolve the ID via getIdByIndex.
   */
  setIndexMapping(index: number, toolCallId: string): void {
    this.indexToIdMap.set(index, toolCallId);
  }

  /**
   * Resolve a tool call ID from a chunk index.
   * Returns undefined if no mapping exists (e.g., provider sends id on every chunk).
   */
  getIdByIndex(index: number): string | undefined {
    return this.indexToIdMap.get(index);
  }

  /**
   * Check if a tool call has accumulated content.
   */
  has(toolCallId: string): boolean {
    return this.argBuffers.has(toolCallId);
  }

  /**
   * Clear the buffer for a specific tool call (after processing).
   */
  clear(toolCallId: string): void {
    this.argBuffers.delete(toolCallId);
    this.toolNames.delete(toolCallId);
    this.lastChunks.delete(toolCallId);
  }

  /**
   * Clear all buffers (at the start of a new callModel cycle).
   */
  clearAll(): void {
    this.argBuffers.clear();
    this.toolNames.clear();
    this.indexToIdMap.clear();
    this.lastChunks.clear();
  }
}