import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
import { serverStreamResultPatchNotificationSchema } from "./types.js";
import { resolveNpxBinary } from "./npx-resolver.js";
import { logger } from "./logger.js";
import { McpOAuthProvider } from "./mcp-oauth-provider.js";
import { supportsOAuth } from "./mcp-auth-flow.js";
import { registerSamplingHandler } from "./sampling-handler.js";
import { interpolateEnvRecord, resolveBearerToken, resolveConfigPath } from "./utils.js";
export class McpServerManager {
    connections = new Map();
    connectPromises = new Map();
    uiStreamListeners = new Map();
    samplingConfig;
    setSamplingConfig(config) {
        this.samplingConfig = config;
    }
    async connect(name, definition) {
        // Dedupe concurrent connection attempts
        if (this.connectPromises.has(name)) {
            return this.connectPromises.get(name);
        }
        // Reuse existing connection if healthy
        const existing = this.connections.get(name);
        if (existing?.status === "connected") {
            existing.lastUsedAt = Date.now();
            return existing;
        }
        const promise = this.createConnection(name, definition);
        this.connectPromises.set(name, promise);
        try {
            const connection = await promise;
            this.connections.set(name, connection);
            return connection;
        }
        finally {
            this.connectPromises.delete(name);
        }
    }
    async createConnection(name, definition) {
        const client = this.createClient(name);
        let transport;
        try {
            if (definition.command) {
                let command = definition.command;
                let args = definition.args ?? [];
                if (command === "npx" || command === "npm") {
                    const resolved = await resolveNpxBinary(command, args);
                    if (resolved) {
                        command = resolved.isJs ? "node" : resolved.binPath;
                        args = resolved.isJs ? [resolved.binPath, ...resolved.extraArgs] : resolved.extraArgs;
                        logger.debug(`${name} resolved to ${resolved.binPath} (skipping npm parent)`);
                    }
                }
                transport = new StdioClientTransport({
                    command,
                    args,
                    env: resolveEnv(definition.env),
                    cwd: resolveConfigPath(definition.cwd),
                    stderr: definition.debug ? "inherit" : "ignore",
                });
            }
            else if (definition.url) {
                // HTTP transport with fallback
                transport = await this.createHttpTransport(definition, name);
            }
            else {
                throw new Error(`Server ${name} has no command or url`);
            }
            await client.connect(transport);
            this.attachAdapterNotificationHandlers(name, client);
            // Discover tools and resources
            const [tools, resources] = await Promise.all([
                this.fetchAllTools(client),
                this.fetchAllResources(client),
            ]);
            return {
                client,
                transport,
                definition,
                tools,
                resources,
                lastUsedAt: Date.now(),
                inFlight: 0,
                status: "connected",
            };
        }
        catch (error) {
            // Clean up client and transport on any error.
            await client.close().catch(() => { });
            if (transport) {
                await transport.close().catch(() => { });
            }
            // Check for UnauthorizedError — server requires OAuth.
            // This can fire from createHttpTransport probe or client.connect.
            // transport may be undefined if createHttpTransport threw during the probe
            // before returning a transport; close() handles this safely.
            if (error instanceof UnauthorizedError && supportsOAuth(definition)) {
                return {
                    client,
                    transport: transport,
                    definition,
                    tools: [],
                    resources: [],
                    lastUsedAt: Date.now(),
                    inFlight: 0,
                    status: "needs-auth",
                };
            }
            throw error;
        }
    }
    createClient(serverName) {
        const client = new Client({ name: `spectral-mcp-${serverName}`, version: "1.0.0" }, this.samplingConfig ? { capabilities: { sampling: {} } } : undefined);
        if (this.samplingConfig) {
            registerSamplingHandler(client, { ...this.samplingConfig, serverName });
        }
        return client;
    }
    async createHttpTransport(definition, serverName) {
        const url = new URL(definition.url);
        // Build headers first (including any bearer token)
        const headers = resolveHeaders(definition.headers) ?? {};
        // For bearer auth, add the token to headers BEFORE creating requestInit
        if (definition.auth === "bearer") {
            const token = resolveBearerToken(definition);
            if (token) {
                headers["Authorization"] = `Bearer ${token}`;
            }
        }
        // Create request init with headers (Authorization now included for bearer auth)
        const requestInit = Object.keys(headers).length > 0 ? { headers } : undefined;
        // For OAuth servers, create an auth provider
        let authProvider;
        if (supportsOAuth(definition)) {
            // Extract OAuth config (handles both object and false cases)
            const oauthConfig = definition.oauth === false ? {} : {
                grantType: definition.oauth?.grantType,
                clientId: definition.oauth?.clientId,
                clientSecret: definition.oauth?.clientSecret,
                scope: definition.oauth?.scope,
                redirectUri: definition.oauth?.redirectUri,
            };
            authProvider = new McpOAuthProvider(serverName, definition.url, oauthConfig, {
                onRedirect: async (_authUrl) => {
                    // URL is captured by startAuth, no need to log
                },
            });
        }
        // Try StreamableHTTP first (modern MCP servers)
        const streamableTransport = new StreamableHTTPClientTransport(url, {
            requestInit,
            authProvider,
        });
        try {
            // Create a test client to verify the transport works
            const testClient = new Client({ name: "spectral-mcp-probe", version: "2.1.2" });
            await testClient.connect(streamableTransport);
            await testClient.close().catch(() => { });
            // Close probe transport before creating fresh one
            await streamableTransport.close().catch(() => { });
            // StreamableHTTP works - create fresh transport for actual use
            return new StreamableHTTPClientTransport(url, { requestInit, authProvider });
        }
        catch (error) {
            // StreamableHTTP failed, close and try SSE fallback
            await streamableTransport.close().catch(() => { });
            // If this was an UnauthorizedError, don't try SSE - the server needs auth
            if (error instanceof UnauthorizedError) {
                throw error;
            }
            // SSE is the legacy transport
            return new SSEClientTransport(url, { requestInit, authProvider });
        }
    }
    async fetchAllTools(client) {
        const allTools = [];
        let cursor;
        do {
            const result = await client.listTools(cursor ? { cursor } : undefined);
            allTools.push(...(result.tools ?? []));
            cursor = result.nextCursor;
        } while (cursor);
        return allTools;
    }
    async fetchAllResources(client) {
        try {
            const allResources = [];
            let cursor;
            do {
                const result = await client.listResources(cursor ? { cursor } : undefined);
                allResources.push(...(result.resources ?? []));
                cursor = result.nextCursor;
            } while (cursor);
            return allResources;
        }
        catch {
            // Server may not support resources
            return [];
        }
    }
    attachAdapterNotificationHandlers(serverName, client) {
        client.setNotificationHandler(serverStreamResultPatchNotificationSchema, (notification) => {
            const listener = this.uiStreamListeners.get(notification.params.streamToken);
            if (!listener)
                return;
            listener(serverName, notification.params);
        });
    }
    registerUiStreamListener(streamToken, listener) {
        this.uiStreamListeners.set(streamToken, listener);
    }
    removeUiStreamListener(streamToken) {
        this.uiStreamListeners.delete(streamToken);
    }
    async readResource(name, uri) {
        const connection = this.connections.get(name);
        if (!connection || connection.status !== "connected") {
            throw new Error(`Server "${name}" is not connected`);
        }
        try {
            this.touch(name);
            this.incrementInFlight(name);
            return await connection.client.readResource({ uri });
        }
        finally {
            this.decrementInFlight(name);
            this.touch(name);
        }
    }
    async close(name) {
        const connection = this.connections.get(name);
        if (!connection)
            return;
        // Delete from map BEFORE async cleanup to prevent a race where a
        // concurrent connect() creates a new connection that our deferred
        // delete() would then remove, orphaning the new server process.
        connection.status = "closed";
        this.connections.delete(name);
        await connection.client.close().catch(() => { });
        if (connection.transport) {
            await connection.transport.close().catch(() => { });
        }
    }
    async closeAll() {
        const names = [...this.connections.keys()];
        await Promise.all(names.map(name => this.close(name)));
    }
    getConnection(name) {
        return this.connections.get(name);
    }
    getAllConnections() {
        return new Map(this.connections);
    }
    touch(name) {
        const connection = this.connections.get(name);
        if (connection) {
            connection.lastUsedAt = Date.now();
        }
    }
    incrementInFlight(name) {
        const connection = this.connections.get(name);
        if (connection) {
            connection.inFlight = (connection.inFlight ?? 0) + 1;
        }
    }
    decrementInFlight(name) {
        const connection = this.connections.get(name);
        if (connection && connection.inFlight) {
            connection.inFlight--;
        }
    }
    isIdle(name, timeoutMs) {
        const connection = this.connections.get(name);
        if (!connection || connection.status !== "connected")
            return false;
        if (connection.inFlight > 0)
            return false;
        return (Date.now() - connection.lastUsedAt) > timeoutMs;
    }
}
/**
 * Resolve environment variables with interpolation.
 */
function resolveEnv(env) {
    // Copy process.env, filtering out undefined values
    const resolved = {};
    for (const [key, value] of Object.entries(process.env)) {
        if (value !== undefined) {
            resolved[key] = value;
        }
    }
    if (!env)
        return resolved;
    const overrides = interpolateEnvRecord(env);
    return overrides ? { ...resolved, ...overrides } : resolved;
}
/**
 * Resolve headers with environment variable interpolation.
 */
function resolveHeaders(headers) {
    return interpolateEnvRecord(headers);
}
