import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { McpRequestMonitor } from "./requestMonitor.js";
import { broadcastSse } from "./monitor.js";

// Tool handlers map for retry functionality (global, shared across sessions)
export const toolHandlers = new Map();

// Registration functions - set by index.js
let _registerPrompts = null;
let _registerTools = null;
let _registerResources = null;

// Shared request monitor instance
let _requestMonitor = null;

/**
 * Set the registration functions (called once from index.js)
 */
export function setRegistrationFunctions({ registerPrompts, registerTools, registerResources }) {
    _registerPrompts = registerPrompts;
    _registerTools = registerTools;
    _registerResources = registerResources;
}

/**
 * Create and configure the MCP server
 * Each session should get its own server instance
 */
export function createMcpServer() {
	const server = new McpServer({
		name: "acai-code-mcp-server",
		version: "1.0.0",
	});

    // Intercept tool registration to capture handlers for retry/resend from monitor
    const originalTool = server.tool.bind(server);
	server.tool = (name, ...args) => {
		const handler = args[args.length - 1];
		toolHandlers.set(name, { handler });
		return originalTool(name, ...args);
	};

	return server;
}

/**
 * Create a fully configured server for a new session
 * This creates a new McpServer instance with all tools/prompts/resources registered
 * IMPORTANT: MCP SDK only supports one transport per server, so each session needs its own server
 */
export function createSessionServer() {
    const server = createMcpServer();

    // Wrap with request monitoring BEFORE registering tools/prompts
    if (_requestMonitor) {
        wrapServerWithMonitor(server, _requestMonitor);
    }

    // Register all tools, prompts, and resources
    if (_registerPrompts) _registerPrompts(server);
    if (_registerTools) _registerTools(server);
    if (_registerResources) _registerResources(server);

    return server;
}

/**
 * Wrap a server's request handlers with monitoring
 */
function wrapServerWithMonitor(server, monitor) {
    const originalSetRequestHandler = server.server.setRequestHandler.bind(server.server);
    server.server.setRequestHandler = (schema, handler) => {
        const method = schema.shape.method.value;
        return originalSetRequestHandler(schema, async (request, extra) => {
            const entry = monitor.start(method, request, extra);
            try {
                const result = await handler(request, extra);
                monitor.finish(entry, result);
                return result;
            } catch (error) {
                monitor.fail(entry, error);
                throw error;
            }
        });
    };
}

/**
 * Create and configure the request monitor
 */
export function createRequestMonitor() {
    _requestMonitor = new McpRequestMonitor();

    // Broadcast summary updates via SSE
    _requestMonitor.on("summary", (summary) => {
        broadcastSse("summary", summary);
    });

    return _requestMonitor;
}