// src/utils/__tests__/toolDiscoveryCache.test.ts
import {
  ToolMessage,
  AIMessageChunk,
  HumanMessage,
  SystemMessage,
} from '@langchain/core/messages';
import type { BaseMessage } from '@langchain/core/messages';
import { ToolDiscoveryCache } from '../toolDiscoveryCache';
import { Constants } from '@/common';

/**
 * Creates a mock tool_search result message.
 */
function createToolSearchResult(
  toolNames: string[],
  toolCallId: string = 'tc_1'
): ToolMessage {
  return new ToolMessage({
    content: `Found ${toolNames.length} tools`,
    tool_call_id: toolCallId,
    name: Constants.TOOL_SEARCH,
    artifact: {
      tool_references: toolNames.map((name) => ({ tool_name: name })),
    },
  });
}

/**
 * Creates a mock AI message with tool calls.
 */
function createAIWithToolCalls(toolCallIds: string[]): AIMessageChunk {
  return new AIMessageChunk({
    content: 'I will search for tools',
    tool_calls: toolCallIds.map((id) => ({
      id,
      name: Constants.TOOL_SEARCH,
      args: { query: 'test' },
    })),
  });
}

describe('ToolDiscoveryCache', () => {
  let cache: ToolDiscoveryCache;

  beforeEach(() => {
    cache = new ToolDiscoveryCache();
  });

  describe('getNewDiscoveries', () => {
    it('returns empty array for empty messages', () => {
      expect(cache.getNewDiscoveries([])).toEqual([]);
    });

    it('discovers tools from tool_search results', () => {
      const messages: BaseMessage[] = [
        new SystemMessage('You are helpful'),
        new HumanMessage('Find tools'),
        createAIWithToolCalls(['tc_1']),
        createToolSearchResult(['web_search', 'file_read'], 'tc_1'),
      ];

      const result = cache.getNewDiscoveries(messages);
      expect(result).toEqual(['web_search', 'file_read']);
      expect(cache.size).toBe(2);
    });

    it('only scans new messages on subsequent calls', () => {
      const messages: BaseMessage[] = [
        new HumanMessage('msg1'),
        createAIWithToolCalls(['tc_1']),
        createToolSearchResult(['tool_a'], 'tc_1'),
      ];

      // First scan
      const first = cache.getNewDiscoveries(messages);
      expect(first).toEqual(['tool_a']);

      // Add more messages
      messages.push(
        new HumanMessage('msg2'),
        createAIWithToolCalls(['tc_2']),
        createToolSearchResult(['tool_b'], 'tc_2')
      );

      // Second scan: only finds tool_b (tool_a already cached)
      const second = cache.getNewDiscoveries(messages);
      expect(second).toEqual(['tool_b']);
      expect(cache.size).toBe(2);
    });

    it('deduplicates tool names across scans', () => {
      const messages: BaseMessage[] = [
        createAIWithToolCalls(['tc_1']),
        createToolSearchResult(['tool_a', 'tool_b'], 'tc_1'),
      ];

      cache.getNewDiscoveries(messages);

      // Add another search that returns tool_a again
      messages.push(
        createAIWithToolCalls(['tc_2']),
        createToolSearchResult(['tool_a', 'tool_c'], 'tc_2')
      );

      const second = cache.getNewDiscoveries(messages);
      // tool_a is already cached, only tool_c is new
      expect(second).toEqual(['tool_c']);
      expect(cache.size).toBe(3);
    });

    it('ignores non-tool-search tool messages', () => {
      const messages: BaseMessage[] = [
        createAIWithToolCalls(['tc_1']),
        new ToolMessage({
          content: 'result',
          tool_call_id: 'tc_1',
          name: 'some_other_tool',
        }),
      ];

      const result = cache.getNewDiscoveries(messages);
      expect(result).toEqual([]);
    });

    it('returns empty when no new messages since last scan', () => {
      const messages: BaseMessage[] = [
        createAIWithToolCalls(['tc_1']),
        createToolSearchResult(['tool_a'], 'tc_1'),
      ];

      cache.getNewDiscoveries(messages);
      // No new messages added
      const second = cache.getNewDiscoveries(messages);
      expect(second).toEqual([]);
    });
  });

  describe('has', () => {
    it('returns true for discovered tools', () => {
      const messages: BaseMessage[] = [
        createAIWithToolCalls(['tc_1']),
        createToolSearchResult(['tool_a'], 'tc_1'),
      ];

      cache.getNewDiscoveries(messages);
      expect(cache.has('tool_a')).toBe(true);
      expect(cache.has('tool_b')).toBe(false);
    });
  });

  describe('getAllDiscoveredTools', () => {
    it('returns all discovered tool names', () => {
      const messages: BaseMessage[] = [
        createAIWithToolCalls(['tc_1']),
        createToolSearchResult(['tool_a', 'tool_b'], 'tc_1'),
      ];

      cache.getNewDiscoveries(messages);
      expect(cache.getAllDiscoveredTools()).toEqual(
        expect.arrayContaining(['tool_a', 'tool_b'])
      );
    });
  });

  describe('seed', () => {
    it('pre-populates the cache with known tool names', () => {
      cache.seed(['tool_x', 'tool_y']);
      expect(cache.size).toBe(2);
      expect(cache.has('tool_x')).toBe(true);
      expect(cache.has('tool_y')).toBe(true);
    });

    it('seeded tools are treated as already discovered', () => {
      cache.seed(['tool_a']);

      const messages: BaseMessage[] = [
        createAIWithToolCalls(['tc_1']),
        createToolSearchResult(['tool_a', 'tool_b'], 'tc_1'),
      ];

      // tool_a is already seeded, only tool_b should be new
      const result = cache.getNewDiscoveries(messages);
      expect(result).toEqual(['tool_b']);
    });
  });

  describe('reset', () => {
    it('clears all state', () => {
      cache.seed(['tool_a']);
      expect(cache.size).toBe(1);

      cache.reset();
      expect(cache.size).toBe(0);
      expect(cache.has('tool_a')).toBe(false);
    });

    it('allows re-discovery after reset', () => {
      const messages: BaseMessage[] = [
        createAIWithToolCalls(['tc_1']),
        createToolSearchResult(['tool_a'], 'tc_1'),
      ];

      cache.getNewDiscoveries(messages);
      expect(cache.size).toBe(1);

      cache.reset();

      // Same messages should produce discoveries again
      const result = cache.getNewDiscoveries(messages);
      expect(result).toEqual(['tool_a']);
    });
  });
});