import { Server } from "./index.js";
import { zodToJsonSchema } from "zod-to-json-schema";
import { z, ZodType, } from "zod";
import { McpError, ErrorCode, ListResourceTemplatesRequestSchema, ReadResourceRequestSchema, ListToolsRequestSchema, CallToolRequestSchema, ListResourcesRequestSchema, ListPromptsRequestSchema, GetPromptRequestSchema, CompleteRequestSchema, } from "../types.js";
import { Completable } from "./completable.js";
import { UriTemplate } from "../shared/uriTemplate.js";
/**
 * High-level MCP server that provides a simpler API for working with resources, tools, and prompts.
 * For advanced usage (like sending notifications or setting custom request handlers), use the underlying
 * Server instance available via the `server` property.
 */
export class McpServer {
    constructor(serverInfo, options) {
        this._registeredResources = {};
        this._registeredResourceTemplates = {};
        this._registeredTools = {};
        this._registeredPrompts = {};
        this._toolHandlersInitialized = false;
        this._completionHandlerInitialized = false;
        this._resourceHandlersInitialized = false;
        this._promptHandlersInitialized = false;
        this.server = new Server(serverInfo, options);
    }
    /**
     * Attaches to the given transport, starts it, and starts listening for messages.
     *
     * The `server` object assumes ownership of the Transport, replacing any callbacks that have already been set, and expects that it is the only user of the Transport instance going forward.
     */
    async connect(transport) {
        return await this.server.connect(transport);
    }
    /**
     * Closes the connection.
     */
    async close() {
        await this.server.close();
    }
    setToolRequestHandlers() {
        if (this._toolHandlersInitialized) {
            return;
        }
        this.server.assertCanSetRequestHandler(ListToolsRequestSchema.shape.method.value);
        this.server.assertCanSetRequestHandler(CallToolRequestSchema.shape.method.value);
        this.server.registerCapabilities({
            tools: {
                listChanged: true
            }
        });
        this.server.setRequestHandler(ListToolsRequestSchema, () => ({
            tools: Object.entries(this._registeredTools).filter(([, tool]) => tool.enabled).map(([name, tool]) => {
                return {
                    name,
                    description: tool.description,
                    inputSchema: tool.inputSchema
                        ? zodToJsonSchema(tool.inputSchema, {
                            strictUnions: true,
                        })
                        : EMPTY_OBJECT_JSON_SCHEMA,
                    annotations: tool.annotations,
                };
            }),
        }));
        this.server.setRequestHandler(CallToolRequestSchema, async (request, extra) => {
            const tool = this._registeredTools[request.params.name];
            if (!tool) {
                throw new McpError(ErrorCode.InvalidParams, `Tool ${request.params.name} not found`);
            }
            if (!tool.enabled) {
                throw new McpError(ErrorCode.InvalidParams, `Tool ${request.params.name} disabled`);
            }
            if (tool.inputSchema) {
                const parseResult = await tool.inputSchema.safeParseAsync(request.params.arguments);
                if (!parseResult.success) {
                    throw new McpError(ErrorCode.InvalidParams, `Invalid arguments for tool ${request.params.name}: ${parseResult.error.message}`);
                }
                const args = parseResult.data;
                const cb = tool.callback;
                try {
                    return await Promise.resolve(cb(args, extra));
                }
                catch (error) {
                    return {
                        content: [
                            {
                                type: "text",
                                text: error instanceof Error ? error.message : String(error),
                            },
                        ],
                        isError: true,
                    };
                }
            }
            else {
                const cb = tool.callback;
                try {
                    return await Promise.resolve(cb(extra));
                }
                catch (error) {
                    return {
                        content: [
                            {
                                type: "text",
                                text: error instanceof Error ? error.message : String(error),
                            },
                        ],
                        isError: true,
                    };
                }
            }
        });
        this._toolHandlersInitialized = true;
    }
    setCompletionRequestHandler() {
        if (this._completionHandlerInitialized) {
            return;
        }
        this.server.assertCanSetRequestHandler(CompleteRequestSchema.shape.method.value);
        this.server.setRequestHandler(CompleteRequestSchema, async (request) => {
            switch (request.params.ref.type) {
                case "ref/prompt":
                    return this.handlePromptCompletion(request, request.params.ref);
                case "ref/resource":
                    return this.handleResourceCompletion(request, request.params.ref);
                default:
                    throw new McpError(ErrorCode.InvalidParams, `Invalid completion reference: ${request.params.ref}`);
            }
        });
        this._completionHandlerInitialized = true;
    }
    async handlePromptCompletion(request, ref) {
        const prompt = this._registeredPrompts[ref.name];
        if (!prompt) {
            throw new McpError(ErrorCode.InvalidParams, `Prompt ${ref.name} not found`);
        }
        if (!prompt.enabled) {
            throw new McpError(ErrorCode.InvalidParams, `Prompt ${ref.name} disabled`);
        }
        if (!prompt.argsSchema) {
            return EMPTY_COMPLETION_RESULT;
        }
        const field = prompt.argsSchema.shape[request.params.argument.name];
        if (!(field instanceof Completable)) {
            return EMPTY_COMPLETION_RESULT;
        }
        const def = field._def;
        const suggestions = await def.complete(request.params.argument.value);
        return createCompletionResult(suggestions);
    }
    async handleResourceCompletion(request, ref) {
        const template = Object.values(this._registeredResourceTemplates).find((t) => t.resourceTemplate.uriTemplate.toString() === ref.uri);
        if (!template) {
            if (this._registeredResources[ref.uri]) {
                // Attempting to autocomplete a fixed resource URI is not an error in the spec (but probably should be).
                return EMPTY_COMPLETION_RESULT;
            }
            throw new McpError(ErrorCode.InvalidParams, `Resource template ${request.params.ref.uri} not found`);
        }
        const completer = template.resourceTemplate.completeCallback(request.params.argument.name);
        if (!completer) {
            return EMPTY_COMPLETION_RESULT;
        }
        const suggestions = await completer(request.params.argument.value);
        return createCompletionResult(suggestions);
    }
    setResourceRequestHandlers() {
        if (this._resourceHandlersInitialized) {
            return;
        }
        this.server.assertCanSetRequestHandler(ListResourcesRequestSchema.shape.method.value);
        this.server.assertCanSetRequestHandler(ListResourceTemplatesRequestSchema.shape.method.value);
        this.server.assertCanSetRequestHandler(ReadResourceRequestSchema.shape.method.value);
        this.server.registerCapabilities({
            resources: {
                listChanged: true
            }
        });
        this.server.setRequestHandler(ListResourcesRequestSchema, async (request, extra) => {
            const resources = Object.entries(this._registeredResources).filter(([_, resource]) => resource.enabled).map(([uri, resource]) => ({
                uri,
                name: resource.name,
                ...resource.metadata,
            }));
            const templateResources = [];
            for (const template of Object.values(this._registeredResourceTemplates)) {
                if (!template.resourceTemplate.listCallback) {
                    continue;
                }
                const result = await template.resourceTemplate.listCallback(extra);
                for (const resource of result.resources) {
                    templateResources.push({
                        ...resource,
                        ...template.metadata,
                    });
                }
            }
            return { resources: [...resources, ...templateResources] };
        });
        this.server.setRequestHandler(ListResourceTemplatesRequestSchema, async () => {
            const resourceTemplates = Object.entries(this._registeredResourceTemplates).map(([name, template]) => ({
                name,
                uriTemplate: template.resourceTemplate.uriTemplate.toString(),
                ...template.metadata,
            }));
            return { resourceTemplates };
        });
        this.server.setRequestHandler(ReadResourceRequestSchema, async (request, extra) => {
            const uri = new URL(request.params.uri);
            // First check for exact resource match
            const resource = this._registeredResources[uri.toString()];
            if (resource) {
                if (!resource.enabled) {
                    throw new McpError(ErrorCode.InvalidParams, `Resource ${uri} disabled`);
                }
                return resource.readCallback(uri, extra);
            }
            // Then check templates
            for (const template of Object.values(this._registeredResourceTemplates)) {
                const variables = template.resourceTemplate.uriTemplate.match(uri.toString());
                if (variables) {
                    return template.readCallback(uri, variables, extra);
                }
            }
            throw new McpError(ErrorCode.InvalidParams, `Resource ${uri} not found`);
        });
        this.setCompletionRequestHandler();
        this._resourceHandlersInitialized = true;
    }
    setPromptRequestHandlers() {
        if (this._promptHandlersInitialized) {
            return;
        }
        this.server.assertCanSetRequestHandler(ListPromptsRequestSchema.shape.method.value);
        this.server.assertCanSetRequestHandler(GetPromptRequestSchema.shape.method.value);
        this.server.registerCapabilities({
            prompts: {
                listChanged: true
            }
        });
        this.server.setRequestHandler(ListPromptsRequestSchema, () => ({
            prompts: Object.entries(this._registeredPrompts).filter(([, prompt]) => prompt.enabled).map(([name, prompt]) => {
                return {
                    name,
                    description: prompt.description,
                    arguments: prompt.argsSchema
                        ? promptArgumentsFromSchema(prompt.argsSchema)
                        : undefined,
                };
            }),
        }));
        this.server.setRequestHandler(GetPromptRequestSchema, async (request, extra) => {
            const prompt = this._registeredPrompts[request.params.name];
            if (!prompt) {
                throw new McpError(ErrorCode.InvalidParams, `Prompt ${request.params.name} not found`);
            }
            if (!prompt.enabled) {
                throw new McpError(ErrorCode.InvalidParams, `Prompt ${request.params.name} disabled`);
            }
            if (prompt.argsSchema) {
                const parseResult = await prompt.argsSchema.safeParseAsync(request.params.arguments);
                if (!parseResult.success) {
                    throw new McpError(ErrorCode.InvalidParams, `Invalid arguments for prompt ${request.params.name}: ${parseResult.error.message}`);
                }
                const args = parseResult.data;
                const cb = prompt.callback;
                return await Promise.resolve(cb(args, extra));
            }
            else {
                const cb = prompt.callback;
                return await Promise.resolve(cb(extra));
            }
        });
        this.setCompletionRequestHandler();
        this._promptHandlersInitialized = true;
    }
    resource(name, uriOrTemplate, ...rest) {
        let metadata;
        if (typeof rest[0] === "object") {
            metadata = rest.shift();
        }
        const readCallback = rest[0];
        if (typeof uriOrTemplate === "string") {
            if (this._registeredResources[uriOrTemplate]) {
                throw new Error(`Resource ${uriOrTemplate} is already registered`);
            }
            const registeredResource = {
                name,
                metadata,
                readCallback: readCallback,
                enabled: true,
                disable: () => registeredResource.update({ enabled: false }),
                enable: () => registeredResource.update({ enabled: true }),
                remove: () => registeredResource.update({ uri: null }),
                update: (updates) => {
                    if (typeof updates.uri !== "undefined" && updates.uri !== uriOrTemplate) {
                        delete this._registeredResources[uriOrTemplate];
                        if (updates.uri)
                            this._registeredResources[updates.uri] = registeredResource;
                    }
                    if (typeof updates.name !== "undefined")
                        registeredResource.name = updates.name;
                    if (typeof updates.metadata !== "undefined")
                        registeredResource.metadata = updates.metadata;
                    if (typeof updates.callback !== "undefined")
                        registeredResource.readCallback = updates.callback;
                    if (typeof updates.enabled !== "undefined")
                        registeredResource.enabled = updates.enabled;
                    this.sendResourceListChanged();
                },
            };
            this._registeredResources[uriOrTemplate] = registeredResource;
            this.setResourceRequestHandlers();
            this.sendResourceListChanged();
            return registeredResource;
        }
        else {
            if (this._registeredResourceTemplates[name]) {
                throw new Error(`Resource template ${name} is already registered`);
            }
            const registeredResourceTemplate = {
                resourceTemplate: uriOrTemplate,
                metadata,
                readCallback: readCallback,
                enabled: true,
                disable: () => registeredResourceTemplate.update({ enabled: false }),
                enable: () => registeredResourceTemplate.update({ enabled: true }),
                remove: () => registeredResourceTemplate.update({ name: null }),
                update: (updates) => {
                    if (typeof updates.name !== "undefined" && updates.name !== name) {
                        delete this._registeredResourceTemplates[name];
                        if (updates.name)
                            this._registeredResourceTemplates[updates.name] = registeredResourceTemplate;
                    }
                    if (typeof updates.template !== "undefined")
                        registeredResourceTemplate.resourceTemplate = updates.template;
                    if (typeof updates.metadata !== "undefined")
                        registeredResourceTemplate.metadata = updates.metadata;
                    if (typeof updates.callback !== "undefined")
                        registeredResourceTemplate.readCallback = updates.callback;
                    if (typeof updates.enabled !== "undefined")
                        registeredResourceTemplate.enabled = updates.enabled;
                    this.sendResourceListChanged();
                },
            };
            this._registeredResourceTemplates[name] = registeredResourceTemplate;
            this.setResourceRequestHandlers();
            this.sendResourceListChanged();
            return registeredResourceTemplate;
        }
    }
    tool(name, ...rest) {
        if (this._registeredTools[name]) {
            throw new Error(`Tool ${name} is already registered`);
        }
        // Helper to check if an object is a Zod schema (ZodRawShape)
        const isZodRawShape = (obj) => {
            if (typeof obj !== "object" || obj === null)
                return false;
            // Check that at least one property is a ZodType instance
            return Object.values(obj).some(v => v instanceof ZodType);
        };
        let description;
        if (typeof rest[0] === "string") {
            description = rest.shift();
        }
        let paramsSchema;
        let annotations;
        // Handle the different overload combinations
        if (rest.length > 1) {
            // We have at least two more args before the callback
            const firstArg = rest[0];
            if (isZodRawShape(firstArg)) {
                // We have a params schema as the first arg
                paramsSchema = rest.shift();
                // Check if the next arg is potentially annotations  
                if (rest.length > 1 && typeof rest[0] === "object" && rest[0] !== null && !(isZodRawShape(rest[0]))) {
                    // Case: tool(name, paramsSchema, annotations, cb)
                    // Or: tool(name, description, paramsSchema, annotations, cb)
                    annotations = rest.shift();
                }
            }
            else if (typeof firstArg === "object" && firstArg !== null) {
                // Not a ZodRawShape, so must be annotations in this position
                // Case: tool(name, annotations, cb)
                // Or: tool(name, description, annotations, cb)
                annotations = rest.shift();
            }
        }
        const cb = rest[0];
        const registeredTool = {
            description,
            inputSchema: paramsSchema === undefined ? undefined : z.object(paramsSchema),
            annotations,
            callback: cb,
            enabled: true,
            disable: () => registeredTool.update({ enabled: false }),
            enable: () => registeredTool.update({ enabled: true }),
            remove: () => registeredTool.update({ name: null }),
            update: (updates) => {
                if (typeof updates.name !== "undefined" && updates.name !== name) {
                    delete this._registeredTools[name];
                    if (updates.name)
                        this._registeredTools[updates.name] = registeredTool;
                }
                if (typeof updates.description !== "undefined")
                    registeredTool.description = updates.description;
                if (typeof updates.paramsSchema !== "undefined")
                    registeredTool.inputSchema = z.object(updates.paramsSchema);
                if (typeof updates.callback !== "undefined")
                    registeredTool.callback = updates.callback;
                if (typeof updates.annotations !== "undefined")
                    registeredTool.annotations = updates.annotations;
                if (typeof updates.enabled !== "undefined")
                    registeredTool.enabled = updates.enabled;
                this.sendToolListChanged();
            },
        };
        this._registeredTools[name] = registeredTool;
        this.setToolRequestHandlers();
        this.sendToolListChanged();
        return registeredTool;
    }
    prompt(name, ...rest) {
        if (this._registeredPrompts[name]) {
            throw new Error(`Prompt ${name} is already registered`);
        }
        let description;
        if (typeof rest[0] === "string") {
            description = rest.shift();
        }
        let argsSchema;
        if (rest.length > 1) {
            argsSchema = rest.shift();
        }
        const cb = rest[0];
        const registeredPrompt = {
            description,
            argsSchema: argsSchema === undefined ? undefined : z.object(argsSchema),
            callback: cb,
            enabled: true,
            disable: () => registeredPrompt.update({ enabled: false }),
            enable: () => registeredPrompt.update({ enabled: true }),
            remove: () => registeredPrompt.update({ name: null }),
            update: (updates) => {
                if (typeof updates.name !== "undefined" && updates.name !== name) {
                    delete this._registeredPrompts[name];
                    if (updates.name)
                        this._registeredPrompts[updates.name] = registeredPrompt;
                }
                if (typeof updates.description !== "undefined")
                    registeredPrompt.description = updates.description;
                if (typeof updates.argsSchema !== "undefined")
                    registeredPrompt.argsSchema = z.object(updates.argsSchema);
                if (typeof updates.callback !== "undefined")
                    registeredPrompt.callback = updates.callback;
                if (typeof updates.enabled !== "undefined")
                    registeredPrompt.enabled = updates.enabled;
                this.sendPromptListChanged();
            },
        };
        this._registeredPrompts[name] = registeredPrompt;
        this.setPromptRequestHandlers();
        this.sendPromptListChanged();
        return registeredPrompt;
    }
    /**
     * Checks if the server is connected to a transport.
     * @returns True if the server is connected
     */
    isConnected() {
        return this.server.transport !== undefined;
    }
    /**
     * Sends a resource list changed event to the client, if connected.
     */
    sendResourceListChanged() {
        if (this.isConnected()) {
            this.server.sendResourceListChanged();
        }
    }
    /**
     * Sends a tool list changed event to the client, if connected.
     */
    sendToolListChanged() {
        if (this.isConnected()) {
            this.server.sendToolListChanged();
        }
    }
    /**
     * Sends a prompt list changed event to the client, if connected.
     */
    sendPromptListChanged() {
        if (this.isConnected()) {
            this.server.sendPromptListChanged();
        }
    }
}
/**
 * A resource template combines a URI pattern with optional functionality to enumerate
 * all resources matching that pattern.
 */
export class ResourceTemplate {
    constructor(uriTemplate, _callbacks) {
        this._callbacks = _callbacks;
        this._uriTemplate =
            typeof uriTemplate === "string"
                ? new UriTemplate(uriTemplate)
                : uriTemplate;
    }
    /**
     * Gets the URI template pattern.
     */
    get uriTemplate() {
        return this._uriTemplate;
    }
    /**
     * Gets the list callback, if one was provided.
     */
    get listCallback() {
        return this._callbacks.list;
    }
    /**
     * Gets the callback for completing a specific URI template variable, if one was provided.
     */
    completeCallback(variable) {
        var _a;
        return (_a = this._callbacks.complete) === null || _a === void 0 ? void 0 : _a[variable];
    }
}
const EMPTY_OBJECT_JSON_SCHEMA = {
    type: "object",
};
function promptArgumentsFromSchema(schema) {
    return Object.entries(schema.shape).map(([name, field]) => ({
        name,
        description: field.description,
        required: !field.isOptional(),
    }));
}
function createCompletionResult(suggestions) {
    return {
        completion: {
            values: suggestions.slice(0, 100),
            total: suggestions.length,
            hasMore: suggestions.length > 100,
        },
    };
}
const EMPTY_COMPLETION_RESULT = {
    completion: {
        values: [],
        hasMore: false,
    },
};
//# sourceMappingURL=mcp.js.map