/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import { CancellationToken } from 'vs/base/common/cancellation';
import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
import { ExtHostLanguageModelsShape, MainContext, MainThreadLanguageModelsShape } from 'vs/workbench/api/common/extHost.protocol';
import * as typeConvert from 'vs/workbench/api/common/extHostTypeConverters';
import * as extHostTypes from 'vs/workbench/api/common/extHostTypes';
import type * as vscode from 'vscode';
import { Progress } from 'vs/platform/progress/common/progress';
import { IChatMessage, IChatResponseFragment, ILanguageModelChatMetadata } from 'vs/workbench/contrib/chat/common/languageModels';
import { ExtensionIdentifier, ExtensionIdentifierMap, ExtensionIdentifierSet, IExtensionDescription } from 'vs/platform/extensions/common/extensions';
import { AsyncIterableSource, Barrier } from 'vs/base/common/async';
import { Emitter, Event } from 'vs/base/common/event';
import { localize } from 'vs/nls';
import { INTERNAL_AUTH_PROVIDER_PREFIX } from 'vs/workbench/services/authentication/common/authentication';
import { CancellationError } from 'vs/base/common/errors';
import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
import { IExtHostRpcService } from 'vs/workbench/api/common/extHostRpcService';
import { IExtHostAuthentication } from 'vs/workbench/api/common/extHostAuthentication';
import { ILogService } from 'vs/platform/log/common/log';
import { Iterable } from 'vs/base/common/iterator';
import { checkProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';

export interface IExtHostLanguageModels extends ExtHostLanguageModels { }

export const IExtHostLanguageModels = createDecorator<IExtHostLanguageModels>('IExtHostLanguageModels');

type LanguageModelData = {
	readonly languageModelId: string;
	readonly extension: ExtensionIdentifier;
	readonly provider: vscode.ChatResponseProvider;
};

class LanguageModelResponseStream {

	readonly stream = new AsyncIterableSource<string>();

	constructor(
		readonly option: number,
		stream?: AsyncIterableSource<string>
	) {
		this.stream = stream ?? new AsyncIterableSource<string>();
	}
}

class LanguageModelResponse {

	readonly apiObject: vscode.LanguageModelChatResponse;

	private readonly _responseStreams = new Map<number, LanguageModelResponseStream>();
	private readonly _defaultStream = new AsyncIterableSource<string>();
	private _isDone: boolean = false;
	private _isStreaming: boolean = false;

	constructor() {

		const that = this;
		this.apiObject = {
			// result: promise,
			stream: that._defaultStream.asyncIterable,
			// streams: AsyncIterable<string>[] // FUTURE responses per N
		};
	}

	private * _streams() {
		if (this._responseStreams.size > 0) {
			for (const [, value] of this._responseStreams) {
				yield value.stream;
			}
		} else {
			yield this._defaultStream;
		}
	}

	handleFragment(fragment: IChatResponseFragment): void {
		if (this._isDone) {
			return;
		}
		this._isStreaming = true;
		let res = this._responseStreams.get(fragment.index);
		if (!res) {
			if (this._responseStreams.size === 0) {
				// the first response claims the default response
				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(): boolean {
		return this._isStreaming;
	}

	reject(err: Error): void {
		this._isDone = true;
		for (const stream of this._streams()) {
			stream.reject(err);
		}
	}

	resolve(): void {
		this._isDone = true;
		for (const stream of this._streams()) {
			stream.resolve();
		}
	}

}

export class ExtHostLanguageModels implements ExtHostLanguageModelsShape {

	declare _serviceBrand: undefined;

	private static _idPool = 1;

	private readonly _proxy: MainThreadLanguageModelsShape;
	private readonly _onDidChangeModelAccess = new Emitter<{ from: ExtensionIdentifier; to: ExtensionIdentifier }>();
	private readonly _onDidChangeProviders = new Emitter<vscode.LanguageModelChangeEvent>();
	readonly onDidChangeProviders = this._onDidChangeProviders.event;

	private readonly _languageModels = new Map<number, LanguageModelData>();
	private readonly _allLanguageModelData = new Map<string, ILanguageModelChatMetadata>(); // these are ALL models, not just the one in this EH
	private readonly _modelAccessList = new ExtensionIdentifierMap<ExtensionIdentifierSet>();
	private readonly _pendingRequest = new Map<number, { languageModelId: string; res: LanguageModelResponse }>();

	constructor(
		@IExtHostRpcService extHostRpc: IExtHostRpcService,
		@ILogService private readonly _logService: ILogService,
		@IExtHostAuthentication private readonly _extHostAuthentication: IExtHostAuthentication,
	) {
		this._proxy = extHostRpc.getProxy(MainContext.MainThreadLanguageModels);
	}

	dispose(): void {
		this._onDidChangeModelAccess.dispose();
		this._onDidChangeProviders.dispose();
	}

	registerLanguageModel(extension: IExtensionDescription, identifier: string, provider: vscode.ChatResponseProvider, metadata: vscode.ChatResponseProviderMetadata): IDisposable {

		const handle = ExtHostLanguageModels._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,
			name: metadata.name ?? '',
			version: metadata.version,
			tokens: metadata.tokens,
			auth
		});

		const responseReceivedListener = provider.onDidReceiveLanguageModelResponse2?.(({ extensionId, participant, tokenCount }) => {
			this._proxy.$whenLanguageModelChatRequestMade(identifier, new ExtensionIdentifier(extensionId), participant, tokenCount);
		});

		return toDisposable(() => {
			this._languageModels.delete(handle);
			this._proxy.$unregisterProvider(handle);
			responseReceivedListener?.dispose();
		});
	}

	async $provideLanguageModelResponse(handle: number, requestId: number, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise<any> {
		const data = this._languageModels.get(handle);
		if (!data) {
			return;
		}
		const progress = new Progress<vscode.ChatResponseFragment>(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 });
		});

		if (data.provider.provideLanguageModelResponse) {
			return data.provider.provideLanguageModelResponse(messages.map(typeConvert.LanguageModelChatMessage.to), options, ExtensionIdentifier.toKey(from), progress, token);
		} else {
			return data.provider.provideLanguageModelResponse2(messages.map(typeConvert.LanguageModelMessage.to), options, ExtensionIdentifier.toKey(from), progress, token);
		}
	}


	//#region --- token counting

	$provideTokenLength(handle: number, value: string, token: CancellationToken): Promise<number> {
		const data = this._languageModels.get(handle);
		if (!data) {
			return Promise.resolve(0);
		}
		return Promise.resolve(data.provider.provideTokenCount(value, token));
	}


	//#region --- making request

	$updateLanguageModels(data: { added?: ILanguageModelChatMetadata[] | undefined; removed?: string[] | undefined }): void {
		const added: string[] = [];
		const removed: string[] = [];
		if (data.added) {
			for (const metadata of data.added) {
				this._allLanguageModelData.set(metadata.identifier, metadata);
				added.push(metadata.identifier);
			}
		}
		if (data.removed) {
			for (const id of data.removed) {
				// clean up
				this._allLanguageModelData.delete(id);
				removed.push(id);

				// cancel pending requests for this model
				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)
		}));

		// TODO@jrieken@TylerLeonhardt - this is a temporary hack to populate the auth providers
		data.added?.forEach(this._fakeAuthPopulate, this);
	}

	getLanguageModelIds(): string[] {
		return Array.from(this._allLanguageModelData.keys());
	}

	$updateModelAccesslist(data: { from: ExtensionIdentifier; to: ExtensionIdentifier; enabled: boolean }[]): void {
		const updated = new Array<{ from: ExtensionIdentifier; to: ExtensionIdentifier }>();
		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: IExtensionDescription, languageModelId: string, messages: (vscode.LanguageModelChatMessage | vscode.LanguageModelChatMessage2)[], options: vscode.LanguageModelChatRequestOptions, token: CancellationToken) {

		const internalMessages: IChatMessage[] = this._convertMessages(extension, messages);

		const from = extension.identifier;
		const metadata = await this._proxy.$prepareChatAccess(from, languageModelId, options.justification);

		if (!metadata || !this._allLanguageModelData.has(languageModelId)) {
			throw extHostTypes.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 extHostTypes.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, internalMessages, options.modelOptions ?? {}, token);

		const barrier = new Barrier();

		const res = new LanguageModelResponse();
		this._pendingRequest.set(requestId, { languageModelId, res });

		let error: Error | undefined;

		requestPromise.catch(err => {
			if (barrier.isOpen()) {
				// we received an error while streaming. this means we need to reject the "stream"
				// because we have already returned the request object
				res.reject(err);
			} else {
				error = err;
			}
		}).finally(() => {
			this._pendingRequest.delete(requestId);
			res.resolve();
			barrier.open();
		});

		await barrier.wait();

		if (error) {
			if (error.name === extHostTypes.LanguageModelError.name) {
				throw error;
			}

			throw new extHostTypes.LanguageModelError(
				`Language model '${languageModelId}' errored, check cause for more details`,
				'Unknown',
				error
			);
		}

		return res.apiObject;
	}

	private _convertMessages(extension: IExtensionDescription, messages: (vscode.LanguageModelChatMessage2 | vscode.LanguageModelChatMessage)[]) {
		const internalMessages: IChatMessage[] = [];
		for (const message of messages) {
			if (message instanceof extHostTypes.LanguageModelChatMessage2) {
				if (message.role as number === extHostTypes.LanguageModelChatMessageRole.System) {
					checkProposedApiEnabled(extension, 'languageModelSystem');
				}
				internalMessages.push(typeConvert.LanguageModelChatMessage.from(message));
			} else {
				if (message instanceof extHostTypes.LanguageModelChatSystemMessage) {
					checkProposedApiEnabled(extension, 'languageModelSystem');
				}
				internalMessages.push(typeConvert.LanguageModelMessage.from(message));
			}
		}
		return internalMessages;
	}

	async $handleResponseFragment(requestId: number, chunk: IChatResponseFragment): Promise<void> {
		const data = this._pendingRequest.get(requestId);//.report(chunk);
		if (data) {
			data.res.handleFragment(chunk);
		}
	}

	// BIG HACK: Using AuthenticationProviders to check access to Language Models
	private async _getAuthAccess(from: IExtensionDescription, to: { identifier: ExtensionIdentifier; displayName: string }, justification: string | undefined, silent: boolean | undefined): Promise<boolean> {
		// This needs to be done in both MainThread & ExtHost ChatProvider
		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
				? localize('chatAccessWithJustification', "To allow access to the language models provided by {0}. Justification:\n\n{1}", to.displayName, justification)
				: localize('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) {
			// ignore
			return false;
		}
	}

	private _isUsingAuth(from: ExtensionIdentifier, toMetadata: ILanguageModelChatMetadata): toMetadata is ILanguageModelChatMetadata & { auth: NonNullable<ILanguageModelChatMetadata['auth']> } {
		// If the 'to' extension uses an auth check
		return !!toMetadata.auth
			// And we're asking from a different extension
			&& !ExtensionIdentifier.equals(toMetadata.extension, from);
	}

	private async _fakeAuthPopulate(metadata: ILanguageModelChatMetadata): Promise<void> {

		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);
			}
		}
	}

	async computeTokenLength(languageModelId: string, value: string | vscode.LanguageModelChatMessage, token: vscode.CancellationToken): Promise<number> {

		const data = this._allLanguageModelData.get(languageModelId);
		if (!data) {
			throw extHostTypes.LanguageModelError.NotFound(`Language model '${languageModelId}' is unknown.`);
		}

		const local = Iterable.find(this._languageModels.values(), candidate => candidate.languageModelId === languageModelId);
		if (local) {
			// stay inside the EH
			return local.provider.provideTokenCount(value, token);
		}

		return this._proxy.$countTokens(data.identifier, (typeof value === 'string' ? value : typeConvert.LanguageModelMessage.from(value)), token);
	}

	getLanguageModelInfo(languageModelId: string): vscode.LanguageModelInformation | undefined {
		const data = this._allLanguageModelData.get(languageModelId);
		if (!data) {
			return undefined;
		}

		return {
			id: data.identifier,
			name: data.name,
			version: data.version,
			tokens: data.tokens,
		};
	}

	private readonly _languageAccessInformationExtensions = new Set<Readonly<IExtensionDescription>>();

	createLanguageModelAccessInformation(from: Readonly<IExtensionDescription>): vscode.LanguageModelAccessInformation {

		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: string): boolean | undefined {

				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);
			}
		};
	}
}