import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { Logger } from '../utils/logger.util.js';
import { formatErrorForMcpTool } from '../utils/error.util.js';
import { OpenTelemetryController } from '../controllers/opentelemetry-llm.controller.js';
import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
import {
	OpenTelemetryLlmInputSchema,
	OpenTelemetryLlmInputSchemaType,
} from '../types/opentelemetry-llm.types.js';

/**
 * @function captureOpenTelemetryLlmObservability
 * @description MCP Tool handler to capture LLM observability events using OpenTelemetry.
 *              It records metrics, traces, and spans for LLM requests that can be sent to any OpenTelemetry-compatible backend.
 * @param {OpenTelemetryLlmInputSchemaType} args - Arguments provided to the tool.
 * @returns {Promise<CallToolResult>} Formatted response for the MCP.
 * @throws {McpError} Formatted error if the controller or service layer encounters an issue.
 */
async function captureOpenTelemetryLlmObservability(
	args: OpenTelemetryLlmInputSchemaType,
): Promise<CallToolResult> {
	const methodLogger = Logger.forContext(
		'opentelemetry.tool',
		'captureOpenTelemetryLlmObservability',
	);
	methodLogger.debug('Capture LLM Observability with OpenTelemetry...', args);

	try {
		// Parse and validate arguments
		const validatedArgs = OpenTelemetryLlmInputSchema.parse(args);

		// Pass validated args to the controller
		const result = await OpenTelemetryController.capture(validatedArgs);
		methodLogger.debug('Got response from controller', result);

		// Format the response for the MCP tool
		return {
			content: [
				{
					type: 'text' as const,
					text: result.content,
				},
			],
		};
	} catch (error) {
		methodLogger.error(
			'Error tracking LLM generation with OpenTelemetry',
			error,
		);
		return formatErrorForMcpTool(error);
	}
}

/**
 * @function registerTools
 * @description Registers the OpenTelemetry LLM observability tool with the MCP server.
 *
 * @param {McpServer} server - The MCP server instance.
 */
function registerTools(server: McpServer) {
	const methodLogger = Logger.forContext(
		'opentelemetry.tool',
		'registerTools',
	);
	methodLogger.debug('Registering OpenTelemetry LLM observability tools...');

	server.tool(
		'llm_observability_otel',
		`Captures LLM usage using OpenTelemetry for observability, including requests, responses, and performance metrics. Works with any OpenTelemetry-compatible backend like Jaeger, New Relic, Grafana, etc.`,
		OpenTelemetryLlmInputSchema.shape,
		captureOpenTelemetryLlmObservability,
	);

	methodLogger.debug('Successfully registered llm_observability_otel tool.');
}

export default { registerTools };