import { __decorate, __param } from '../../../../../../external/tslib/tslib.es6.js';
import { toDisposable } from '../../../base/common/lifecycle.js';
import { MainContext } from './extHost.protocol.js';
import { LanguageModelMessage } from './extHostTypeConverters.js';
import { LanguageModelError } from './extHostTypes.js';
import { Progress } from '../../../platform/progress/common/progress.js';
import { ExtensionIdentifierMap, ExtensionIdentifier, ExtensionIdentifierSet } from '../../../platform/extensions/common/extensions.js';
import { Barrier, AsyncIterableSource } from '../../../base/common/async.js';
import { Emitter, Event } from '../../../base/common/event.js';
import { localizeWithPath } from '../../../nls.js';
import { INTERNAL_AUTH_PROVIDER_PREFIX } from '../../services/authentication/common/authentication.js';
import { CancellationError } from '../../../base/common/errors.js';
import { createDecorator } from '../../../platform/instantiation/common/instantiation.js';
import { IExtHostRpcService } from './extHostRpcService.js';
import { IExtHostAuthentication } from './extHostAuthentication.js';
import { ILogService } from '../../../platform/log/common/log.js';
var ExtHostLanguageModels_1;
const IExtHostLanguageModels = ( createDecorator('IExtHostLanguageModels'));
class LanguageModelResponseStream {
    constructor(option, stream) {
        this.option = option;
        this.stream = ( new AsyncIterableSource());
        this.stream = stream ?? ( new AsyncIterableSource());
    }
}
class LanguageModelResponse {
    constructor() {
        this._responseStreams = ( new Map());
        this._defaultStream = ( new AsyncIterableSource());
        this._isDone = false;
        this._isStreaming = false;
        const that = this;
        this.apiObject = {
            stream: that._defaultStream.asyncIterable,
        };
    }
    *_streams() {
        if (this._responseStreams.size > 0) {
            for (const [, value] of this._responseStreams) {
                yield value.stream;
            }
        }
        else {
            yield this._defaultStream;
        }
    }
    handleFragment(fragment) {
        if (this._isDone) {
            return;
        }
        this._isStreaming = true;
        let res = this._responseStreams.get(fragment.index);
        if (!res) {
            if (this._responseStreams.size === 0) {
                res = ( new LanguageModelResponseStream(fragment.index, this._defaultStream));
            }
            else {
                res = ( new LanguageModelResponseStream(fragment.index));
            }
            this._responseStreams.set(fragment.index, res);
        }
        res.stream.emitOne(fragment.part);
    }
    get isStreaming() {
        return this._isStreaming;
    }
    reject(err) {
        this._isDone = true;
        for (const stream of this._streams()) {
            stream.reject(err);
        }
    }
    resolve() {
        this._isDone = true;
        for (const stream of this._streams()) {
            stream.resolve();
        }
    }
}
let ExtHostLanguageModels = class ExtHostLanguageModels {
    static { ExtHostLanguageModels_1 = this; }
    static { this._idPool = 1; }
    constructor(extHostRpc, _logService, _extHostAuthentication) {
        this._logService = _logService;
        this._extHostAuthentication = _extHostAuthentication;
        this._onDidChangeModelAccess = ( new Emitter());
        this._onDidChangeProviders = ( new Emitter());
        this.onDidChangeProviders = this._onDidChangeProviders.event;
        this._languageModels = ( new Map());
        this._allLanguageModelData = ( new Map());
        this._modelAccessList = ( new ExtensionIdentifierMap());
        this._pendingRequest = ( new Map());
        this._languageAccessInformationExtensions = ( new Set());
        this._proxy = ( extHostRpc.getProxy(MainContext.MainThreadLanguageModels));
    }
    dispose() {
        this._onDidChangeModelAccess.dispose();
        this._onDidChangeProviders.dispose();
    }
    registerLanguageModel(extension, identifier, provider, metadata) {
        const handle = ExtHostLanguageModels_1._idPool++;
        this._languageModels.set(handle, { extension: extension.identifier, provider, languageModelId: identifier });
        let auth;
        if (metadata.auth) {
            auth = {
                providerLabel: extension.displayName || extension.name,
                accountLabel: typeof metadata.auth === 'object' ? metadata.auth.label : undefined
            };
        }
        this._proxy.$registerLanguageModelProvider(handle, identifier, {
            extension: extension.identifier,
            identifier: identifier,
            model: metadata.name ?? '',
            auth
        });
        return toDisposable(() => {
            this._languageModels.delete(handle);
            this._proxy.$unregisterProvider(handle);
        });
    }
    async $provideLanguageModelResponse(handle, requestId, from, messages, options, token) {
        const data = this._languageModels.get(handle);
        if (!data) {
            return;
        }
        const progress = ( new Progress(async (fragment) => {
            if (token.isCancellationRequested) {
                this._logService.warn(`[CHAT](${data.extension.value}) CANNOT send progress because the REQUEST IS CANCELLED`);
                return;
            }
            this._proxy.$handleProgressChunk(requestId, { index: fragment.index, part: fragment.part });
        }));
        return data.provider.provideLanguageModelResponse2(( messages.map(LanguageModelMessage.to)), options, ExtensionIdentifier.toKey(from), progress, token);
    }
    $updateLanguageModels(data) {
        const added = [];
        const removed = [];
        if (data.added) {
            for (const metadata of data.added) {
                this._allLanguageModelData.set(metadata.identifier, metadata);
                added.push(metadata.model);
            }
        }
        if (data.removed) {
            for (const id of data.removed) {
                this._allLanguageModelData.delete(id);
                removed.push(id);
                for (const [key, value] of this._pendingRequest) {
                    if (value.languageModelId === id) {
                        value.res.reject(( new CancellationError()));
                        this._pendingRequest.delete(key);
                    }
                }
            }
        }
        this._onDidChangeProviders.fire(( Object.freeze({
            added: ( Object.freeze(added)),
            removed: ( Object.freeze(removed))
        })));
        data.added?.forEach(this._fakeAuthPopulate, this);
    }
    getLanguageModelIds() {
        return Array.from(( this._allLanguageModelData.keys()));
    }
    $updateModelAccesslist(data) {
        const updated = ( new Array());
        for (const { from, to, enabled } of data) {
            const set = this._modelAccessList.get(from) ?? ( new ExtensionIdentifierSet());
            const oldValue = ( set.has(to));
            if (oldValue !== enabled) {
                if (enabled) {
                    set.add(to);
                }
                else {
                    set.delete(to);
                }
                this._modelAccessList.set(from, set);
                const newItem = { from, to };
                updated.push(newItem);
                this._onDidChangeModelAccess.fire(newItem);
            }
        }
    }
    async sendChatRequest(extension, languageModelId, messages, options, token) {
        const from = extension.identifier;
        const metadata = await this._proxy.$prepareChatAccess(from, languageModelId, options.justification);
        if (!metadata || !( this._allLanguageModelData.has(languageModelId))) {
            throw LanguageModelError.NotFound(`Language model '${languageModelId}' is unknown.`);
        }
        if (this._isUsingAuth(from, metadata)) {
            const success = await this._getAuthAccess(extension, { identifier: metadata.extension, displayName: metadata.auth.providerLabel }, options.justification, options.silent);
            if (!success || !this._modelAccessList.get(from)?.has(metadata.extension)) {
                throw LanguageModelError.NoPermissions(`Language model '${languageModelId}' cannot be used by '${from.value}'.`);
            }
        }
        const requestId = (Math.random() * 1e6) | 0;
        const requestPromise = this._proxy.$fetchResponse(from, languageModelId, requestId, ( messages.map(LanguageModelMessage.from)), options.modelOptions ?? {}, token);
        const barrier = ( new Barrier());
        const res = ( new LanguageModelResponse());
        this._pendingRequest.set(requestId, { languageModelId, res });
        let error;
        requestPromise.catch(err => {
            if (barrier.isOpen()) {
                res.reject(err);
            }
            else {
                error = err;
            }
        }).finally(() => {
            this._pendingRequest.delete(requestId);
            res.resolve();
            barrier.open();
        });
        await barrier.wait();
        if (error) {
            throw new LanguageModelError(`Language model '${languageModelId}' errored, check cause for more details`, 'Unknown', error);
        }
        return res.apiObject;
    }
    async $handleResponseFragment(requestId, chunk) {
        const data = this._pendingRequest.get(requestId);
        if (data) {
            data.res.handleFragment(chunk);
        }
    }
    async _getAuthAccess(from, to, justification, silent) {
        const providerId = INTERNAL_AUTH_PROVIDER_PREFIX + to.identifier.value;
        const session = await this._extHostAuthentication.getSession(from, providerId, [], { silent: true });
        if (session) {
            this.$updateModelAccesslist([{ from: from.identifier, to: to.identifier, enabled: true }]);
            return true;
        }
        if (silent) {
            return false;
        }
        try {
            const detail = justification
                ? ( localizeWithPath(
                'vs/workbench/api/common/extHostLanguageModels',
                'chatAccessWithJustification',
                "To allow access to the language models provided by {0}. Justification:\n\n{1}",
                to.displayName,
                justification
            ))
                : ( localizeWithPath(
                'vs/workbench/api/common/extHostLanguageModels',
                'chatAccess',
                "To allow access to the language models provided by {0}",
                to.displayName
            ));
            await this._extHostAuthentication.getSession(from, providerId, [], { forceNewSession: { detail } });
            this.$updateModelAccesslist([{ from: from.identifier, to: to.identifier, enabled: true }]);
            return true;
        }
        catch (err) {
            return false;
        }
    }
    _isUsingAuth(from, toMetadata) {
        return !!toMetadata.auth
            && !ExtensionIdentifier.equals(toMetadata.extension, from);
    }
    async _fakeAuthPopulate(metadata) {
        for (const from of this._languageAccessInformationExtensions) {
            try {
                await this._getAuthAccess(from, { identifier: metadata.extension, displayName: '' }, undefined, true);
            }
            catch (err) {
                this._logService.error('Fake Auth request failed');
                this._logService.error(err);
            }
        }
    }
    createLanguageModelAccessInformation(from) {
        this._languageAccessInformationExtensions.add(from);
        const that = this;
        const _onDidChangeAccess = Event.signal(Event.filter(this._onDidChangeModelAccess.event, e => ExtensionIdentifier.equals(e.from, from.identifier)));
        const _onDidAddRemove = Event.signal(this._onDidChangeProviders.event);
        return {
            get onDidChange() {
                return Event.any(_onDidChangeAccess, _onDidAddRemove);
            },
            canSendRequest(languageModelId) {
                const data = that._allLanguageModelData.get(languageModelId);
                if (!data) {
                    return undefined;
                }
                if (!that._isUsingAuth(from.identifier, data)) {
                    return true;
                }
                const list = that._modelAccessList.get(from.identifier);
                if (!list) {
                    return undefined;
                }
                return ( list.has(data.extension));
            }
        };
    }
};
ExtHostLanguageModels = ExtHostLanguageModels_1 = ( __decorate([
    ( __param(0, IExtHostRpcService)),
    ( __param(1, ILogService)),
    ( __param(2, IExtHostAuthentication))
], ExtHostLanguageModels));
export { ExtHostLanguageModels, IExtHostLanguageModels };
