import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { StdioServerParameters } from '@modelcontextprotocol/sdk/client/stdio.js';
import { StreamableHTTPClientTransportOptions } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
import { Root } from '@modelcontextprotocol/sdk/types.js';
import { DynamicActionProviderAction, Genkit, MultipartToolAction, ToolAction, DynamicResourceAction, PromptGenerateOptions, ExecutablePrompt } from 'genkit';

/**
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

interface McpServerRef {
    client: Client;
    transport: Transport;
    error?: string;
}
interface McpServerControls {
    /** when true, the server will be stopped and its registered components will
     * not appear in lists/plugins/etc */
    disabled?: boolean;
    roots?: Root[];
}
type McpStdioServerConfig = StdioServerParameters & {
    url?: never;
    transport?: never;
};
type McpStreamableHttpConfig = {
    url: string;
    command?: never;
    transport?: never;
} & Omit<StreamableHTTPClientTransportOptions, 'sessionId'>;
type McpTransportServerConfig = {
    transport: Transport;
    command?: never;
    url?: never;
};
/**
 * Configuration for an individual MCP server. The interface should be familiar
 * and compatible with existing tool configurations e.g. Cursor or Claude
 * Desktop.
 *
 * In addition to stdio servers, remote servers are supported via URL and
 * custom/arbitary transports are supported as well.
 */
type McpServerConfig = (McpStdioServerConfig | McpStreamableHttpConfig | McpTransportServerConfig) & McpServerControls;
/**
 * Configuration options for an individual `GenkitMcpClient` instance.
 * This defines how the client connects to a single MCP server and how it behaves.
 */
type McpClientOptions<M extends boolean = false> = {
    /** Client name to advertise to the server. */
    name: string;
    /** Name for the server, defaults to the server's advertised name. */
    serverName?: string;
    /**
     * An optional version number for this client. This is primarily for logging
     * and identification purposes. Defaults to '1.0.0'.
     */
    version?: string;
    /**
     * If true, tool responses from the MCP server will be returned in their raw
     * MCP format. Otherwise (default), they are processed and potentially
     * simplified for better compatibility with Genkit's typical data structures.
     */
    rawToolResponses?: boolean;
    /** If true, tools will be registered as multipart tool.v2 actions. */
    multipart?: M;
    /** The server configuration to connect. */
    mcpServer: McpServerConfig;
    /** Manually supply a session id for HTTP streaming clients if desired. */
    sessionId?: string;
};
type McpClientOptionsWithCache<M extends boolean = false> = McpClientOptions<M> & {
    cacheTtlMillis?: number;
};
/**
 * Represents a client connection to a single MCP (Model Context Protocol) server.
 * It handles the lifecycle of the connection (connect, disconnect, disable, re-enable, reconnect)
 * and provides methods to fetch tools from the connected server.
 *
 * An instance of `GenkitMcpClient` is typically managed by a `GenkitMcpHost`
 * when dealing with multiple MCP server connections.
 */
declare class GenkitMcpClient<Multipart extends boolean = false> {
    _server?: McpServerRef;
    private _dynamicActionProvider;
    sessionId?: string;
    readonly name: string;
    readonly suppliedServerName?: string;
    private version;
    private serverConfig;
    private rawToolResponses?;
    private multipart?;
    private disabled;
    private roots?;
    private _readyListeners;
    private _ready;
    constructor(options: McpClientOptions<Multipart>);
    set dynamicActionProvider(dap: DynamicActionProviderAction);
    _invalidateDapCache(): void;
    get serverName(): string;
    updateRoots(roots: Root[]): Promise<void>;
    /**
     * Sets up a connection based on a provided map of server configurations.
     * - Reconnects existing servers if their configuration appears to have
     *   changed (implicitly handled by `connectServer`).
     * - Sets the client's ready state once all connection attempts are complete.
     * @param mcpServers A record mapping server names to their configurations.
     */
    private _initializeConnection;
    /**
     * Returns a Promise that resolves when the client has attempted to connect
     * to all configured servers, or rejects if a critical error occurs during
     * the initial connection phase.
     */
    ready(): Promise<void>;
    /**
     * Connects to a single MCP server defined by the provided configuration.
     * @param config The configuration object for the server.
     */
    private _connect;
    /**
     * Disconnects the MCP server and removes its registration
     * from this client instance.
     */
    _disconnect(): Promise<void>;
    /**
     * Disables a server. Closes the underlying transport and server's configuration. Does nothing if the server is
     * already disabled.
     */
    disable(): Promise<void>;
    /**
     * Whether this client-server connection is enabled or not.
     */
    isEnabled(): boolean;
    /**
     * Enables a server connection, including previously disabled ones. Does nothing if the
     * server is not disabled.
     */
    enable(): Promise<void>;
    /**
     * Closes and then restarts the transport connection for the specified server.
     * Useful for attempting to recover from connection issues without full
     * reconfiguration.
     */
    restart(): Promise<void>;
    /**
     * Fetches all tools available through this client, if the server
     * configuration is not disabled.
     */
    getActiveTools(ai: Genkit): Promise<(Multipart extends true ? MultipartToolAction : ToolAction)[]>;
    /**
     * Fetches all resources available through this client, if the server
     * configuration is not disabled.
     */
    getActiveResources(ai: Genkit): Promise<DynamicResourceAction[]>;
    /**
     * Fetches all active prompts available through this client, if the server
     * configuration supports prompts.
     * @param ai The Genkit instance.
     * @param options Optional prompt generation options.
     * @returns A promise that resolves to an array of ExecutablePrompt.
     */
    getActivePrompts(ai: Genkit, options?: PromptGenerateOptions): Promise<ExecutablePrompt[]>;
    /**
     * Get the specified prompt as an `ExecutablePrompt` available through this
     * client. If no such prompt is found, return undefined.
     */
    getPrompt(ai: Genkit, promptName: string, opts?: PromptGenerateOptions): Promise<ExecutablePrompt | undefined>;
    /** Returns the underlying MCP SDK client if one has been initialized. */
    get mcpClient(): Client | undefined;
}

export { GenkitMcpClient, type McpClientOptions, type McpClientOptionsWithCache, type McpServerConfig, type McpServerControls, type McpStdioServerConfig, type McpStreamableHttpConfig, type McpTransportServerConfig };