import { z } from 'zod/v3';
import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import type { IndemnClient } from '../sdk/client.js';
import { AgentsSDK } from '../sdk/agents.js';
import { ConfigSDK } from '../sdk/config.js';
import { FunctionsSDK } from '../sdk/functions.js';
import { KBSDK } from '../sdk/kb.js';
import { RubricsSDK } from '../sdk/rubrics.js';
import { TestSetsSDK } from '../sdk/test-sets.js';
import { EvalSDK } from '../sdk/eval.js';
import { ChatSession } from '../sdk/chat.js';
import { exportEvalMarkdown, exportEvalReport, exportAgentCard } from '../report/export.js';

// Helper: format a successful MCP tool result
function success(data: unknown) {
  return { content: [{ type: 'text' as const, text: JSON.stringify(data, null, 2) }] };
}

// Helper: format an error MCP tool result
function error(err: unknown) {
  const message = err instanceof Error ? err.message : String(err);
  return { content: [{ type: 'text' as const, text: `Error: ${message}` }], isError: true };
}

export function registerTools(server: McpServer, client: IndemnClient): void {
  const agents = new AgentsSDK(client);
  const config = new ConfigSDK(client);
  const functions = new FunctionsSDK(client);
  const kb = new KBSDK(client);
  const rubrics = new RubricsSDK(client);
  const testSets = new TestSetsSDK(client);
  const evals = new EvalSDK(client);

  // ---------------------------------------------------------------------------
  // Agent tools
  // ---------------------------------------------------------------------------

  server.tool(
    'list_agents',
    'List all agents in the organization',
    { limit: z.number().optional().describe('Maximum number of agents to return') },
    async (args) => {
      try {
        let result = await agents.listAgents();
        if (args.limit) result = result.slice(0, args.limit);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'get_agent',
    'Get details of a specific agent',
    { agent_id: z.string().describe('The agent ID') },
    async (args) => {
      try {
        const result = await agents.getAgent(args.agent_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'create_agent',
    'Create a new agent',
    {
      name: z.string().describe('Agent name'),
      channels: z.array(z.string()).optional().describe('Channels (e.g. ["WEB"])'),
    },
    async (args) => {
      try {
        const result = await agents.createAgent({ name: args.name, channels: args.channels });
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'update_agent',
    'Update an existing agent',
    {
      agent_id: z.string().describe('The agent ID'),
      name: z.string().optional().describe('New agent name'),
      description: z.string().optional().describe('New agent description'),
      webhook_url: z.string().optional().describe('Webhook URL'),
    },
    async (args) => {
      try {
        const { agent_id, ...updates } = args;
        const input: Record<string, unknown> = {};
        if (updates.name !== undefined) input.name = updates.name;
        if (updates.description !== undefined) input.description = updates.description;
        if (updates.webhook_url !== undefined) input.webhook_url = updates.webhook_url;
        const result = await agents.updateAgent(agent_id, input);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'delete_agent',
    'Delete an agent',
    { agent_id: z.string().describe('The agent ID to delete') },
    async (args) => {
      try {
        await agents.deleteAgent(args.agent_id);
        return success({ deleted: true, agent_id: args.agent_id });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'clone_agent',
    'Clone an existing agent',
    { agent_id: z.string().describe('The agent ID to clone') },
    async (args) => {
      try {
        const result = await agents.cloneAgent(args.agent_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  // ---------------------------------------------------------------------------
  // Config tools
  // ---------------------------------------------------------------------------

  server.tool(
    'get_config',
    'Get the full configuration for an agent',
    { agent_id: z.string().describe('The agent ID') },
    async (args) => {
      try {
        const result = await config.getConfig(args.agent_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'set_config',
    'Update configuration for an agent',
    {
      agent_id: z.string().describe('The agent ID'),
      system_prompt: z.string().optional().describe('System prompt text'),
      first_message: z.string().optional().describe('First message text shown to users'),
    },
    async (args) => {
      try {
        const input: Record<string, unknown> = {};
        if (args.system_prompt !== undefined) input.system_prompt = args.system_prompt;
        if (args.first_message !== undefined) input.first_message = { text: args.first_message };
        const result = await config.setConfig(args.agent_id, input);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'get_prompt',
    'Get the system prompt for an agent',
    { agent_id: z.string().describe('The agent ID') },
    async (args) => {
      try {
        const result = await config.getPrompt(args.agent_id);
        return success({ agent_id: args.agent_id, system_prompt: result ?? null });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'set_prompt',
    'Set the system prompt for an agent',
    {
      agent_id: z.string().describe('The agent ID'),
      prompt: z.string().describe('The system prompt text'),
    },
    async (args) => {
      try {
        const result = await config.setPrompt(args.agent_id, args.prompt);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  // ---------------------------------------------------------------------------
  // Function tools
  // ---------------------------------------------------------------------------

  server.tool(
    'list_functions',
    'List all functions (tools) for an agent',
    { agent_id: z.string().describe('The agent ID') },
    async (args) => {
      try {
        const result = await functions.listFunctions(args.agent_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'create_function',
    'Create a new function (tool) for an agent',
    {
      agent_id: z.string().describe('The agent ID'),
      type: z.string().describe('Function type (e.g. "rest_api")'),
      description: z.string().optional().describe('What the function does'),
      name: z.string().optional().describe('Function name'),
      url: z.string().optional().describe('Endpoint URL'),
      method: z.string().optional().describe('HTTP method (GET, POST, etc.)'),
    },
    async (args) => {
      try {
        const { agent_id, ...input } = args;
        const result = await functions.createFunction(agent_id, input);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'update_function',
    'Update an existing function (tool) for an agent',
    {
      agent_id: z.string().describe('The agent ID'),
      function_id: z.string().describe('The function ID'),
      description: z.string().optional().describe('Updated description'),
      name: z.string().optional().describe('Updated function name'),
      url: z.string().optional().describe('Updated endpoint URL'),
      method: z.string().optional().describe('Updated HTTP method'),
    },
    async (args) => {
      try {
        const { agent_id, function_id, ...input } = args;
        const result = await functions.updateFunction(agent_id, function_id, input);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'delete_function',
    'Delete a function (tool) from an agent',
    {
      agent_id: z.string().describe('The agent ID'),
      function_id: z.string().describe('The function ID to delete'),
    },
    async (args) => {
      try {
        await functions.deleteFunction(args.agent_id, args.function_id);
        return success({ deleted: true, agent_id: args.agent_id, function_id: args.function_id });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'test_function',
    'Test a function (tool) with optional parameters',
    {
      agent_id: z.string().describe('The agent ID'),
      function_id: z.string().describe('The function ID to test'),
      params: z.record(z.unknown()).optional().describe('Test parameters to send'),
    },
    async (args) => {
      try {
        const result = await functions.testFunction(
          args.agent_id,
          args.function_id,
          args.params as Record<string, unknown> | undefined,
        );
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'export_functions',
    'Export all functions from an agent',
    { agent_id: z.string().describe('The agent ID') },
    async (args) => {
      try {
        const result = await functions.exportFunctions(args.agent_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'import_functions',
    'Import functions into an agent from a URL (e.g. OpenAPI spec)',
    {
      agent_id: z.string().describe('The agent ID'),
      url: z.string().describe('URL to import functions from'),
    },
    async (args) => {
      try {
        const result = await functions.importFunctions(args.agent_id, args.url);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'list_function_parameters',
    'List parameters for a function',
    {
      agent_id: z.string().describe('The agent ID'),
      function_id: z.string().describe('The function ID'),
    },
    async (args) => {
      try {
        const result = await functions.listParameters(args.agent_id, args.function_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'add_function_parameter',
    'Add a parameter to a function',
    {
      agent_id: z.string().describe('The agent ID'),
      function_id: z.string().describe('The function ID'),
      name: z.string().describe('Parameter name'),
      type: z.string().describe('Parameter type (string, number, boolean)'),
      description: z.string().optional().describe('Parameter description'),
      required: z.boolean().optional().describe('Whether the parameter is required'),
    },
    async (args) => {
      try {
        const { agent_id, function_id, ...param } = args;
        const result = await functions.addParameter(agent_id, function_id, param);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'update_function_parameter',
    'Update a function parameter',
    {
      agent_id: z.string().describe('The agent ID'),
      function_id: z.string().describe('The function ID'),
      parameter_id: z.string().describe('The parameter ID'),
      name: z.string().optional().describe('Updated name'),
      type: z.string().optional().describe('Updated type'),
      description: z.string().optional().describe('Updated description'),
      required: z.boolean().optional().describe('Updated required flag'),
    },
    async (args) => {
      try {
        const { agent_id, function_id, parameter_id, ...param } = args;
        const result = await functions.updateParameter(agent_id, function_id, parameter_id, param);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'delete_function_parameter',
    'Delete a function parameter',
    {
      agent_id: z.string().describe('The agent ID'),
      function_id: z.string().describe('The function ID'),
      parameter_id: z.string().describe('The parameter ID'),
    },
    async (args) => {
      try {
        await functions.deleteParameter(args.agent_id, args.function_id, args.parameter_id);
        return success({ deleted: true });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'list_master_functions',
    'List available master function templates',
    {},
    async () => {
      try {
        const result = await functions.listMasterFunctions();
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'list_llm_models',
    'List available LLM models',
    {},
    async () => {
      try {
        const result = await functions.listLLMModels();
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  // ---------------------------------------------------------------------------
  // Knowledge base tools
  // ---------------------------------------------------------------------------

  server.tool(
    'list_knowledge_bases',
    'List all knowledge bases in the organization',
    { limit: z.number().optional().describe('Maximum number of KBs to return') },
    async (args) => {
      try {
        let result = await kb.listKBs();
        if (args.limit) result = result.slice(0, args.limit);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'get_knowledge_base',
    'Get details of a specific knowledge base',
    { kb_id: z.string().describe('The knowledge base ID') },
    async (args) => {
      try {
        const result = await kb.getKB(args.kb_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'create_knowledge_base',
    'Create a new knowledge base',
    {
      name: z.string().describe('Knowledge base name'),
      type: z.string().describe('Knowledge base type'),
    },
    async (args) => {
      try {
        const result = await kb.createKB({ name: args.name, type: args.type });
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'update_knowledge_base',
    'Update an existing knowledge base',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      name: z.string().optional().describe('Updated name'),
      description: z.string().optional().describe('Updated description'),
    },
    async (args) => {
      try {
        const input: Record<string, unknown> = {};
        if (args.name !== undefined) input.name = args.name;
        if (args.description !== undefined) input.description = args.description;
        const result = await kb.updateKB(args.kb_id, input);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'delete_knowledge_base',
    'Delete a knowledge base',
    { kb_id: z.string().describe('The knowledge base ID to delete') },
    async (args) => {
      try {
        await kb.deleteKB(args.kb_id);
        return success({ deleted: true, kb_id: args.kb_id });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'list_kb_data',
    'List data sources in a knowledge base',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      page: z.number().optional().describe('Page number'),
      limit: z.number().optional().describe('Items per page'),
    },
    async (args) => {
      try {
        const query: Record<string, string> = {};
        if (args.page !== undefined) query.page = String(args.page);
        if (args.limit !== undefined) query.limit = String(args.limit);
        const result = await kb.listDataSources(args.kb_id, Object.keys(query).length ? query : undefined);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'add_kb_data',
    'Add a data source entry to a knowledge base. For URL scraping use scrape_url instead. For file uploads use upload_kb_file.',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      question: z.string().optional().describe('FAQ question'),
      answer: z.string().optional().describe('FAQ answer'),
      source_url: z.string().optional().describe('Source URL (initiates scraping via import endpoint)'),
      content: z.string().optional().describe('Raw content'),
    },
    async (args) => {
      try {
        const { kb_id, ...input } = args;
        if (input.source_url) {
          const result = await kb.scrapeUrl(kb_id, input.source_url as string);
          return success(result);
        }
        const result = await kb.addDataSource(kb_id, input);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'scrape_url',
    'Scrape a website URL and add it as a data source to a knowledge base. Initiates async crawling.',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      url: z.string().describe('URL to scrape'),
      crawl_mode: z.string().optional().describe('Crawl mode: "full" (entire site) or "selected"'),
    },
    async (args) => {
      try {
        const result = await kb.scrapeUrl(args.kb_id, args.url, args.crawl_mode);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'upload_kb_file',
    'Upload a file to a knowledge base',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      file_path: z.string().describe('Path to the file to upload'),
    },
    async (args) => {
      try {
        const result = await kb.uploadFile(args.kb_id, args.file_path);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'update_kb_data',
    'Update a data source entry in a knowledge base',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      source_id: z.string().describe('The data source ID'),
      question: z.string().optional().describe('Updated question'),
      answer: z.string().optional().describe('Updated answer'),
      source_url: z.string().optional().describe('Updated source URL'),
      content: z.string().optional().describe('Updated content'),
    },
    async (args) => {
      try {
        const { kb_id, source_id, ...input } = args;
        const result = await kb.updateDataSource(kb_id, source_id, input);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'delete_kb_data',
    'Delete a data source entry from a knowledge base',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      source_id: z.string().describe('The data source ID to delete'),
    },
    async (args) => {
      try {
        await kb.deleteDataSource(args.kb_id, args.source_id);
        return success({ deleted: true, kb_id: args.kb_id, source_id: args.source_id });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'link_kb',
    'Link knowledge bases to an agent',
    {
      agent_id: z.string().describe('The agent ID'),
      kb_ids: z.array(z.string()).describe('Array of knowledge base IDs to link'),
    },
    async (args) => {
      try {
        const result = await kb.linkKB(args.agent_id, args.kb_ids);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'unlink_kb',
    'Unlink a knowledge base from an agent',
    {
      agent_id: z.string().describe('The agent ID'),
      kb_id: z.string().describe('The knowledge base ID to unlink'),
    },
    async (args) => {
      try {
        await kb.unlinkKB(args.agent_id, args.kb_id);
        return success({ unlinked: true, agent_id: args.agent_id, kb_id: args.kb_id });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'export_kb',
    'Export knowledge base data',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      format: z.string().optional().describe('Export format (e.g. "json", "csv")'),
    },
    async (args) => {
      try {
        const result = await kb.exportKB(args.kb_id, args.format || 'json');
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'import_kb',
    'Import data into a knowledge base',
    {
      kb_id: z.string().describe('The knowledge base ID'),
      data: z.record(z.unknown()).describe('Data to import'),
    },
    async (args) => {
      try {
        const result = await kb.importKB(args.kb_id, args.data);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  // ---------------------------------------------------------------------------
  // Chat tool
  // ---------------------------------------------------------------------------

  server.tool(
    'chat_with_agent',
    'Send a message to an agent and get a response. Not interactive — sends one message and returns the reply.',
    {
      agent_id: z.string().describe('The agent ID (bot type) to chat with'),
      message: z.string().describe('The message to send'),
      session_id: z.string().optional().describe('Optional session ID to continue a conversation'),
    },
    async (args) => {
      try {
        const middlewareUrl = client.getHosts().middleware;
        const chat = new ChatSession(args.agent_id, middlewareUrl, args.session_id);

        // Track streaming state to distinguish greeting from response
        let receiving = false;
        let messageCount = 0;

        chat.on('stream', (data: { isFirstChunk?: boolean; streamFinished?: boolean }) => {
          if (data.isFirstChunk) receiving = true;
          if (data.streamFinished) receiving = false;
        });

        const greetingMessages: string[] = [];
        chat.on('message', (text: string) => {
          greetingMessages.push(text);
          messageCount++;
        });

        const sessionId = await chat.connect();

        // Wait for greeting to finish streaming (up to 8s)
        await new Promise<void>((resolve) => {
          const check = setInterval(() => {
            if (!receiving) { clearInterval(check); resolve(); }
          }, 100);
          setTimeout(() => { clearInterval(check); resolve(); }, 8000);
        });

        // Now swap the listener — next message is the actual response
        const greetingCount = messageCount;
        chat.removeAllListeners('message');

        const response = await new Promise<string>((resolve, reject) => {
          const timeout = setTimeout(() => {
            chat.disconnect();
            reject(new Error('Agent did not respond within 60 seconds'));
          }, 60000);

          chat.on('message', (text: string) => {
            clearTimeout(timeout);
            resolve(text);
          });

          chat.send(args.message);
        });

        chat.disconnect();

        return success({ session_id: sessionId, greeting: greetingMessages.slice(0, greetingCount).join('\n'), response });
      } catch (err) { return error(err); }
    },
  );

  // ---------------------------------------------------------------------------
  // Eval tools
  // ---------------------------------------------------------------------------

  server.tool(
    'run_evaluation',
    'Trigger an evaluation run for a bot against a test set',
    {
      bot_id: z.string().describe('The bot ID to evaluate'),
      test_set_id: z.string().describe('The test set ID'),
      rubric_id: z.string().optional().describe('Optional rubric ID'),
      concurrency: z.number().optional().describe('Number of concurrent evaluations'),
      eval_model: z.string().optional().describe('Model to use for evaluation judging'),
      keep: z.boolean().optional().describe('Keep evaluation data after completion'),
      limit: z.number().optional().describe('Limit items to evaluate'),
      component_scope: z.string().optional().describe('Component scope filter'),
      component_ids: z.array(z.string()).optional().describe('Component ID filter'),
      mode: z.string().optional().describe('Evaluation mode'),
    },
    async (args) => {
      try {
        const input: Record<string, unknown> = {
          bot_id: args.bot_id,
          test_set_id: args.test_set_id,
        };
        if (args.rubric_id !== undefined) input.rubric_id = args.rubric_id;
        if (args.concurrency !== undefined) input.concurrency = args.concurrency;
        if (args.eval_model !== undefined) input.eval_model = args.eval_model;
        if (args.keep !== undefined) input.keep = args.keep;
        if (args.limit !== undefined) input.limit = args.limit;
        if (args.component_scope !== undefined) input.component_scope = args.component_scope;
        if (args.component_ids !== undefined) input.component_ids = args.component_ids;
        if (args.mode !== undefined) input.mode = args.mode;
        const result = await evals.triggerEvaluation(input as any);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'get_eval_status',
    'Get the status of an evaluation run',
    { run_id: z.string().describe('The evaluation run ID') },
    async (args) => {
      try {
        const result = await evals.getRunStatus(args.run_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'get_eval_results',
    'Get the detailed results of a completed evaluation run',
    { run_id: z.string().describe('The evaluation run ID') },
    async (args) => {
      try {
        const result = await evals.getRunResults(args.run_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'get_bot_context',
    'Get the full context of a bot (prompt, tools, KBs, LLM config)',
    { bot_id: z.string().describe('The bot ID') },
    async (args) => {
      try {
        const result = await evals.getBotContext(args.bot_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  // ---------------------------------------------------------------------------
  // Rubric tools
  // ---------------------------------------------------------------------------

  server.tool(
    'list_rubrics',
    'List evaluation rubrics, optionally filtered by agent',
    { agent_id: z.string().optional().describe('Filter rubrics by agent ID') },
    async (args) => {
      try {
        const result = await rubrics.listRubrics(args.agent_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'get_rubric',
    'Get details of a specific rubric',
    { rubric_id: z.string().describe('The rubric ID') },
    async (args) => {
      try {
        const result = await rubrics.getRubric(args.rubric_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'create_rubric',
    'Create a new evaluation rubric',
    {
      name: z.string().describe('Rubric name'),
      rules: z.array(z.record(z.unknown())).describe('Array of rubric rules'),
      agent_id: z.string().optional().describe('Agent ID to associate the rubric with'),
      description: z.string().optional().describe('Rubric description'),
      severity_definitions: z.record(z.string()).optional().describe('Severity level definitions'),
    },
    async (args) => {
      try {
        const input: Record<string, unknown> = { name: args.name, rules: args.rules };
        if (args.agent_id !== undefined) input.agent_id = args.agent_id;
        if (args.description !== undefined) input.description = args.description;
        if (args.severity_definitions !== undefined) input.severity_definitions = args.severity_definitions;
        const result = await rubrics.createRubric(input as any);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'update_rubric',
    'Update an existing rubric',
    {
      rubric_id: z.string().describe('The rubric ID'),
      name: z.string().optional().describe('Updated name'),
      description: z.string().optional().describe('Updated description'),
      rules: z.array(z.record(z.unknown())).optional().describe('Updated rules'),
    },
    async (args) => {
      try {
        const input: Record<string, unknown> = {};
        if (args.name !== undefined) input.name = args.name;
        if (args.description !== undefined) input.description = args.description;
        if (args.rules !== undefined) input.rules = args.rules;
        const result = await rubrics.updateRubric(args.rubric_id, input as any);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'delete_rubric',
    'Delete a rubric',
    { rubric_id: z.string().describe('The rubric ID to delete') },
    async (args) => {
      try {
        await rubrics.deleteRubric(args.rubric_id);
        return success({ deleted: true, rubric_id: args.rubric_id });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'list_rubric_versions',
    'List all versions of a rubric',
    { rubric_id: z.string().describe('The rubric ID') },
    async (args) => {
      try {
        const result = await rubrics.listVersions(args.rubric_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  // ---------------------------------------------------------------------------
  // Test set tools
  // ---------------------------------------------------------------------------

  server.tool(
    'list_test_sets',
    'List test sets, optionally filtered by agent',
    { agent_id: z.string().optional().describe('Filter test sets by agent ID') },
    async (args) => {
      try {
        const result = await testSets.listTestSets(args.agent_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'get_test_set',
    'Get details of a specific test set',
    { test_set_id: z.string().describe('The test set ID') },
    async (args) => {
      try {
        const result = await testSets.getTestSet(args.test_set_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'create_test_set',
    'Create a new test set',
    {
      name: z.string().describe('Test set name'),
      items: z.array(z.record(z.unknown())).describe('Array of test items'),
      agent_id: z.string().optional().describe('Agent ID to associate the test set with'),
      description: z.string().optional().describe('Test set description'),
    },
    async (args) => {
      try {
        const input: Record<string, unknown> = { name: args.name, items: args.items };
        if (args.agent_id !== undefined) input.agent_id = args.agent_id;
        if (args.description !== undefined) input.description = args.description;
        const result = await testSets.createTestSet(input as any);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'update_test_set',
    'Update an existing test set',
    {
      test_set_id: z.string().describe('The test set ID'),
      name: z.string().optional().describe('Updated name'),
      description: z.string().optional().describe('Updated description'),
      items: z.array(z.record(z.unknown())).optional().describe('Updated test items'),
    },
    async (args) => {
      try {
        const input: Record<string, unknown> = {};
        if (args.name !== undefined) input.name = args.name;
        if (args.description !== undefined) input.description = args.description;
        if (args.items !== undefined) input.items = args.items;
        const result = await testSets.updateTestSet(args.test_set_id, input as any);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'delete_test_set',
    'Delete a test set',
    { test_set_id: z.string().describe('The test set ID to delete') },
    async (args) => {
      try {
        await testSets.deleteTestSet(args.test_set_id);
        return success({ deleted: true, test_set_id: args.test_set_id });
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'list_test_set_versions',
    'List all versions of a test set',
    { test_set_id: z.string().describe('The test set ID') },
    async (args) => {
      try {
        const result = await testSets.listVersions(args.test_set_id);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  // ---------------------------------------------------------------------------
  // Export tools
  // ---------------------------------------------------------------------------

  server.tool(
    'export_eval_markdown',
    'Export evaluation results as a structured markdown file with conversations, scores, and criteria',
    {
      run_id: z.string().describe('The evaluation run ID'),
      output_path: z.string().optional().describe('Output file path (defaults to current directory)'),
    },
    async (args) => {
      try {
        const result = await exportEvalMarkdown(args.run_id, args.output_path, client);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'export_eval_report',
    'Generate a branded PDF evaluation report with cover page, scores, and per-item details',
    {
      run_id: z.string().describe('The evaluation run ID'),
      output_path: z.string().optional().describe('Output file path (defaults to current directory)'),
    },
    async (args) => {
      try {
        const result = await exportEvalReport(args.run_id, args.output_path, client);
        return success(result);
      } catch (err) { return error(err); }
    },
  );

  server.tool(
    'export_agent_card',
    'Generate a branded PDF agent card with system prompt, tools, KBs, LLM config, and evaluation history',
    {
      agent_id: z.string().describe('The agent ID'),
      output_path: z.string().optional().describe('Output file path (defaults to current directory)'),
    },
    async (args) => {
      try {
        const result = await exportAgentCard(args.agent_id, args.output_path, client);
        return success(result);
      } catch (err) { return error(err); }
    },
  );
}