
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 } from '../../../platform/log/common/log.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';

const IExtHostMpcService = ( createDecorator('IExtHostMpcService'));
let ExtHostMcpService = class ExtHostMcpService extends Disposable {
    constructor(extHostRpc) {
        super();
        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)));
            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?.modelContextServerCollections?.find(m => m.id === id);
        if (!metadata) {
            throw ( new Error(
                `MCP configuration providers must be registered in the contributes.modelContextServerCollections 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,
        };
        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,
                    launch: McpServerDefinition.from(item)
                });
            }
            this._proxy.$upsertMcpCollection(mcp, servers);
        };
        store.add(toDisposable(() => {
            this._unresolvedMcpServers.delete(mcp.id);
            this._proxy.$deleteMcpCollection(mcp.id);
        }));
        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))
], ExtHostMcpService));
var HttpMode;
(function (HttpMode) {
    HttpMode[HttpMode["Unknown"] = 0] = "Unknown";
    HttpMode[HttpMode["Http"] = 1] = "Http";
    HttpMode[HttpMode["SSE"] = 2] = "SSE";
})(HttpMode || (HttpMode = {}));
class McpHTTPHandle extends Disposable {
    constructor(_id, _launch, _proxy) {
        super();
        this._id = _id;
        this._launch = _launch;
        this._proxy = _proxy;
        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 {
            await this._requestSequencer.queue(() => {
                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);
                }
            });
        }
        catch (err) {
            const msg = `Error sending message to ${this._launch.uri}: ${String(err)}`;
            this._proxy.$onDidChangeState(this._id, { state: McpConnectionState.Kind.Error, message: msg });
        }
    }
    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;
        }
        const res = await fetch(( this._launch.uri.toString(true)), {
            method: 'POST',
            signal: this._abortCtrl.signal,
            headers,
            body: asBytes,
        });
        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) {
            this._log(LogLevel.Info, `${res.status} status sending message to ${this._launch.uri}, will attempt to fall back to legacy SSE`);
            const endpoint = await this._attachSSE();
            if (endpoint) {
                this._mode = { value: HttpMode.SSE, endpoint };
                await this._sendLegacySSE(endpoint, message);
            }
            return;
        }
        if (res.status >= 300) {
            const retryWithSessionId = this._mode.value === HttpMode.Http && !!this._mode.sessionId;
            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();
        }
        this._handleSuccessfulStreamableHttp(res);
    }
    async _handleSuccessfulStreamableHttp(res) {
        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);
                    }
                }));
                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',
                };
                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 fetch(( this._launch.uri.toString(true)), {
                    method: 'GET',
                    signal: this._abortCtrl.signal,
                    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;
            }
            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());
        let res;
        try {
            res = await fetch(( this._launch.uri.toString(true)), {
                method: 'GET',
                signal: this._abortCtrl.signal,
                headers: {
                    ...Object.fromEntries(this._launch.headers),
                    'Accept': 'text/event-stream',
                },
            });
            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 res = await fetch(url, {
            method: 'POST',
            signal: this._abortCtrl.signal,
            headers: {
                ...Object.fromEntries(this._launch.headers),
                'Content-Type': 'application/json',
                'Content-Length': String(asBytes.length),
            },
            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);
    }
    _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;
        }
    }
}
function isJSON(str) {
    try {
        JSON.parse(str);
        return true;
    }
    catch (e) {
        return false;
    }
}

export { ExtHostMcpService, IExtHostMpcService };
