// src/messages/toolDiscovery.ts
import { AIMessageChunk, ToolMessage } from '@langchain/core/messages';
import type { BaseMessage } from '@langchain/core/messages';
import { Constants, MessageTypes } from '@/common';
import { findLastIndex } from './core';

type ToolSearchArtifact = {
  tool_references?: Array<{ tool_name: string }>;
};

/**
 * Extracts discovered tool names from tool search results in the current turn.
 * Only processes tool search messages after the latest AI message with tool calls.
 *
 * Similar pattern to formatArtifactPayload - finds relevant messages efficiently
 * by identifying the latest AI parent and only processing subsequent tool messages.
 *
 * @param messages - All messages in the conversation
 * @returns Array of discovered tool names (empty if no new discoveries)
 */
export function extractToolDiscoveries(messages: BaseMessage[]): string[] {
  if (messages.length === 0) return [];
  const lastMessage = messages[messages.length - 1];
  // Use getType() instead of instanceof to avoid module mismatch issues
  if (lastMessage.getType() !== MessageTypes.TOOL) return [];
  const lastToolMessage = lastMessage as ToolMessage;

  // Find the latest AIMessage with tool_calls that this tool message belongs to
  const latestAIParentIndex = findLastIndex(
    messages,
    (msg) =>
      (msg instanceof AIMessageChunk &&
        (msg.tool_calls?.length ?? 0) > 0 &&
        msg.tool_calls?.some((tc) => tc.id === lastToolMessage.tool_call_id)) ??
      false
  );

  if (latestAIParentIndex === -1) return [];

  // Collect tool_call_ids from the AI message
  const aiMessage = messages[latestAIParentIndex] as AIMessageChunk;
  const toolCallIds = new Set(aiMessage.tool_calls?.map((tc) => tc.id) ?? []);

  // Only process tool search results after the AI message that belong to this turn
  const discoveredNames: string[] = [];
  for (let i = latestAIParentIndex + 1; i < messages.length; i++) {
    const msg = messages[i];
    // Use getType() instead of instanceof to avoid module mismatch issues
    if (msg.getType() !== MessageTypes.TOOL) continue;
    const toolMsg = msg as ToolMessage;
    if (toolMsg.name !== Constants.TOOL_SEARCH) continue;
    if (!toolCallIds.has(toolMsg.tool_call_id)) continue;

    // This is a tool search result from the current turn
    if (typeof toolMsg.artifact === 'object' && toolMsg.artifact != null) {
      const artifact = toolMsg.artifact as ToolSearchArtifact;
      if (artifact.tool_references && artifact.tool_references.length > 0) {
        for (const ref of artifact.tool_references) {
          discoveredNames.push(ref.tool_name);
        }
      }
    }
  }

  return discoveredNames;
}

/**
 * Checks if the current turn has any tool search results.
 * Quick check to avoid full extraction when not needed.
 */
export function hasToolSearchInCurrentTurn(messages: BaseMessage[]): boolean {
  if (messages.length === 0) return false;
  const lastMessage = messages[messages.length - 1];
  // Use getType() instead of instanceof to avoid module mismatch issues
  if (lastMessage.getType() !== MessageTypes.TOOL) return false;
  const lastToolMessage = lastMessage as ToolMessage;

  // Find the latest AIMessage with tool_calls
  const latestAIParentIndex = findLastIndex(
    messages,
    (msg) =>
      (msg instanceof AIMessageChunk &&
        (msg.tool_calls?.length ?? 0) > 0 &&
        msg.tool_calls?.some((tc) => tc.id === lastToolMessage.tool_call_id)) ??
      false
  );

  if (latestAIParentIndex === -1) return false;

  const aiMessage = messages[latestAIParentIndex] as AIMessageChunk;
  const toolCallIds = new Set(aiMessage.tool_calls?.map((tc) => tc.id) ?? []);

  // Check if any tool search results exist after the AI message
  // Use getType() instead of instanceof to avoid module mismatch issues
  for (let i = latestAIParentIndex + 1; i < messages.length; i++) {
    const msg = messages[i];
    if (
      msg.getType() === MessageTypes.TOOL &&
      (msg as ToolMessage).name === Constants.TOOL_SEARCH &&
      toolCallIds.has((msg as ToolMessage).tool_call_id)
    ) {
      return true;
    }
  }

  return false;
}