
import { __decorate, __param } from '../../../../../../external/tslib/tslib.es6.js';
import { Sequencer, DeferredPromise, timeout, raceCancellationError } from '../../../base/common/async.js';
import { CancellationToken, CancellationTokenSource } from '../../../base/common/cancellation.js';
import { Disposable, DisposableMap, DisposableStore, toDisposable } from '../../../base/common/lifecycle.js';
import { SSEParser } from '../../../base/common/sseParser.js';
import { ExtensionIdentifier } from '../../../platform/extensions/common/extensions.js';
import { createDecorator } from '../../../platform/instantiation/common/instantiation.js';
import { LogLevel, canLog } from '../../../platform/log/common/log.js';
import { ILogService } from '../../../platform/log/common/log.service.js';
import { StorageScope } from '../../../platform/storage/common/storage.js';
import { McpServerLaunch, McpServerTransportType, McpConnectionState, extensionPrefixedIdentifier } from '../../contrib/mcp/common/mcpTypes.js';
import { MainContext } from './extHost.protocol.js';
import { IExtHostRpcService } from './extHostRpcService.js';
import { McpServerDefinition } from './extHostTypeConverters.js';
import { parseWWWAuthenticateHeader, getDefaultMetadataForUrl, isAuthorizationProtectedResourceMetadata, AUTH_SERVER_METADATA_DISCOVERY_PATH, OPENID_CONNECT_DISCOVERY_PATH, isAuthorizationServerMetadata } from '../../../base/common/oauth.js';
import { URI } from '../../../base/common/uri.js';
import { MCP } from '../../contrib/mcp/common/modelContextProtocol.js';
import { CancellationError } from '../../../base/common/errors.js';
import { ConfigurationTarget } from '../../../platform/configuration/common/configuration.js';
import { IExtHostInitDataService } from './extHostInitDataService.js';

const IExtHostMpcService = ( createDecorator('IExtHostMpcService'));
let ExtHostMcpService = class ExtHostMcpService extends Disposable {
    constructor(extHostRpc, _logService, _extHostInitData) {
        super();
        this._logService = _logService;
        this._extHostInitData = _extHostInitData;
        this._initialProviderPromises = ( new Set());
        this._sseEventSources = this._register(( new DisposableMap()));
        this._unresolvedMcpServers = ( new Map());
        this._proxy = ( extHostRpc.getProxy(MainContext.MainThreadMcp));
    }
    $startMcp(id, launch) {
        this._startMcp(id, McpServerLaunch.fromSerialized(launch));
    }
    _startMcp(id, launch) {
        if (launch.type === McpServerTransportType.HTTP) {
            this._sseEventSources.set(id, ( new McpHTTPHandle(id, launch, this._proxy, this._logService)));
            return;
        }
        throw ( new Error('not implemented'));
    }
    $stopMcp(id) {
        if (( this._sseEventSources.has(id))) {
            this._sseEventSources.deleteAndDispose(id);
            this._proxy.$onDidChangeState(id, { state: McpConnectionState.Kind.Stopped });
        }
    }
    $sendMessage(id, message) {
        this._sseEventSources.get(id)?.send(message);
    }
    async $waitForInitialCollectionProviders() {
        await Promise.all(this._initialProviderPromises);
    }
    async $resolveMcpLaunch(collectionId, label) {
        const rec = this._unresolvedMcpServers.get(collectionId);
        if (!rec) {
            return;
        }
        const server = rec.servers.find(s => s.label === label);
        if (!server) {
            return;
        }
        if (!rec.provider.resolveMcpServerDefinition) {
            return McpServerDefinition.from(server);
        }
        const resolved = await rec.provider.resolveMcpServerDefinition(server, CancellationToken.None);
        return resolved ? McpServerDefinition.from(resolved) : undefined;
    }
    registerMcpConfigurationProvider(extension, id, provider) {
        const store = ( new DisposableStore());
        const metadata = extension.contributes?.mcpServerDefinitionProviders?.find(m => m.id === id);
        if (!metadata) {
            throw ( new Error(
                `MCP configuration providers must be registered in the contributes.mcpServerDefinitionProviders array within your package.json, but "${id}" was not`
            ));
        }
        const mcp = {
            id: extensionPrefixedIdentifier(extension.identifier, id),
            isTrustedByDefault: true,
            label: metadata?.label ?? extension.displayName ?? extension.name,
            scope: StorageScope.WORKSPACE,
            canResolveLaunch: typeof provider.resolveMcpServerDefinition === 'function',
            extensionId: extension.identifier.value,
            configTarget: this._extHostInitData.remote.isRemote ? ConfigurationTarget.USER_REMOTE : ConfigurationTarget.USER,
        };
        const update = async () => {
            const list = await provider.provideMcpServerDefinitions(CancellationToken.None);
            this._unresolvedMcpServers.set(mcp.id, { servers: list ?? [], provider });
            const servers = [];
            for (const item of list ?? []) {
                let id = ExtensionIdentifier.toKey(extension.identifier) + '/' + item.label;
                if (( servers.some(s => s.id === id))) {
                    let i = 2;
                    while (( servers.some(s => s.id === id + i))) {
                        i++;
                    }
                    id = id + i;
                }
                servers.push({
                    id,
                    label: item.label,
                    cacheNonce: item.version || '$$NONE',
                    launch: McpServerDefinition.from(item),
                });
            }
            this._proxy.$upsertMcpCollection(mcp, servers);
        };
        store.add(toDisposable(() => {
            this._unresolvedMcpServers.delete(mcp.id);
            this._proxy.$deleteMcpCollection(mcp.id);
        }));
        if (provider.onDidChangeMcpServerDefinitions) {
            store.add(provider.onDidChangeMcpServerDefinitions(update));
        }
        if (provider.onDidChangeServerDefinitions) {
            store.add(provider.onDidChangeServerDefinitions(update));
        }
        if (provider.onDidChange) {
            store.add(provider.onDidChange(update));
        }
        const promise = ( new Promise(resolve => {
            setTimeout(() => update().finally(() => {
                this._initialProviderPromises.delete(promise);
                resolve();
            }), 0);
        }));
        this._initialProviderPromises.add(promise);
        return store;
    }
};
ExtHostMcpService = ( __decorate([
    ( __param(0, IExtHostRpcService)),
    ( __param(1, ILogService)),
    ( __param(2, IExtHostInitDataService))
], ExtHostMcpService));
var HttpMode;
(function (HttpMode) {
    HttpMode[HttpMode["Unknown"] = 0] = "Unknown";
    HttpMode[HttpMode["Http"] = 1] = "Http";
    HttpMode[HttpMode["SSE"] = 2] = "SSE";
})(HttpMode || (HttpMode = {}));
const MAX_FOLLOW_REDIRECTS = 5;
const REDIRECT_STATUS_CODES = [301, 302, 303, 307, 308];
class McpHTTPHandle extends Disposable {
    constructor(_id, _launch, _proxy, _logService) {
        super();
        this._id = _id;
        this._launch = _launch;
        this._proxy = _proxy;
        this._logService = _logService;
        this._requestSequencer = ( new Sequencer());
        this._postEndpoint = ( new DeferredPromise());
        this._mode = { value: HttpMode.Unknown };
        this._cts = ( new CancellationTokenSource());
        this._abortCtrl = ( new AbortController());
        this._register(toDisposable(() => {
            this._abortCtrl.abort();
            this._cts.dispose(true);
        }));
        this._proxy.$onDidChangeState(this._id, { state: McpConnectionState.Kind.Running });
    }
    async send(message) {
        try {
            if (this._mode.value === HttpMode.Unknown) {
                await this._requestSequencer.queue(() => this._send(message));
            }
            else {
                await this._send(message);
            }
        }
        catch (err) {
            const msg = `Error sending message to ${this._launch.uri}: ${String(err)}`;
            this._proxy.$onDidChangeState(this._id, { state: McpConnectionState.Kind.Error, message: msg });
        }
    }
    _send(message) {
        if (this._mode.value === HttpMode.SSE) {
            return this._sendLegacySSE(this._mode.endpoint, message);
        }
        else {
            return this._sendStreamableHttp(message, this._mode.value === HttpMode.Http ? this._mode.sessionId : undefined);
        }
    }
    async _sendStreamableHttp(message, sessionId) {
        const asBytes = ( new TextEncoder()).encode(message);
        const headers = {
            ...Object.fromEntries(this._launch.headers),
            'Content-Type': 'application/json',
            'Content-Length': String(asBytes.length),
            Accept: 'text/event-stream, application/json',
        };
        if (sessionId) {
            headers['Mcp-Session-Id'] = sessionId;
        }
        await this._addAuthHeader(headers);
        const res = await this._fetchWithAuthRetry(( this._launch.uri.toString(true)), {
            method: 'POST',
            headers,
            body: asBytes,
        }, headers);
        const wasUnknown = this._mode.value === HttpMode.Unknown;
        const nextSessionId = res.headers.get('Mcp-Session-Id');
        if (nextSessionId) {
            this._mode = { value: HttpMode.Http, sessionId: nextSessionId };
        }
        if (this._mode.value === HttpMode.Unknown &&
            res.status >= 400 && res.status < 500
            && res.status !== 401 && res.status !== 403) {
            this._log(LogLevel.Info, `${res.status} status sending message to ${this._launch.uri}, will attempt to fall back to legacy SSE`);
            this._sseFallbackWithMessage(message);
            return;
        }
        if (res.status >= 300) {
            const retryWithSessionId = this._mode.value === HttpMode.Http && !!this._mode.sessionId && (res.status === 400 || res.status === 404);
            this._proxy.$onDidChangeState(this._id, {
                state: McpConnectionState.Kind.Error,
                message: `${res.status} status sending message to ${this._launch.uri}: ${await this._getErrText(res)}` + (retryWithSessionId ? `; will retry with new session ID` : ''),
                shouldRetry: retryWithSessionId,
            });
            return;
        }
        if (this._mode.value === HttpMode.Unknown) {
            this._mode = { value: HttpMode.Http, sessionId: undefined };
        }
        if (wasUnknown) {
            this._attachStreamableBackchannel();
        }
        await this._handleSuccessfulStreamableHttp(res, message);
    }
    async _sseFallbackWithMessage(message) {
        const endpoint = await this._attachSSE();
        if (endpoint) {
            this._mode = { value: HttpMode.SSE, endpoint };
            await this._sendLegacySSE(endpoint, message);
        }
    }
    async _populateAuthMetadata(url, originalResponse) {
        let resourceMetadataChallenge;
        if (( originalResponse.headers.has('WWW-Authenticate'))) {
            const authHeader = originalResponse.headers.get('WWW-Authenticate');
            const { scheme, params } = parseWWWAuthenticateHeader(authHeader);
            if (scheme === 'Bearer' && params['resource_metadata']) {
                resourceMetadataChallenge = params['resource_metadata'];
            }
        }
        let serverMetadataUrl;
        let scopesSupported;
        let resource;
        if (resourceMetadataChallenge) {
            const resourceMetadata = await this._getResourceMetadata(resourceMetadataChallenge);
            if (( ( new URL(resourceMetadata.resource)).toString()) !== ( ( new URL(url)).toString())) {
                throw ( new Error(
                    `Protected Resource Metadata resource "${resourceMetadata.resource}" does not match MCP server resolved resource "${url}". The MCP server must follow OAuth spec https://datatracker.ietf.org/doc/html/rfc9728#PRConfigurationValidation`
                ));
            }
            serverMetadataUrl = resourceMetadata.authorization_servers?.[0];
            scopesSupported = resourceMetadata.scopes_supported;
            resource = resourceMetadata;
        }
        const baseUrl = ( new URL(originalResponse.url)).origin;
        let addtionalHeaders = {};
        if (!serverMetadataUrl) {
            serverMetadataUrl = baseUrl;
            addtionalHeaders = {
                ...Object.fromEntries(this._launch.headers)
            };
        }
        try {
            const serverMetadataResponse = await this._getAuthorizationServerMetadata(serverMetadataUrl, addtionalHeaders);
            this._authMetadata = {
                authorizationServer: ( URI.parse(serverMetadataUrl)),
                serverMetadata: serverMetadataResponse,
                resourceMetadata: resource
            };
            return;
        }
        catch (e) {
            this._log(LogLevel.Warning, `Error populating auth metadata: ${String(e)}`);
        }
        const defaultMetadata = getDefaultMetadataForUrl(( new URL(baseUrl)));
        defaultMetadata.scopes_supported = scopesSupported ?? defaultMetadata.scopes_supported ?? [];
        this._authMetadata = {
            authorizationServer: ( URI.parse(serverMetadataUrl)),
            serverMetadata: defaultMetadata,
            resourceMetadata: resource
        };
    }
    async _getResourceMetadata(resourceMetadata) {
        const resourceMetadataUrl = ( new URL(resourceMetadata));
        const mcpServerUrl = ( new URL(( this._launch.uri.toString(true))));
        let additionalHeaders = {};
        if (resourceMetadataUrl.origin === mcpServerUrl.origin) {
            additionalHeaders = {
                ...Object.fromEntries(this._launch.headers)
            };
        }
        const resourceMetadataResponse = await this._fetch(resourceMetadata, {
            method: 'GET',
            headers: {
                ...additionalHeaders,
                'Accept': 'application/json',
                'MCP-Protocol-Version': MCP.LATEST_PROTOCOL_VERSION
            }
        });
        if (resourceMetadataResponse.status !== 200) {
            throw ( new Error(
                `Failed to fetch resource metadata: ${resourceMetadataResponse.status} ${await this._getErrText(resourceMetadataResponse)}`
            ));
        }
        const body = await resourceMetadataResponse.json();
        if (isAuthorizationProtectedResourceMetadata(body)) {
            return body;
        }
        else {
            throw ( new Error(`Invalid resource metadata: ${JSON.stringify(body)}`));
        }
    }
    async _getAuthorizationServerMetadata(authorizationServer, addtionalHeaders) {
        const authorizationServerUrl = ( new URL(authorizationServer));
        const extraPath = authorizationServerUrl.pathname === '/' ? '' : authorizationServerUrl.pathname;
        const pathToFetch = ( ( new URL(AUTH_SERVER_METADATA_DISCOVERY_PATH, authorizationServer)).toString()) + extraPath;
        let authServerMetadataResponse = await this._fetch(pathToFetch, {
            method: 'GET',
            headers: {
                ...addtionalHeaders,
                'Accept': 'application/json',
                'MCP-Protocol-Version': MCP.LATEST_PROTOCOL_VERSION,
            }
        });
        if (authServerMetadataResponse.status !== 200) {
            const openidPathInsertionUrl = ( ( new URL(OPENID_CONNECT_DISCOVERY_PATH, authorizationServer)).toString()) + extraPath;
            authServerMetadataResponse = await this._fetch(openidPathInsertionUrl, {
                method: 'GET',
                headers: {
                    ...addtionalHeaders,
                    'Accept': 'application/json',
                    'MCP-Protocol-Version': MCP.LATEST_PROTOCOL_VERSION
                }
            });
            if (authServerMetadataResponse.status !== 200) {
                authServerMetadataResponse = await this._fetch(( URI.joinPath(( URI.parse(authorizationServer)), OPENID_CONNECT_DISCOVERY_PATH).toString(true)), {
                    method: 'GET',
                    headers: {
                        ...addtionalHeaders,
                        'Accept': 'application/json',
                        'MCP-Protocol-Version': MCP.LATEST_PROTOCOL_VERSION
                    }
                });
                if (authServerMetadataResponse.status !== 200) {
                    throw ( new Error(
                        `Failed to fetch authorization server metadata: ${authServerMetadataResponse.status} ${await this._getErrText(authServerMetadataResponse)}`
                    ));
                }
            }
        }
        const body = await authServerMetadataResponse.json();
        if (isAuthorizationServerMetadata(body)) {
            return body;
        }
        throw ( new Error(`Invalid authorization server metadata: ${JSON.stringify(body)}`));
    }
    async _handleSuccessfulStreamableHttp(res, message) {
        if (res.status === 202) {
            return;
        }
        switch (res.headers.get('Content-Type')?.toLowerCase()) {
            case 'text/event-stream': {
                const parser = ( new SSEParser(event => {
                    if (event.type === 'message') {
                        this._proxy.$onDidReceiveMessage(this._id, event.data);
                    }
                    else if (event.type === 'endpoint') {
                        this._log(LogLevel.Warning, `Received SSE endpoint from a POST to ${this._launch.uri}, will fall back to legacy SSE`);
                        this._sseFallbackWithMessage(message);
                        throw ( new CancellationError());
                    }
                }));
                try {
                    await this._doSSE(parser, res);
                }
                catch (err) {
                    this._log(LogLevel.Warning, `Error reading SSE stream: ${String(err)}`);
                }
                break;
            }
            case 'application/json':
                this._proxy.$onDidReceiveMessage(this._id, await res.text());
                break;
            default: {
                const responseBody = await res.text();
                if (isJSON(responseBody)) {
                    this._proxy.$onDidReceiveMessage(this._id, responseBody);
                }
                else {
                    this._log(LogLevel.Warning, `Unexpected ${res.status} response for request: ${responseBody}`);
                }
            }
        }
    }
    async _attachStreamableBackchannel() {
        let lastEventId;
        for (let retry = 0; !this._store.isDisposed; retry++) {
            await timeout(Math.min(retry * 1000, 30_000), this._cts.token);
            let res;
            try {
                const headers = {
                    ...Object.fromEntries(this._launch.headers),
                    'Accept': 'text/event-stream',
                };
                await this._addAuthHeader(headers);
                if (this._mode.value === HttpMode.Http && this._mode.sessionId !== undefined) {
                    headers['Mcp-Session-Id'] = this._mode.sessionId;
                }
                if (lastEventId) {
                    headers['Last-Event-ID'] = lastEventId;
                }
                res = await this._fetchWithAuthRetry(( this._launch.uri.toString(true)), {
                    method: 'GET',
                    headers,
                }, headers);
            }
            catch (e) {
                this._log(LogLevel.Info, `Error connecting to ${this._launch.uri} for async notifications, will retry`);
                continue;
            }
            if (res.status >= 400) {
                this._log(LogLevel.Debug, `${res.status} status connecting to ${this._launch.uri} for async notifications; they will be disabled: ${await this._getErrText(res)}`);
                return;
            }
            if (res.headers.get('content-type')?.toLowerCase().includes('text/event-stream')) {
                retry = 0;
            }
            const parser = ( new SSEParser(event => {
                if (event.type === 'message') {
                    this._proxy.$onDidReceiveMessage(this._id, event.data);
                }
                if (event.id) {
                    lastEventId = event.id;
                }
            }));
            try {
                await this._doSSE(parser, res);
            }
            catch (e) {
                this._log(LogLevel.Info, `Error reading from async stream, we will reconnect: ${e}`);
            }
        }
    }
    async _attachSSE() {
        const postEndpoint = ( new DeferredPromise());
        const headers = {
            ...Object.fromEntries(this._launch.headers),
            'Accept': 'text/event-stream',
        };
        await this._addAuthHeader(headers);
        let res;
        try {
            res = await this._fetchWithAuthRetry(( this._launch.uri.toString(true)), {
                method: 'GET',
                headers,
            }, headers);
            if (res.status >= 300) {
                this._proxy.$onDidChangeState(this._id, { state: McpConnectionState.Kind.Error, message: `${res.status} status connecting to ${this._launch.uri} as SSE: ${await this._getErrText(res)}` });
                return;
            }
        }
        catch (e) {
            this._proxy.$onDidChangeState(this._id, { state: McpConnectionState.Kind.Error, message: `Error connecting to ${this._launch.uri} as SSE: ${e}` });
            return;
        }
        const parser = ( new SSEParser(event => {
            if (event.type === 'message') {
                this._proxy.$onDidReceiveMessage(this._id, event.data);
            }
            else if (event.type === 'endpoint') {
                postEndpoint.complete(( ( new URL(event.data, ( this._launch.uri.toString(true)))).toString()));
            }
        }));
        this._register(toDisposable(() => postEndpoint.cancel()));
        this._doSSE(parser, res).catch(err => {
            this._proxy.$onDidChangeState(this._id, { state: McpConnectionState.Kind.Error, message: `Error reading SSE stream: ${String(err)}` });
        });
        return postEndpoint.p;
    }
    async _sendLegacySSE(url, message) {
        const asBytes = ( new TextEncoder()).encode(message);
        const headers = {
            ...Object.fromEntries(this._launch.headers),
            'Content-Type': 'application/json',
            'Content-Length': String(asBytes.length),
        };
        await this._addAuthHeader(headers);
        const res = await this._fetch(url, {
            method: 'POST',
            headers,
            body: asBytes,
        });
        if (res.status >= 300) {
            this._log(LogLevel.Warning, `${res.status} status sending message to ${this._postEndpoint}: ${await this._getErrText(res)}`);
        }
    }
    async _doSSE(parser, res) {
        if (!res.body) {
            return;
        }
        const reader = res.body.getReader();
        let chunk;
        do {
            try {
                chunk = await raceCancellationError(reader.read(), this._cts.token);
            }
            catch (err) {
                reader.cancel();
                if (this._store.isDisposed) {
                    return;
                }
                else {
                    throw err;
                }
            }
            if (chunk.value) {
                parser.feed(chunk.value);
            }
        } while (!chunk.done);
    }
    async _addAuthHeader(headers) {
        if (this._authMetadata) {
            try {
                const token = await this._proxy.$getTokenFromServerMetadata(this._id, this._authMetadata.authorizationServer, this._authMetadata.serverMetadata, this._authMetadata.resourceMetadata);
                if (token) {
                    headers['Authorization'] = `Bearer ${token}`;
                }
            }
            catch (e) {
                this._log(LogLevel.Warning, `Error getting token from server metadata: ${String(e)}`);
            }
        }
        return headers;
    }
    _log(level, message) {
        if (!this._store.isDisposed) {
            this._proxy.$onDidPublishLog(this._id, level, message);
        }
    }
    async _getErrText(res) {
        try {
            return await res.text();
        }
        catch {
            return res.statusText;
        }
    }
    async _fetchWithAuthRetry(url, init, headers) {
        const doFetch = () => this._fetch(url, init);
        let res = await doFetch();
        if (res.status === 401) {
            if (!this._authMetadata) {
                await this._populateAuthMetadata(url, res);
                await this._addAuthHeader(headers);
                if (headers['Authorization']) {
                    init.headers = headers;
                    res = await doFetch();
                }
            }
        }
        return res;
    }
    async _fetch(url, init) {
        if (canLog(this._logService.getLevel(), LogLevel.Trace)) {
            const traceObj = { ...init, headers: { ...init.headers } };
            if (traceObj.body) {
                traceObj.body = ( new TextDecoder()).decode(traceObj.body);
            }
            if (traceObj.headers?.Authorization) {
                traceObj.headers.Authorization = '***';
            }
            this._log(LogLevel.Trace, `Fetching ${url} with options: ${JSON.stringify(traceObj)}`);
        }
        let currentUrl = url;
        let response;
        for (let redirectCount = 0; redirectCount < MAX_FOLLOW_REDIRECTS; redirectCount++) {
            response = await fetch(currentUrl, {
                ...init,
                signal: this._abortCtrl.signal,
                redirect: 'manual'
            });
            if (!REDIRECT_STATUS_CODES.includes(response.status)) {
                break;
            }
            const location = response.headers.get('location');
            if (!location) {
                break;
            }
            const nextUrl = ( ( new URL(location, currentUrl)).toString());
            this._log(LogLevel.Trace, `Redirect (${response.status}) from ${currentUrl} to ${nextUrl}`);
            currentUrl = nextUrl;
            if (response.status === 303 || ((response.status === 301 || response.status === 302) && init.method === 'POST')) {
                init.method = 'GET';
                delete init.body;
            }
        }
        if (canLog(this._logService.getLevel(), LogLevel.Trace)) {
            const headers = {};
            response.headers.forEach((value, key) => { headers[key] = value; });
            this._log(LogLevel.Trace, `Fetched ${currentUrl}: ${JSON.stringify({
                status: response.status,
                headers: headers,
            })}`);
        }
        return response;
    }
}
function isJSON(str) {
    try {
        JSON.parse(str);
        return true;
    }
    catch (e) {
        return false;
    }
}

export { ExtHostMcpService, IExtHostMpcService };
