import { tool } from 'ai';
import { z } from 'zod';
import { edgeService } from '@/services/database/edges';
import { formatNodeForChat } from '../infrastructure/nodeFormatter';

export const queryEdgeTool = tool({
  description: 'Find edges by node/direction/source/ID',
  inputSchema: z.object({
    filters: z.object({
      node_id: z.number().optional().describe('Get all edges connected to this specific node (both incoming and outgoing)'),
      from_node_id: z.number().optional().describe('Get edges originating from this specific node'),
      to_node_id: z.number().optional().describe('Get edges pointing to this specific node'),
      source: z.enum(['user', 'ai_similarity', 'helper_name']).optional().describe('Filter edges by their source type'),
      edge_id: z.number().optional().describe('Get a specific edge by its ID'),
      limit: z.number().min(1).max(100).default(20).describe('Maximum number of results to return')
    }).optional().describe('Filters to apply when querying edges')
  }),
  execute: async ({ filters = {} }) => {
    console.log('🔍 QueryEdge tool called with filters:', JSON.stringify(filters, null, 2));
    
    try {
      let result;
      let message = '';

      // Handle specific edge ID lookup
      if (filters.edge_id) {
        const edge = await edgeService.getEdgeById(filters.edge_id);
        return {
          success: true,
          data: {
            edges: edge ? [edge] : [],
            count: edge ? 1 : 0,
            filters_applied: filters
          },
          message: edge ? `Found edge ${filters.edge_id}` : `Edge ${filters.edge_id} not found`
        };
      }

      // Handle node connections (most common use case)
      if (filters.node_id) {
        const connections = await edgeService.getNodeConnections(filters.node_id);
        const edges = connections.map(conn => conn.edge);
        
        // Apply additional filters if specified
        let filteredEdges = edges;
        if (filters.source) {
          filteredEdges = edges.filter(edge => edge.source === filters.source);
        }
        
        // Apply limit and format connected nodes
        const limitedConnections = connections.slice(0, filters.limit || 20);
        const formattedConnections = limitedConnections.map(connection => {
          const formattedNode = formatNodeForChat({
            id: connection.connected_node.id,
            title: connection.connected_node.title,
            dimensions: connection.connected_node.dimensions || []
          });
          
          return {
            ...connection,
            connected_node: {
              ...connection.connected_node,
              formatted_display: formattedNode
            }
          };
        });
        
        // Create message with formatted connected nodes
        const connectedNodeLabels = formattedConnections.map(conn => conn.connected_node.formatted_display).join(', ');
        const message = `Found ${filteredEdges.length} edges for node ${filters.node_id}${connectedNodeLabels ? `. Connected nodes: ${connectedNodeLabels}` : ''}`;

        return {
          success: true,
          data: {
            edges: filteredEdges.slice(0, filters.limit || 20),
            connections: formattedConnections,
            count: filteredEdges.length,
            filters_applied: filters
          },
          message: message
        };
      }

      // Handle directional queries or get all edges
      const allEdges = await edgeService.getEdges();
      let filteredEdges = allEdges;

      // Apply filters
      if (filters.from_node_id) {
        filteredEdges = filteredEdges.filter(edge => edge.from_node_id === filters.from_node_id);
        message += `from node ${filters.from_node_id} `;
      }
      
      if (filters.to_node_id) {
        filteredEdges = filteredEdges.filter(edge => edge.to_node_id === filters.to_node_id);
        message += `to node ${filters.to_node_id} `;
      }
      
      if (filters.source) {
        filteredEdges = filteredEdges.filter(edge => edge.source === filters.source);
        message += `with source ${filters.source} `;
      }

      // Apply limit
      const limitedEdges = filteredEdges.slice(0, filters.limit || 20);

      return {
        success: true,
        data: {
          edges: limitedEdges,
          count: filteredEdges.length,
          total_available: allEdges.length,
          filters_applied: filters
        },
        message: `Found ${filteredEdges.length} edges ${message}`.trim()
      };
    } catch (error) {
      console.error('QueryEdge tool error:', error);
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Failed to query edges',
        data: {
          edges: [],
          count: 0,
          filters_applied: filters
        }
      };
    }
  }
});
