// src/utils/toolDiscoveryCache.ts
import type { BaseMessage } from '@langchain/core/messages';
import { Constants, MessageTypes } from '@/common';
import { TOOL_DISCOVERY_CACHE_MAX_SIZE } from '@/common/constants';

/**
 * Cached tool discovery entry.
 * Stores the tool name and the message index where it was discovered,
 * enabling efficient lookups without re-parsing conversation history.
 */
export interface ToolDiscoveryEntry {
  /** The tool name that was discovered */
  toolName: string;
  /** Message index in conversation history where discovery occurred */
  discoveredAtIndex: number;
}

/**
 * ToolDiscoveryCache provides a run-scoped cache of tool search results.
 *
 * Problem: Without caching, every LLM iteration re-parses the full message
 * history via extractToolDiscoveries() to find tool_search results. In long
 * conversations with many tool iterations, this is redundant work.
 *
 * Solution: Cache discovered tool names by message index. On each iteration,
 * only scan messages AFTER the last scanned index. Already-seen discoveries
 * are returned from cache instantly.
 *
 * This mirrors the pattern used by VS Code Copilot Chat where tool search
 * results from prior turns are cached to avoid re-discovery.
 *
 * @example
 * ```ts
 * const cache = new ToolDiscoveryCache();
 *
 * // First call: scans all messages
 * const newTools = cache.getNewDiscoveries(messages);
 * // Returns: ['web_search', 'file_read']
 *
 * // Second call (3 new messages added): only scans new messages
 * const moreTools = cache.getNewDiscoveries(messages);
 * // Returns: ['code_exec'] (only newly discovered)
 * ```
 */
export class ToolDiscoveryCache {
  /** Set of all discovered tool names (deduped) */
  private _discoveredTools: Set<string> = new Set();
  /** Last message index that was scanned */
  private _lastScannedIndex: number = -1;

  /**
   * Scan messages for new tool_search results since the last scan.
   * Only processes messages after `_lastScannedIndex` to avoid redundant work.
   *
   * @param messages - Full conversation message array
   * @returns Array of newly discovered tool names (not previously cached)
   */
  getNewDiscoveries(messages: BaseMessage[]): string[] {
    if (messages.length === 0) {
      return [];
    }

    const startIndex = this._lastScannedIndex + 1;
    if (startIndex >= messages.length) {
      return [];
    }

    const newDiscoveries: string[] = [];

    for (let i = startIndex; i < messages.length; i++) {
      const msg = messages[i];
      if (msg.getType() !== MessageTypes.TOOL) {
        continue;
      }

      // Check if this is a tool_search result
      if ((msg as { name?: string }).name !== Constants.TOOL_SEARCH) {
        continue;
      }

      // Extract tool references from artifact
      const artifact = (msg as { artifact?: unknown }).artifact;
      if (typeof artifact === 'object' && artifact != null) {
        const refs = (
          artifact as { tool_references?: Array<{ tool_name: string }> }
        ).tool_references;
        if (refs && refs.length > 0) {
          for (const ref of refs) {
            if (!this._discoveredTools.has(ref.tool_name)) {
              // Enforce cache size limit
              if (this._discoveredTools.size >= TOOL_DISCOVERY_CACHE_MAX_SIZE) {
                break;
              }
              this._discoveredTools.add(ref.tool_name);
              newDiscoveries.push(ref.tool_name);
            }
          }
        }
      }
    }

    this._lastScannedIndex = messages.length - 1;
    return newDiscoveries;
  }

  /**
   * Returns all tool names discovered so far (across all scans).
   */
  getAllDiscoveredTools(): string[] {
    return [...this._discoveredTools];
  }

  /**
   * Check if a specific tool has been discovered.
   */
  has(toolName: string): boolean {
    return this._discoveredTools.has(toolName);
  }

  /**
   * Number of unique tools discovered.
   */
  get size(): number {
    return this._discoveredTools.size;
  }

  /**
   * Reset the cache (e.g., on graph reset).
   */
  reset(): void {
    this._discoveredTools.clear();
    this._lastScannedIndex = -1;
  }

  /**
   * Seed the cache with previously known tool names (e.g., from prior conversation turns).
   * Does not affect _lastScannedIndex — the next getNewDiscoveries call will still
   * scan all messages from the beginning.
   *
   * @param toolNames - Tool names to pre-seed into the cache
   */
  seed(toolNames: string[]): void {
    for (const name of toolNames) {
      if (this._discoveredTools.size >= TOOL_DISCOVERY_CACHE_MAX_SIZE) {
        break;
      }
      this._discoveredTools.add(name);
    }
  }
}