/**
 * Tool Discovery System
 *
 * Enables agents to discover and request tools on-demand, reducing initial context size.
 */

import { createTool } from '@mastra/core';
import { z } from 'zod';
import {
  estimateTokenCost,
  getRecommendedTools,
  getToolMetadata,
  getToolsByCategory,
  searchTools,
  TOOL_CATEGORIES,
  type ToolCategory,
} from './registry';

/**
 * Discover Tools by Category or Search Query
 *
 * Allows agents to find available tools without loading them into context.
 */
export const discoverToolsTool = createTool({
  id: 'discover-tools',
  description:
    'Discover available tools by category or search query. Use this to find tools you need before requesting them.',
  inputSchema: z.object({
    category: z
      .enum([
        'file-discovery',
        'file-read',
        'file-write',
        'spreadsheet',
        'text-document',
        'presentation',
        'e2b-execution',
        'e2b-files',
        'e2b-environment',
        'govcon',
        'classifiers',
        'fpds',
        'gsa',
        'rag',
        'search',
      ])
      .optional()
      .describe('Filter by tool category'),
    searchQuery: z.string().optional().describe('Search tools by description or use case'),
    showRecommended: z
      .boolean()
      .optional()
      .describe('Show recommended tools based on currently loaded tools'),
  }),
  outputSchema: z.object({
    tools: z.array(
      z.object({
        name: z.string(),
        category: z.string(),
        description: z.string(),
        useCases: z.array(z.string()),
        tokenCost: z.enum(['low', 'medium', 'high']),
        dependencies: z.array(z.string()).optional(),
      })
    ),
    totalTools: z.number(),
    categories: z.array(z.string()).optional(),
  }),
  execute: async ({ context }) => {
    const { category, searchQuery, showRecommended } = context;

    let tools;

    if (showRecommended) {
      // Get currently loaded tools from runtime context (would need to track this)
      // For now, return empty array as example
      const loadedTools: string[] = [];
      tools = getRecommendedTools(loadedTools);
    } else if (category) {
      tools = getToolsByCategory(category as ToolCategory);
    } else if (searchQuery) {
      tools = searchTools(searchQuery);
    } else {
      // Return all categories
      const categories = Object.keys(TOOL_CATEGORIES);
      return {
        tools: [],
        totalTools: Object.keys(TOOL_CATEGORIES).reduce(
          (sum, cat) => sum + TOOL_CATEGORIES[cat as ToolCategory].length,
          0
        ),
        categories,
      };
    }

    return {
      tools: tools.map((t) => ({
        name: t.name,
        category: t.category,
        description: t.description,
        useCases: t.useCases,
        tokenCost: t.tokenCost,
        dependencies: t.dependencies,
      })),
      totalTools: tools.length,
    };
  },
});

/**
 * Get Tool Details
 *
 * Returns detailed information about a specific tool including usage examples.
 */
export const getToolDetailsTool = createTool({
  id: 'get-tool-details',
  description:
    'Get detailed information about a specific tool including usage patterns and examples. Use this before using a tool for the first time.',
  inputSchema: z.object({
    toolName: z.string().describe('Name of the tool to get details for'),
  }),
  outputSchema: z.object({
    name: z.string(),
    category: z.string(),
    description: z.string(),
    useCases: z.array(z.string()),
    tokenCost: z.enum(['low', 'medium', 'high']),
    dependencies: z.array(z.string()).optional(),
    estimatedTokens: z.number(),
    knowledgeBaseRef: z.string().optional(),
  }),
  execute: async ({ context }) => {
    const { toolName } = context;
    const tool = getToolMetadata(toolName);

    if (!tool) {
      throw new Error(`Tool '${toolName}' not found in registry`);
    }

    // Map tool names to knowledge base references
    const knowledgeBaseRefs: Record<string, string> = {
      'execute-python': 'e2b-architecture.md#execution-tools',
      'run-python-code': 'e2b-architecture.md#repl-state-persistence',
      'upload-to-sandbox': 'e2b-architecture.md#vfs-sandbox-integration',
      'download-from-sandbox': 'e2b-architecture.md#vfs-sandbox-integration',
      'read-file': 'file-operations.md#smart-reading',
      'create-text-document': 'file-operations.md#text-documents',
      'apply-ai-patch': 'file-operations.md#morphllm-patches',
      'create-spreadsheet': 'spreadsheet-formulas.md#creation',
      'update-sheet-cell': 'spreadsheet-formulas.md#formulas',
    };

    return {
      name: tool.name,
      category: tool.category,
      description: tool.description,
      useCases: tool.useCases,
      tokenCost: tool.tokenCost,
      dependencies: tool.dependencies,
      estimatedTokens: estimateTokenCost([toolName]),
      knowledgeBaseRef: knowledgeBaseRefs[toolName],
    };
  },
});

/**
 * Estimate Context Size
 *
 * Helps agents understand token budget before loading tools.
 */
export const estimateContextSizeTool = createTool({
  id: 'estimate-context-size',
  description:
    'Estimate total token cost for a set of tools. Use this to plan which tools to load within token budget.',
  inputSchema: z.object({
    toolNames: z.array(z.string()).describe('List of tool names to estimate'),
  }),
  outputSchema: z.object({
    totalEstimatedTokens: z.number(),
    breakdown: z.array(
      z.object({
        toolName: z.string(),
        tokenCost: z.enum(['low', 'medium', 'high']),
        estimatedTokens: z.number(),
      })
    ),
    withinBudget: z.boolean(),
    recommendation: z.string().optional(),
  }),
  execute: async ({ context }) => {
    const { toolNames } = context;
    const costs = { low: 100, medium: 300, high: 800 };

    const breakdown = toolNames.map((name) => {
      const tool = getToolMetadata(name);
      return {
        toolName: name,
        tokenCost: tool?.tokenCost || ('medium' as const),
        estimatedTokens: tool ? costs[tool.tokenCost] : 300,
      };
    });

    const totalEstimatedTokens = breakdown.reduce((sum, item) => sum + item.estimatedTokens, 0);

    // Cerebras free tier has 8192 token limit
    // Reserve 2000 for instructions, 3000 for conversation
    const toolBudget = 3000;
    const withinBudget = totalEstimatedTokens <= toolBudget;

    let recommendation;
    if (!withinBudget) {
      recommendation = `Total ${totalEstimatedTokens} tokens exceeds recommended budget of ${toolBudget}. Consider loading high-cost tools only when needed.`;
    }

    return {
      totalEstimatedTokens,
      breakdown,
      withinBudget,
      recommendation,
    };
  },
});