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

import product from '../../../../platform/product/common/product.js';
import { Barrier } from '../../../../base/common/async.js';
import { CancellationToken, CancellationTokenSource } from '../../../../base/common/cancellation.js';
import { Emitter, Event } from '../../../../base/common/event.js';
import { Lazy } from '../../../../base/common/lazy.js';
import { Disposable, MutableDisposable } from '../../../../base/common/lifecycle.js';
import { IRequestContext } from '../../../../base/parts/request/common/request.js';
import { localize } from '../../../../nls.js';
import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
import { ChatAIDisabledSettingId } from '../../../../platform/chat/common/chatSettings.js';
import { IContextKey, IContextKeyService, RawContextKey } from '../../../../platform/contextkey/common/contextkey.js';
import { IDialogService } from '../../../../platform/dialogs/common/dialogs.js';
import { createDecorator, IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
import { ILogService, LogLevel } from '../../../../platform/log/common/log.js';
import { IProductService } from '../../../../platform/product/common/productService.js';
import { asText, IRequestService } from '../../../../platform/request/common/request.js';
import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
import { ITelemetryService, TelemetryLevel } from '../../../../platform/telemetry/common/telemetry.js';
import { AuthenticationSession, IAuthenticationService } from '../../authentication/common/authentication.js';
import { IOpenerService } from '../../../../platform/opener/common/opener.js';
import { URI } from '../../../../base/common/uri.js';
import Severity from '../../../../base/common/severity.js';
import { IWorkbenchEnvironmentService } from '../../environment/common/environmentService.js';
import { isWeb } from '../../../../base/common/platform.js';
import { ILifecycleService } from '../../lifecycle/common/lifecycle.js';
import { Mutable } from '../../../../base/common/types.js';
import { InstantiationType, registerSingleton } from '../../../../platform/instantiation/common/extensions.js';
import { IObservable, observableFromEvent } from '../../../../base/common/observable.js';
import { IDefaultAccountService } from '../../../../platform/defaultAccount/common/defaultAccount.js';
import { IDefaultAccount, IEntitlementsData } from '../../../../base/common/defaultAccount.js';
import { isInternalAccount } from '../../../../platform/assignment/common/assignment.js';

export namespace ChatEntitlementContextKeys {

	export const Setup = {
		hidden: new RawContextKey<boolean>('chatSetupHidden', false, true), 		// True when chat setup is explicitly hidden.
		installed: new RawContextKey<boolean>('chatSetupInstalled', false, true),  	// True when the chat extension is installed and enabled.
		disabled: new RawContextKey<boolean>('chatSetupDisabled', false, true),  	// True when the chat extension is disabled due to any other reason than workspace trust.
		disabledInWorkspace: new RawContextKey<boolean>('chatSetupDisabledInWorkspace', false, true),	// True when chat is disabled at the workspace level via settings.
		untrusted: new RawContextKey<boolean>('chatSetupUntrusted', false, true),  	// True when the chat extension is disabled due to workspace trust.
		later: new RawContextKey<boolean>('chatSetupLater', false, true),  			// True when the user wants to finish setup later.
		registered: new RawContextKey<boolean>('chatSetupRegistered', false, true), // True when the user has registered as Free or Pro user.
		completed: new RawContextKey<boolean>('chatSetupCompleted', false, true)	// True when the user has completed the setup flow, regardless of the outcome.
	};

	export const Entitlement = {
		signedOut: new RawContextKey<boolean>('chatEntitlementSignedOut', false, true), 				// True when user is signed out.
		canSignUp: new RawContextKey<boolean>('chatPlanCanSignUp', false, true), 						// True when user can sign up to be a chat free user.

		planFree: new RawContextKey<boolean>('chatPlanFree', false, true),								// True when user is a chat free user.
		planPro: new RawContextKey<boolean>('chatPlanPro', false, true),								// True when user is a chat pro user.
		planEdu: new RawContextKey<boolean>('chatPlanEdu', false, true),								// True when user is a chat edu user.
		planProPlus: new RawContextKey<boolean>('chatPlanProPlus', false, true), 						// True when user is a chat pro plus user.
		planMax: new RawContextKey<boolean>('chatPlanMax', false, true), 								// True when user is a chat max user.
		planBusiness: new RawContextKey<boolean>('chatPlanBusiness', false, true), 						// True when user is a chat business user.
		planEnterprise: new RawContextKey<boolean>('chatPlanEnterprise', false, true), 					// True when user is a chat enterprise user.

		organisations: new RawContextKey<string[]>('chatEntitlementOrganisations', undefined, true), 	// The organizations the user belongs to.
		internal: new RawContextKey<boolean>('chatEntitlementInternal', false, true), 					// True when the user is GitHub/Microsoft staff or belongs to an internal organisation.
		sku: new RawContextKey<string>('chatEntitlementSku', undefined, true), 							// The SKU of the user.
	};

	export const chatQuotaExceeded = new RawContextKey<boolean>('chatQuotaExceeded', false, true);
	export const completionsQuotaExceeded = new RawContextKey<boolean>('completionsQuotaExceeded', false, true);

	export const chatAnonymous = new RawContextKey<boolean>('chatAnonymous', false, true);

	export const clientByokEnabled = new RawContextKey<boolean>('github.copilot.clientByokEnabled', true, true);

	export const hasByokModels = new RawContextKey<boolean>('github.copilot.hasByokModels', false, true);
}

export const IChatEntitlementService = createDecorator<IChatEntitlementService>('chatEntitlementService');

export enum ChatEntitlement {
	/** Signed out */
	Unknown = 1,
	/** Signed in but not yet resolved */
	Unresolved = 2,
	/** Signed in and entitled to Free */
	Available = 3,
	/** Signed in but not entitled to Free */
	Unavailable = 4,
	/** Signed-up to Free */
	Free = 5,
	/** Signed-up to EDU */
	EDU = 10,
	/** Signed-up to Pro */
	Pro = 6,
	/** Signed-up to Pro Plus */
	ProPlus = 7,
	/** Signed-up to Business */
	Business = 8,
	/** Signed-up to Enterprise */
	Enterprise = 9,
	/** Signed-up to Max */
	Max = 11,
}

export interface IChatSentiment {

	/**
	 * Whether the user has completed the setup flow or not, regardless of the outcome
	 */
	completed?: boolean;

	/**
	 * User has Chat installed.
	 */
	installed?: boolean;

	/**
	 * User signals no intent in using Chat.
	 *
	 * Note: in contrast to `disabled`, this should not only disable
	 * Chat but also hide all of its UI.
	 */
	hidden?: boolean;

	/**
	 * User signals intent to disable Chat.
	 *
	 * Note: in contrast to `hidden`, this should not hide
	 * Chat but but disable its functionality.
	 */
	disabled?: boolean;

	/**
	 * Chat is disabled at the workspace level
	 *
	 * Note: in contrast to `hidden` (which hides all UI globally),
	 * this only disables Chat in the current workspace while
	 * keeping its UI visible so the user can re-enable it.
	 */
	disabledInWorkspace?: boolean;

	/**
	 * Chat is disabled due to missing workspace trust.
	 *
	 * Note: even though this disables Chat, we want to treat it
	 * different from the `disabled` state that is by explicit
	 * user choice.
	 */
	untrusted?: boolean;

	/**
	 * User signals intent to use Chat later.
	 */
	later?: boolean;

	/**
	 * User has registered as Free or Pro user.
	 */
	registered?: boolean;
}

/**
 * The inputs needed to decide whether Chat still requires the user to run setup
 * (sign in / sign up / trust / enable) before it can service a request.
 */
export interface IChatSetupRequirement {
	/** Whether the setup flow has been completed (any outcome). */
	readonly completed: boolean;
	/** Whether the chat extension is disabled for a reason other than trust. */
	readonly disabled: boolean;
	/** Whether the chat extension is disabled because the workspace is untrusted. */
	readonly untrusted: boolean;
	/** The user's last known or resolved entitlement. */
	readonly entitlement: ChatEntitlement;
	/** Whether anonymous (signed-out) Chat access is enabled. */
	readonly anonymous: boolean;
	/** Whether BYOK models are available. */
	readonly hasByokModels: boolean;
}

/**
 * Returns whether Chat requires setup before it can service a request.
 * The model picker uses a narrower condition that only surfaces interactive setup.
 */
export function chatRequiresSetup(context: IChatSetupRequirement): boolean {
	return (
		(!context.completed && !context.hasByokModels) ||			// Setup not completed (unless BYOK models are available)
		context.disabled ||											// Extension disabled: run setup to enable
		context.untrusted ||										// Workspace untrusted: run setup to ask for trust
		context.entitlement === ChatEntitlement.Available ||		// Entitlement available: run setup to sign up
		(
			context.entitlement === ChatEntitlement.Unknown &&		// Entitlement unknown: run setup to sign in / sign up
			!context.anonymous &&									// unless anonymous access is enabled
			!context.hasByokModels									// unless BYOK models are available
		)
	);
}

export interface IChatEntitlementService {

	_serviceBrand: undefined;

	readonly onDidChangeEntitlement: Event<void>;

	readonly entitlement: ChatEntitlement;
	readonly entitlementObs: IObservable<ChatEntitlement>;

	readonly clientByokEnabled: boolean;
	readonly hasByokModels: boolean;

	readonly organisations: string[] | undefined;

	/** True for GitHub/Microsoft staff or members of an internal organisation. */
	readonly isInternal: boolean;
	readonly sku: string | undefined;
	readonly copilotTrackingId: string | undefined;

	readonly onDidChangeQuotaExceeded: Event<void>;
	readonly onDidChangeQuotaRemaining: Event<void>;
	readonly onDidChangeUsageBasedBilling: Event<void>;

	readonly quotas: IQuotas;

	readonly onDidChangeSentiment: Event<void>;

	readonly sentiment: IChatSentiment;
	readonly sentimentObs: IObservable<IChatSentiment>;

	// TODO@bpasero eventually this will become enabled by default
	// and in that case we only need to check on entitlements change
	// between `unknown` and any other entitlement.
	readonly onDidChangeAnonymous: Event<void>;
	readonly anonymous: boolean;
	readonly anonymousObs: IObservable<boolean>;

	acceptQuotas(quotas: IQuotas): void;

	/**
	 * Clear all quota state.
	 */
	clearQuotas(): void;

	markAnonymousRateLimited(): void;

	/**
	 * Mark the chat setup flow as completed.
	 */
	markSetupCompleted(): void;

	/**
	 * Force the hidden state on or off, overriding the normal entitlement logic.
	 * Used by the account policy gate to hide all AI features when the gate is
	 * active and unsatisfied.
	 */
	setForceHidden(hidden: boolean): void;

	update(token: CancellationToken): Promise<void>;
}

//#region Helper Functions

/**
 * Checks the chat entitlements to see if the user falls into the paid category
 * @param chatEntitlement The chat entitlement to check
 * @returns Whether or not they are a paid user
 */
export function isProUser(chatEntitlement: ChatEntitlement): boolean {
	return chatEntitlement === ChatEntitlement.EDU ||
		chatEntitlement === ChatEntitlement.Pro ||
		chatEntitlement === ChatEntitlement.ProPlus ||
		chatEntitlement === ChatEntitlement.Max ||
		chatEntitlement === ChatEntitlement.Business ||
		chatEntitlement === ChatEntitlement.Enterprise;
}

/**
 * Gets the full plan name for the given chat entitlement
 * @param chatEntitlement The chat entitlement to get the plan name for
 * @returns The localized full plan name (e.g., "Copilot Pro", "Copilot Free")
 */
export function getChatPlanName(chatEntitlement: ChatEntitlement): string {
	switch (chatEntitlement) {
		case ChatEntitlement.EDU:
			return localize('plan.eduName', 'Copilot Student');
		case ChatEntitlement.Pro:
			return localize('plan.proName', 'Copilot Pro');
		case ChatEntitlement.ProPlus:
			return localize('plan.proPlusName', 'Copilot Pro+');
		case ChatEntitlement.Max:
			return localize('plan.maxName', 'Copilot Max');
		case ChatEntitlement.Business:
			return localize('plan.businessName', 'Copilot Business');
		case ChatEntitlement.Enterprise:
			return localize('plan.enterpriseName', 'Copilot Enterprise');
		default:
			return localize('plan.freeName', 'Copilot Free');
	}
}

//#region Service Implementation

const defaultChatAgent = {
	upgradePlanUrl: product.defaultChatAgent?.upgradePlanUrl ?? '',
	providerUriSetting: product.defaultChatAgent?.providerUriSetting ?? '',
	entitlementSignupLimitedUrl: product.defaultChatAgent?.entitlementSignupLimitedUrl ?? '',
	chatQuotaExceededContext: product.defaultChatAgent?.chatQuotaExceededContext ?? '',
	completionsQuotaExceededContext: product.defaultChatAgent?.completionsQuotaExceededContext ?? ''
};

interface IChatQuotasAccessor {
	clearQuotas(): void;
	acceptQuotas(quotas: IQuotas): void;
}

const CHAT_ALLOW_ANONYMOUS_CONFIGURATION_KEY = 'chat.allowAnonymousAccess';

function isAnonymous(configurationService: IConfigurationService, entitlement: ChatEntitlement, sentiment: IChatSentiment): boolean {
	if (configurationService.getValue(CHAT_ALLOW_ANONYMOUS_CONFIGURATION_KEY) !== true) {
		return false; // only enabled behind an experimental setting
	}

	if (entitlement !== ChatEntitlement.Unknown) {
		return false; // only consider signed out users
	}

	if (sentiment.hidden || sentiment.disabledInWorkspace) {
		return false; // only consider enabled scenarios
	}

	return true;
}

type ChatEntitlementClassification = {
	owner: 'bpasero';
	comment: 'Provides insight into chat entitlements.';
	chatHidden: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether chat is hidden or not.' };
	chatEntitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The current chat entitlement of the user.' };
	chatAnonymous: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user is anonymously using chat.' };
	chatRegistered: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user is registered for chat.' };
	chatDisabled: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether chat is disabled or not.' };
};
type ChatEntitlementEvent = {
	chatHidden: boolean;
	chatEntitlement: ChatEntitlement;
	chatAnonymous: boolean;
	chatRegistered: boolean;
	chatDisabled: boolean;
};

function logChatEntitlements(state: IChatEntitlementContextState, configurationService: IConfigurationService, telemetryService: ITelemetryService): void {
	telemetryService.publicLog2<ChatEntitlementEvent, ChatEntitlementClassification>('chatEntitlements', {
		chatHidden: Boolean(state.hidden),
		chatDisabled: Boolean(state.disabled),
		chatEntitlement: state.entitlement,
		chatRegistered: Boolean(state.registered),
		chatAnonymous: isAnonymous(configurationService, state.entitlement, state)
	});
}

type ChatAdditionalSpendConfigurationClassification = {
	owner: 'pwang347';
	comment: 'Tracks when a user enables or disables additional spend.';
	enabled: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether additional spend is now enabled or disabled.' };
	entitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The current chat entitlement of the user.' };
};
type ChatAdditionalSpendConfigurationEvent = {
	enabled: boolean;
	entitlement: ChatEntitlement;
};

type ChatAdditionalSpendActiveClassification = {
	owner: 'pwang347';
	comment: 'Tracks when a user enters additional spend (included quota exhausted while additional spend is enabled).';
	entitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The current chat entitlement of the user.' };
	additionalUsageCount: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The number of additional spend interactions used so far.' };
};
type ChatAdditionalSpendActiveEvent = {
	entitlement: ChatEntitlement;
	additionalUsageCount: number;
};

export class ChatEntitlementService extends Disposable implements IChatEntitlementService {

	declare _serviceBrand: undefined;

	private static readonly CACHED_UBB_STORAGE_KEY = 'chat.usageBasedBilling';

	readonly context: Lazy<ChatEntitlementContext> | undefined;
	readonly requests: Lazy<ChatEntitlementRequests> | undefined;

	constructor(
		@IInstantiationService instantiationService: IInstantiationService,
		@IProductService productService: IProductService,
		@IWorkbenchEnvironmentService environmentService: IWorkbenchEnvironmentService,
		@IContextKeyService private readonly contextKeyService: IContextKeyService,
		@IConfigurationService private readonly configurationService: IConfigurationService,
		@ITelemetryService private readonly telemetryService: ITelemetryService,
		@ILogService private readonly logService: ILogService,
		@IStorageService private readonly storageService: IStorageService,
	) {
		super();

		const cachedUBB = this.storageService.getBoolean(ChatEntitlementService.CACHED_UBB_STORAGE_KEY, StorageScope.PROFILE);
		this._quotas = cachedUBB !== undefined ? { usageBasedBilling: cachedUBB } : {};

		this.chatQuotaExceededContextKey = ChatEntitlementContextKeys.chatQuotaExceeded.bindTo(this.contextKeyService);
		this.completionsQuotaExceededContextKey = ChatEntitlementContextKeys.completionsQuotaExceeded.bindTo(this.contextKeyService);

		this.anonymousContextKey = ChatEntitlementContextKeys.chatAnonymous.bindTo(this.contextKeyService);
		this.anonymousContextKey.set(this.anonymous);

		// Only apply the workbench-side default if no other source (e.g. the Copilot extension)
		// has already set this key; binding would otherwise reset it to the declared default.
		if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.clientByokEnabled.key) === undefined) {
			ChatEntitlementContextKeys.clientByokEnabled.bindTo(this.contextKeyService);
		}

		this.onDidChangeEntitlement = Event.map(
			Event.filter(
				this.contextKeyService.onDidChangeContext, e => e.affectsSome(new Set([
					ChatEntitlementContextKeys.Entitlement.planEdu.key,
					ChatEntitlementContextKeys.Entitlement.planPro.key,
					ChatEntitlementContextKeys.Entitlement.planBusiness.key,
					ChatEntitlementContextKeys.Entitlement.planEnterprise.key,
					ChatEntitlementContextKeys.Entitlement.planProPlus.key,
					ChatEntitlementContextKeys.Entitlement.planMax.key,
					ChatEntitlementContextKeys.Entitlement.planFree.key,
					ChatEntitlementContextKeys.Entitlement.canSignUp.key,
					ChatEntitlementContextKeys.Entitlement.signedOut.key,
					ChatEntitlementContextKeys.Entitlement.organisations.key,
					ChatEntitlementContextKeys.Entitlement.internal.key,
					ChatEntitlementContextKeys.Entitlement.sku.key
				])), this._store
			), () => { }, this._store
		);
		this.entitlementObs = observableFromEvent(this.onDidChangeEntitlement, () => this.entitlement);

		this.onDidChangeSentiment = Event.map(
			Event.filter(
				this.contextKeyService.onDidChangeContext, e => e.affectsSome(new Set([
					ChatEntitlementContextKeys.Setup.completed.key,
					ChatEntitlementContextKeys.Setup.hidden.key,
					ChatEntitlementContextKeys.Setup.disabled.key,
					ChatEntitlementContextKeys.Setup.untrusted.key,
					ChatEntitlementContextKeys.Setup.installed.key,
					ChatEntitlementContextKeys.Setup.later.key,
					ChatEntitlementContextKeys.Setup.registered.key
				])), this._store
			), () => { }, this._store
		);
		this.sentimentObs = observableFromEvent(this.onDidChangeSentiment, () => this.sentiment);

		if ((isWeb && !environmentService.remoteAuthority && !environmentService.isSessionsWindow)) {
			ChatEntitlementContextKeys.Setup.hidden.bindTo(this.contextKeyService).set(true); // hide copilot UI on web if unsupported
			return;
		}

		if (!productService.defaultChatAgent) {
			return; // we need a default chat agent configured going forward from here
		}

		const context = this.context = new Lazy(() => this._register(instantiationService.createInstance(ChatEntitlementContext)));
		this.requests = new Lazy(() => this._register(instantiationService.createInstance(ChatEntitlementRequests, context.value, {
			clearQuotas: () => this.clearQuotas(),
			acceptQuotas: quotas => this.acceptQuotas(quotas)
		})));

		this.registerListeners();
	}

	//#region --- Entitlements

	readonly onDidChangeEntitlement: Event<void>;
	readonly entitlementObs: IObservable<ChatEntitlement>;

	get entitlement(): ChatEntitlement {
		if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.planEdu.key) === true) {
			return ChatEntitlement.EDU;
		} else if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.planPro.key) === true) {
			return ChatEntitlement.Pro;
		} else if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.planBusiness.key) === true) {
			return ChatEntitlement.Business;
		} else if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.planEnterprise.key) === true) {
			return ChatEntitlement.Enterprise;
		} else if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.planProPlus.key) === true) {
			return ChatEntitlement.ProPlus;
		} else if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.planMax.key) === true) {
			return ChatEntitlement.Max;
		} else if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.planFree.key) === true) {
			return ChatEntitlement.Free;
		} else if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.canSignUp.key) === true) {
			return ChatEntitlement.Available;
		} else if (this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.signedOut.key) === true) {
			return ChatEntitlement.Unknown;
		}

		return ChatEntitlement.Unresolved;
	}

	get isInternal(): boolean {
		return this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Entitlement.internal.key) === true;
	}

	get organisations(): string[] | undefined {
		return this.contextKeyService.getContextKeyValue<string[]>(ChatEntitlementContextKeys.Entitlement.organisations.key);
	}

	get sku(): string | undefined {
		return this.contextKeyService.getContextKeyValue<string>(ChatEntitlementContextKeys.Entitlement.sku.key);
	}

	get copilotTrackingId(): string | undefined {
		return this.context?.value.state.copilotTrackingId;
	}

	get clientByokEnabled(): boolean {
		return this.contextKeyService.getContextKeyValue<boolean>('github.copilot.clientByokEnabled') === true;
	}

	get hasByokModels(): boolean {
		return this.contextKeyService.getContextKeyValue<boolean>('github.copilot.hasByokModels') === true;
	}

	//#endregion

	//#region --- Quotas

	private readonly _onDidChangeQuotaExceeded = this._register(new Emitter<void>());
	readonly onDidChangeQuotaExceeded = this._onDidChangeQuotaExceeded.event;

	private readonly _onDidChangeQuotaRemaining = this._register(new Emitter<void>());
	readonly onDidChangeQuotaRemaining = this._onDidChangeQuotaRemaining.event;

	private readonly _onDidChangeUsageBasedBilling = this._register(new Emitter<void>());
	readonly onDidChangeUsageBasedBilling = this._onDidChangeUsageBasedBilling.event;

	private _quotas: IQuotas;
	private quotaCopilotTrackingId: string | undefined;
	get quotas() { return this._quotas; }

	private readonly chatQuotaExceededContextKey: IContextKey<boolean>;
	private readonly completionsQuotaExceededContextKey: IContextKey<boolean>;

	private ExtensionQuotaContextKeys = {
		chatQuotaExceeded: defaultChatAgent.chatQuotaExceededContext,
		completionsQuotaExceeded: defaultChatAgent.completionsQuotaExceededContext,
	};

	private registerListeners(): void {
		const quotaExceededSet = new Set([this.ExtensionQuotaContextKeys.chatQuotaExceeded, this.ExtensionQuotaContextKeys.completionsQuotaExceeded]);

		const cts = this._register(new MutableDisposable<CancellationTokenSource>());
		this._register(this.contextKeyService.onDidChangeContext(e => {
			if (e.affectsSome(quotaExceededSet)) {
				if (cts.value) {
					cts.value.cancel();
				}
				cts.value = new CancellationTokenSource();
				this.update(cts.value.token);
			}
		}));

		let anonymousUsage = this.anonymous;

		const updateAnonymousUsage = () => {
			const newAnonymousUsage = this.anonymous;
			if (newAnonymousUsage !== anonymousUsage) {
				anonymousUsage = newAnonymousUsage;
				this.anonymousContextKey.set(newAnonymousUsage);

				if (this.context?.hasValue) {
					logChatEntitlements(this.context.value.state, this.configurationService, this.telemetryService);
				}

				this._onDidChangeAnonymous.fire();
			}
		};

		this._register(this.configurationService.onDidChangeConfiguration(e => {
			if (e.affectsConfiguration(CHAT_ALLOW_ANONYMOUS_CONFIGURATION_KEY)) {
				updateAnonymousUsage();
			}
		}));

		this._register(this.onDidChangeEntitlement(() => updateAnonymousUsage()));
		this._register(this.onDidChangeSentiment(() => updateAnonymousUsage()));
	}

	acceptQuotas(incomingQuotas: IQuotas): void {
		const oldQuota = this._quotas;
		const cachedQuota = this.quotaCopilotTrackingId === this.copilotTrackingId ? oldQuota : {};
		const quotas: IQuotas = {
			...incomingQuotas,
			chat: incomingQuotas.chat ? mergeDefinedSnapshot(cachedQuota.chat, incomingQuotas.chat) : undefined,
			completions: incomingQuotas.completions ? mergeDefinedSnapshot(cachedQuota.completions, incomingQuotas.completions) : undefined,
			premiumChat: incomingQuotas.premiumChat ? mergeDefinedSnapshot(cachedQuota.premiumChat, incomingQuotas.premiumChat) : undefined,
			sessionRateLimit: incomingQuotas.sessionRateLimit ? mergeDefinedSnapshot(cachedQuota.sessionRateLimit, incomingQuotas.sessionRateLimit) : undefined,
			weeklyRateLimit: incomingQuotas.weeklyRateLimit ? mergeDefinedSnapshot(cachedQuota.weeklyRateLimit, incomingQuotas.weeklyRateLimit) : undefined,
		};
		this.quotaCopilotTrackingId = this.copilotTrackingId;
		this._quotas = quotas;
		this.updateContextKeys();

		if (oldQuota.usageBasedBilling !== quotas.usageBasedBilling) {
			if (quotas.usageBasedBilling !== undefined) {
				this.storageService.store(ChatEntitlementService.CACHED_UBB_STORAGE_KEY, quotas.usageBasedBilling, StorageScope.PROFILE, StorageTarget.MACHINE);
			} else {
				this.storageService.remove(ChatEntitlementService.CACHED_UBB_STORAGE_KEY, StorageScope.PROFILE);
			}
		}

		if (this.logService.getLevel() === LogLevel.Trace) {
			this.logService.trace(`[chat entitlement]: acceptQuotas: ${JSON.stringify(quotas)}`);
		}

		const { changed: chatChanged } = this.compareQuotas(oldQuota.chat, quotas.chat);
		const { changed: completionsChanged } = this.compareQuotas(oldQuota.completions, quotas.completions);
		const { changed: premiumChatChanged } = this.compareQuotas(oldQuota.premiumChat, quotas.premiumChat);

		if (chatChanged.exceeded || completionsChanged.exceeded || premiumChatChanged.exceeded) {
			this._onDidChangeQuotaExceeded.fire();
		}

		const sessionRateLimitChanged = this.compareQuotas(oldQuota.sessionRateLimit, quotas.sessionRateLimit).changed.remaining;
		const weeklyRateLimitChanged = this.compareQuotas(oldQuota.weeklyRateLimit, quotas.weeklyRateLimit).changed.remaining;

		if (chatChanged.remaining || completionsChanged.remaining || premiumChatChanged.remaining || sessionRateLimitChanged || weeklyRateLimitChanged || oldQuota.usageBasedBilling !== quotas.usageBasedBilling || oldQuota.additionalUsageEnabled !== quotas.additionalUsageEnabled) {
			this._onDidChangeQuotaRemaining.fire();
		}

		if (oldQuota.usageBasedBilling !== quotas.usageBasedBilling) {
			this._onDidChangeUsageBasedBilling.fire();
		}

		// Track additional spend configuration changes (only when both values come from server snapshots)
		if (oldQuota.additionalUsageEnabled !== undefined && quotas.additionalUsageEnabled !== undefined && oldQuota.additionalUsageEnabled !== quotas.additionalUsageEnabled) {
			this.telemetryService.publicLog2<ChatAdditionalSpendConfigurationEvent, ChatAdditionalSpendConfigurationClassification>('chatAdditionalSpendConfiguration', {
				enabled: quotas.additionalUsageEnabled ?? false,
				entitlement: this.entitlement,
			});
		}

		// Track entering additional spend: included quota just exhausted while additional spend is enabled
		if (quotas.additionalUsageEnabled && quotas.premiumChat?.percentRemaining === 0
			&& oldQuota.premiumChat?.percentRemaining !== undefined && oldQuota.premiumChat.percentRemaining > 0) {
			this.telemetryService.publicLog2<ChatAdditionalSpendActiveEvent, ChatAdditionalSpendActiveClassification>('chatAdditionalSpendActive', {
				entitlement: this.entitlement,
				additionalUsageCount: quotas.additionalUsageCount ?? 0,
			});
		}
	}

	private compareQuotas(oldQuota: IQuotaSnapshot | undefined, newQuota: IQuotaSnapshot | undefined): { changed: { exceeded: boolean; remaining: boolean } } {
		return {
			changed: {
				exceeded: (oldQuota?.percentRemaining === 0) !== (newQuota?.percentRemaining === 0),
				remaining: oldQuota?.percentRemaining !== newQuota?.percentRemaining
					// Pooled or newly limited allowances can become unusable without a percentage change.
					|| oldQuota?.hasQuota !== newQuota?.hasQuota
					|| oldQuota?.unlimited !== newQuota?.unlimited
					|| oldQuota?.usageBasedBilling !== newQuota?.usageBasedBilling
					// Unlimited plans report a constant percentage, so consumed credits are the only signal that usage moved.
					|| oldQuota?.creditsUsed !== newQuota?.creditsUsed
			}
		};
	}

	clearQuotas(): void {
		this.acceptQuotas({});
	}

	private updateContextKeys(): void {
		this.chatQuotaExceededContextKey.set(isChatQuotaExceeded(this.entitlement, this._quotas));
		this.completionsQuotaExceededContextKey.set(this._quotas.completions?.percentRemaining === 0);
	}

	//#endregion

	//#region --- Sentiment

	readonly onDidChangeSentiment: Event<void>;
	readonly sentimentObs: IObservable<IChatSentiment>;

	get sentiment(): IChatSentiment {
		return {
			completed: this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Setup.completed.key) === true,
			installed: this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Setup.installed.key) === true,
			hidden: this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Setup.hidden.key) === true,
			disabledInWorkspace: this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Setup.disabledInWorkspace.key) === true,
			disabled: this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Setup.disabled.key) === true,
			untrusted: this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Setup.untrusted.key) === true,
			later: this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Setup.later.key) === true,
			registered: this.contextKeyService.getContextKeyValue<boolean>(ChatEntitlementContextKeys.Setup.registered.key) === true
		};
	}

	//#endregion

	//region --- Anonymous

	private readonly anonymousContextKey: IContextKey<boolean>;

	private readonly _onDidChangeAnonymous = this._register(new Emitter<void>());
	readonly onDidChangeAnonymous = this._onDidChangeAnonymous.event;

	readonly anonymousObs = observableFromEvent(this.onDidChangeAnonymous, () => this.anonymous);

	get anonymous(): boolean {
		return isAnonymous(this.configurationService, this.entitlement, this.sentiment);
	}

	//#endregion

	markAnonymousRateLimited(): void {
		if (!this.anonymous) {
			return;
		}

		this.chatQuotaExceededContextKey.set(true);
		this._onDidChangeQuotaExceeded.fire();
	}

	markSetupCompleted(): void {
		this.context?.value.update({ completed: true });
	}

	setForceHidden(hidden: boolean): void {
		if (this.context) {
			this.context.value.setForceHidden(hidden);
		} else {
			// No ChatEntitlementContext (e.g. no defaultChatAgent in product.json).
			// Set the context key directly as a fallback.
			ChatEntitlementContextKeys.Setup.hidden.bindTo(this.contextKeyService).set(hidden);
		}
	}

	async update(token: CancellationToken): Promise<void> {
		await this.requests?.value.forceResolveEntitlement(token);
	}
}

//#endregion

//#region Chat Entitlement Request Service

type EntitlementClassification = {
	tid: { classification: 'EndUserPseudonymizedInformation'; purpose: 'BusinessInsight'; comment: 'The anonymized analytics id returned by the service'; endpoint: 'GoogleAnalyticsId' };
	entitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Flag indicating the chat entitlement state' };
	sku: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The SKU of the chat entitlement' };
	quotaChatUnlimited: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user has unlimited chat requests' };
	quotaChatHasQuota: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user currently has chat quota available' };
	quotaChatEntitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The raw chat quota entitlement count' };
	quotaPremiumChat: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The percentage of premium chat requests remaining for the user' };
	quotaPremiumChatUnlimited: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user has unlimited premium chat requests' };
	quotaPremiumChatHasQuota: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user currently has premium chat quota available' };
	quotaPremiumChatEntitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The raw premium chat quota entitlement count' };
	quotaCompletions: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The percentage of completions remaining for the user' };
	quotaCompletionsUnlimited: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user has unlimited completions' };
	quotaCompletionsHasQuota: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user currently has completions quota available' };
	quotaCompletionsEntitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The raw completions quota entitlement count' };
	quotaResetDate: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The date the quota will reset' };
	usageBasedBilling: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user is on usage-based billing' };
	additionalUsageEnabled: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether overage / additional spend is enabled' };
	additionalUsageCount: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The number of overage interactions used' };
	canUpgradePlan: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user is eligible to upgrade their plan' };
	owner: 'bpasero';
	comment: 'Reporting chat entitlements';
};

type EntitlementEvent = {
	entitlement: ChatEntitlement;
	tid: string;
	sku: string | undefined;
	quotaChatUnlimited: boolean | undefined;
	quotaChatHasQuota: boolean | undefined;
	quotaChatEntitlement: number | undefined;
	quotaPremiumChat: number | undefined;
	quotaPremiumChatUnlimited: boolean | undefined;
	quotaPremiumChatHasQuota: boolean | undefined;
	quotaPremiumChatEntitlement: number | undefined;
	quotaCompletions: number | undefined;
	quotaCompletionsUnlimited: boolean | undefined;
	quotaCompletionsHasQuota: boolean | undefined;
	quotaCompletionsEntitlement: number | undefined;
	quotaResetDate: string | undefined;
	usageBasedBilling: boolean | undefined;
	additionalUsageEnabled: boolean | undefined;
	additionalUsageCount: number | undefined;
	canUpgradePlan: boolean | undefined;
};

interface IEntitlements {
	readonly entitlement: ChatEntitlement;
	readonly organisations?: string[];
	readonly isStaff?: boolean;
	readonly sku?: string;
	readonly copilotTrackingId?: string;
	readonly quotas?: IQuotas;
}

/** Shared exhaustion semantics for both controls and continuation suggestions. */
export function isChatQuotaExceeded(entitlement: ChatEntitlement, quotas: IChatEntitlementService['quotas']): boolean {
	const premium = quotas.premiumChat;
	const managed = entitlement === ChatEntitlement.Business || entitlement === ChatEntitlement.Enterprise;
	const exhausted = premium?.unlimited ? premium.hasQuota === false : premium?.percentRemaining === 0;
	return quotas.chat?.percentRemaining === 0 || (managed && premium?.hasQuota === false) || (exhausted && (managed || !quotas.additionalUsageEnabled));
}

/** Unknown premium allowance cannot establish that a paid route is usable. */
export function hasUsableCopilotPremiumQuota(entitlement: ChatEntitlement, quotas: IChatEntitlementService['quotas']): boolean {
	const premium = quotas.premiumChat;
	return isProUser(entitlement) && !!premium
		&& (premium.unlimited || (Number.isFinite(premium.percentRemaining) && premium.percentRemaining >= 0 && premium.percentRemaining <= 100))
		&& !isChatQuotaExceeded(entitlement, quotas);
}

export interface IQuotaSnapshot {
	readonly percentRemaining: number;
	readonly unlimited: boolean;
	readonly hasQuota?: boolean;
	/** When this quota resets, as a Unix timestamp in *seconds*. */
	readonly resetAt?: number;
	readonly usageBasedBilling?: boolean;
	readonly entitlement?: number;
	readonly quotaRemaining?: number;
	/**
	 * Aggregate credits consumed, only for users whose org sets no user-level budget. Drawn from
	 * an unmeterable pool, so it has no denominator. See microsoft/vscode#319589.
	 */
	readonly creditsUsed?: number;
}

export const enum QuotaUsageKind {

	/**
	 * Consumption measured against a known entitlement, best surfaced as a percentage.
	 */
	Percentage,

	/**
	 * Absolute credits consumed. Used for plans without a cap to measure a percentage against.
	 */
	CreditsUsed
}

export type IQuotaUsage = {
	readonly kind: QuotaUsageKind.Percentage;
	readonly usedPercentage: number;
	/**
	 * Credits consumed and the entitlement they are measured against. Both are
	 * `undefined` when the plan reports no entitlement to break the percentage down into.
	 */
	readonly used: number | undefined;
	readonly total: number | undefined;
} | {
	readonly kind: QuotaUsageKind.CreditsUsed;
	readonly creditsUsed: number;
};

/**
 * Derives how a quota snapshot should be surfaced, so that every quota display agrees on
 * what "used" means. Returns `undefined` when the snapshot carries no usage worth showing,
 * i.e. an unlimited plan that reports no credits, or a depleted pooled organization quota.
 */
export function getQuotaUsage(quota: IQuotaSnapshot | undefined): IQuotaUsage | undefined {
	if (!quota) {
		return undefined;
	}

	if (quota.unlimited) {
		if (quota.hasQuota === false || typeof quota.creditsUsed !== 'number') {
			return undefined;
		}

		return { kind: QuotaUsageKind.CreditsUsed, creditsUsed: quota.creditsUsed };
	}

	// An entitlement of `0` carries no ratio to report, so it is treated as absent.
	const total = quota.entitlement || undefined;
	let used: number | undefined;
	if (total !== undefined) {
		// `creditsUsed` has no denominator, so dividing it by `entitlement` disagrees with
		// `percentRemaining`. `quotaRemaining` shares the entitlement's basis.
		used = Math.max(0, quota.quotaRemaining !== undefined
			? total - quota.quotaRemaining
			: total * (100 - quota.percentRemaining) / 100);
	}

	return {
		kind: QuotaUsageKind.Percentage,
		usedPercentage: Math.max(0, 100 - quota.percentRemaining),
		used,
		total
	};
}

export interface IQuotaReset {
	readonly date: Date;
	/** Whether the underlying source carries a time of day worth surfacing. */
	readonly hasTime: boolean;
}

/**
 * Resolves which reset applies to a quota. A snapshot's own `resetAt` is authoritative for
 * that category and wins over the coarser account-level reset, so that a quota is never
 * paired with a clock that does not govern it.
 */
export function getQuotaReset(quota: IQuotaSnapshot | undefined, accountReset: { readonly resetDate?: string; readonly resetDateHasTime?: boolean }): IQuotaReset | undefined {
	if (quota?.resetAt) {
		return { date: new Date(quota.resetAt * 1000), hasTime: true };
	}

	if (!accountReset.resetDate) {
		return undefined;
	}

	const parsed = Date.parse(accountReset.resetDate);
	if (isNaN(parsed)) {
		return undefined;
	}

	return { date: new Date(parsed), hasTime: !!accountReset.resetDateHasTime };
}

export interface IRateLimitSnapshot {
	readonly percentRemaining: number;
	readonly unlimited: boolean;
	readonly resetDate?: string;
}

interface IQuotas {
	readonly resetDate?: string;
	readonly resetDateHasTime?: boolean;

	readonly usageBasedBilling?: boolean;
	readonly canUpgradePlan?: boolean;

	readonly chat?: IQuotaSnapshot;
	readonly completions?: IQuotaSnapshot;
	readonly premiumChat?: IQuotaSnapshot;
	readonly additionalUsageEnabled?: boolean;
	readonly additionalUsageCount?: number;
	readonly additionalUsageEntitlement?: number;

	readonly sessionRateLimit?: IRateLimitSnapshot;
	readonly weeklyRateLimit?: IRateLimitSnapshot;
}

function mergeDefinedSnapshot<T extends object>(previous: T | undefined, current: T): T {
	const result = { ...previous, ...current };
	for (const key of Object.keys(current) as (keyof T)[]) {
		if (current[key] === undefined && previous?.[key] !== undefined) {
			result[key] = previous[key];
		}
	}
	return result;
}

export function parseQuotas(entitlementsData: IEntitlementsData): IQuotas {
	const quotas: Mutable<IQuotas> = {
		resetDate: entitlementsData.quota_reset_date_utc ?? entitlementsData.quota_reset_date ?? entitlementsData.limited_user_reset_date,
		resetDateHasTime: typeof entitlementsData.quota_reset_date_utc === 'string',
		usageBasedBilling: entitlementsData.token_based_billing,
		canUpgradePlan: entitlementsData.can_upgrade_plan,
	};

	// Legacy Free SKU Quota
	if (entitlementsData.monthly_quotas?.chat && typeof entitlementsData.limited_user_quotas?.chat === 'number') {
		quotas.chat = {
			percentRemaining: Math.min(100, Math.max(0, (entitlementsData.limited_user_quotas.chat / entitlementsData.monthly_quotas.chat) * 100)),
			unlimited: false
		};
	}

	if (entitlementsData.monthly_quotas?.completions && typeof entitlementsData.limited_user_quotas?.completions === 'number') {
		quotas.completions = {
			percentRemaining: Math.min(100, Math.max(0, (entitlementsData.limited_user_quotas.completions / entitlementsData.monthly_quotas.completions) * 100)),
			unlimited: false
		};
	}

	// New Quota Snapshot
	if (entitlementsData.quota_snapshots) {
		for (const quotaType of ['chat', 'completions', 'premium_interactions'] as const) {
			const rawQuotaSnapshot = entitlementsData.quota_snapshots[quotaType];
			if (!rawQuotaSnapshot) {
				continue;
			}
			const parsedEntitlement = rawQuotaSnapshot.entitlement !== undefined ? Number(rawQuotaSnapshot.entitlement) : undefined;
			const parsedCreditsUsed = rawQuotaSnapshot.credits_used !== undefined ? Number(rawQuotaSnapshot.credits_used) : undefined;

			// Skip snapshots where the user has no allocated entitlement for this
			// category (e.g. free tier premium_interactions with 0 credits). Under
			// TBB, has_quota is always false at the per-snapshot level so we cannot
			// rely on it; instead check the actual entitlement value.
			if (!rawQuotaSnapshot.unlimited && parsedEntitlement === 0) {
				continue;
			}

			const parsedQuotaRemaining = rawQuotaSnapshot.quota_remaining !== undefined ? Number(rawQuotaSnapshot.quota_remaining) : undefined;
			const quotaSnapshot: IQuotaSnapshot = {
				percentRemaining: Math.min(100, Math.max(0, rawQuotaSnapshot.percent_remaining)),
				unlimited: rawQuotaSnapshot.unlimited,
				hasQuota: rawQuotaSnapshot.has_quota,
				usageBasedBilling: entitlementsData.token_based_billing,
				resetAt: rawQuotaSnapshot.quota_reset_at || undefined,
				entitlement: parsedEntitlement !== undefined && Number.isFinite(parsedEntitlement) && parsedEntitlement >= 0 ? parsedEntitlement : undefined,
				quotaRemaining: parsedQuotaRemaining !== undefined && Number.isFinite(parsedQuotaRemaining) && parsedQuotaRemaining >= 0 ? parsedQuotaRemaining : undefined,
				creditsUsed: parsedCreditsUsed !== undefined && Number.isFinite(parsedCreditsUsed) && parsedCreditsUsed >= 0 ? parsedCreditsUsed : undefined,
			};

			switch (quotaType) {
				case 'chat':
					quotas.chat = quotaSnapshot;
					break;
				case 'completions':
					quotas.completions = quotaSnapshot;
					break;
				case 'premium_interactions':
					quotas.premiumChat = quotaSnapshot;
					break;
			}
		}

		const overageSource = entitlementsData.quota_snapshots['premium_interactions'];
		quotas.additionalUsageEnabled = overageSource?.overage_permitted ?? false;
		quotas.additionalUsageCount = overageSource?.overage_count ?? 0;
		quotas.additionalUsageEntitlement = overageSource?.overage_entitlement ?? 0;
	}
	return quotas;
}

export class ChatEntitlementRequests extends Disposable {

	private state: IEntitlements;

	private pendingResolveCts = new CancellationTokenSource();

	constructor(
		private readonly context: ChatEntitlementContext,
		private readonly chatQuotasAccessor: IChatQuotasAccessor,
		@ITelemetryService private readonly telemetryService: ITelemetryService,
		@ILogService private readonly logService: ILogService,
		@IRequestService private readonly requestService: IRequestService,
		@IDialogService private readonly dialogService: IDialogService,
		@IOpenerService private readonly openerService: IOpenerService,
		@ILifecycleService private readonly lifecycleService: ILifecycleService,
		@IDefaultAccountService private readonly defaultAccountService: IDefaultAccountService,
		@IAuthenticationService private readonly authenticationService: IAuthenticationService,
	) {
		super();

		this.state = { entitlement: this.context.state.entitlement };

		this.registerListeners();

		this.resolve();
	}

	private registerListeners(): void {
		this._register(this.defaultAccountService.onDidChangeDefaultAccount(() => this.resolve()));

		this._register(this.context.onDidChange(() => {
			if (this.context.state.disabled || this.context.state.entitlement === ChatEntitlement.Unknown) {
				// When the extension is disabled or the user is not entitled
				// make sure to clear quotas so that any indicators are also gone
				this.state = { entitlement: this.state.entitlement, quotas: undefined };
				this.chatQuotasAccessor.clearQuotas();
			}
		}));
	}

	private async resolve(): Promise<void> {
		this.pendingResolveCts.dispose(true);
		const cts = this.pendingResolveCts = new CancellationTokenSource();

		const defaultAccount = await this.defaultAccountService.getDefaultAccount();
		if (cts.token.isCancellationRequested) {
			return;
		}

		// Immediately signal whether we have a session or not
		let state: IEntitlements | undefined = undefined;
		if (defaultAccount) {
			// Do not overwrite any state we have already
			if (this.state.entitlement === ChatEntitlement.Unknown) {
				state = { entitlement: ChatEntitlement.Unresolved };
			}
		} else {
			state = { entitlement: ChatEntitlement.Unknown };
		}
		if (state) {
			this.update(state);
		}

		if (defaultAccount) {
			// Afterwards resolve entitlement with a network request
			// but only unless it was not already resolved before.
			await this.resolveEntitlement(defaultAccount, cts.token);
		}
	}

	private async resolveEntitlement(defaultAccount: IDefaultAccount, token: CancellationToken): Promise<IEntitlements | undefined> {
		const entitlements = await this.doResolveEntitlement(defaultAccount, token);
		if (typeof entitlements?.entitlement === 'number' && !token.isCancellationRequested) {
			this.update(entitlements);
		}
		return entitlements;
	}

	private async doResolveEntitlement(defaultAccount: IDefaultAccount, token: CancellationToken): Promise<IEntitlements | undefined> {
		if (token.isCancellationRequested) {
			return undefined;
		}

		const entitlementsData = defaultAccount.entitlementsData;
		if (!entitlementsData) {
			this.logService.trace('[chat entitlement]: no entitlements data available on default account');
			return { entitlement: entitlementsData === null ? ChatEntitlement.Unknown : ChatEntitlement.Unresolved };
		}

		let entitlement: ChatEntitlement;
		if (entitlementsData.access_type_sku === 'free_limited_copilot') {
			entitlement = ChatEntitlement.Free;
		} else if (entitlementsData.access_type_sku === 'free_educational_quota') {
			entitlement = ChatEntitlement.EDU;
		} else if (entitlementsData.can_signup_for_limited) {
			entitlement = ChatEntitlement.Available;
		} else if (entitlementsData.copilot_plan === 'individual_edu') {
			entitlement = ChatEntitlement.EDU;
		} else if (entitlementsData.copilot_plan === 'individual') {
			entitlement = ChatEntitlement.Pro;
		} else if (entitlementsData.copilot_plan === 'individual_pro') {
			entitlement = ChatEntitlement.ProPlus;
		} else if (entitlementsData.copilot_plan === 'individual_max') {
			entitlement = ChatEntitlement.Max;
		} else if (entitlementsData.copilot_plan === 'business') {
			entitlement = ChatEntitlement.Business;
		} else if (entitlementsData.copilot_plan === 'enterprise') {
			entitlement = ChatEntitlement.Enterprise;
		} else {
			entitlement = ChatEntitlement.Unavailable;
		}

		const entitlements: IEntitlements = {
			entitlement,
			organisations: entitlementsData.organization_login_list,
			isStaff: entitlementsData.is_staff,
			quotas: this.toQuotas(entitlementsData),
			sku: entitlementsData.access_type_sku,
			copilotTrackingId: entitlementsData.analytics_tracking_id
		};

		this.logService.trace(`[chat entitlement]: resolved to ${entitlements.entitlement}, quotas: ${JSON.stringify(entitlements.quotas)}`);
		this.telemetryService.publicLog2<EntitlementEvent, EntitlementClassification>('chatInstallEntitlement', {
			entitlement: entitlements.entitlement,
			tid: entitlementsData.analytics_tracking_id,
			sku: entitlements.sku,
			quotaChatUnlimited: entitlements.quotas?.chat?.unlimited,
			quotaChatHasQuota: entitlements.quotas?.chat?.hasQuota,
			quotaChatEntitlement: entitlements.quotas?.chat?.entitlement,
			quotaPremiumChat: entitlements.quotas?.premiumChat?.percentRemaining,
			quotaPremiumChatUnlimited: entitlements.quotas?.premiumChat?.unlimited,
			quotaPremiumChatHasQuota: entitlements.quotas?.premiumChat?.hasQuota,
			quotaPremiumChatEntitlement: entitlements.quotas?.premiumChat?.entitlement,
			quotaCompletions: entitlements.quotas?.completions?.percentRemaining,
			quotaCompletionsUnlimited: entitlements.quotas?.completions?.unlimited,
			quotaCompletionsHasQuota: entitlements.quotas?.completions?.hasQuota,
			quotaCompletionsEntitlement: entitlements.quotas?.completions?.entitlement,
			quotaResetDate: entitlements.quotas?.resetDate,
			usageBasedBilling: entitlements.quotas?.usageBasedBilling,
			additionalUsageEnabled: entitlements.quotas?.additionalUsageEnabled,
			additionalUsageCount: entitlements.quotas?.additionalUsageCount,
			canUpgradePlan: entitlements.quotas?.canUpgradePlan
		});

		return entitlements;
	}

	private toQuotas(entitlementsData: IEntitlementsData): IQuotas {
		return parseQuotas(entitlementsData);
	}

	private async request(url: string, type: 'GET', body: undefined, sessions: AuthenticationSession[], token: CancellationToken, callSite: string): Promise<IRequestContext | undefined>;
	private async request(url: string, type: 'POST', body: object, sessions: AuthenticationSession[], token: CancellationToken, callSite: string): Promise<IRequestContext | undefined>;
	private async request(url: string, type: 'GET' | 'POST', body: object | undefined, sessions: AuthenticationSession[], token: CancellationToken, callSite: string): Promise<IRequestContext | undefined> {
		let lastRequest: IRequestContext | undefined;

		for (const session of sessions) {
			if (token.isCancellationRequested) {
				return lastRequest;
			}

			try {
				const response = await this.requestService.request({
					type,
					url,
					data: type === 'POST' ? JSON.stringify(body) : undefined,
					disableCache: true,
					headers: {
						'Authorization': `Bearer ${session.accessToken}`
					},
					callSite
				}, token);

				const status = response.res.statusCode;
				if (status && status !== 200) {
					lastRequest = response;
					continue; // try next session
				}

				return response;
			} catch (error) {
				if (!token.isCancellationRequested) {
					this.logService.error(`[chat entitlement] request: error ${error}`);
				}
			}
		}

		return lastRequest;
	}

	private update(state: IEntitlements): void {
		this.state = state;

		this.context.update({ entitlement: this.state.entitlement, organisations: this.state.organisations, isStaff: this.state.isStaff, sku: this.state.sku, copilotTrackingId: this.state.copilotTrackingId });

		if (state.quotas) {
			this.chatQuotasAccessor.acceptQuotas(state.quotas);
		}
	}

	async forceResolveEntitlement(token = CancellationToken.None): Promise<IEntitlements | undefined> {
		const defaultAccount = await this.defaultAccountService.refresh({ refreshEntitlements: true });
		if (!defaultAccount) {
			return undefined;
		}

		return this.resolveEntitlement(defaultAccount, token);
	}

	async signUpFree(): Promise<true /* signed up */ | false /* already signed up */ | { errorCode: number } /* error */ | undefined /* no session */> {
		const sessions = await this.getSessions();
		if (sessions.length === 0) {
			return undefined;
		}
		return this.doSignUpFree(sessions);
	}

	private async doSignUpFree(sessions: AuthenticationSession[]): Promise<true /* signed up */ | false /* already signed up */ | { errorCode: number } /* error */> {
		const body = {
			restricted_telemetry: this.telemetryService.telemetryLevel === TelemetryLevel.NONE ? 'disabled' : 'enabled',
			public_code_suggestions: 'enabled'
		};

		const response = await this.request(defaultChatAgent.entitlementSignupLimitedUrl, 'POST', body, sessions, CancellationToken.None, 'chatEntitlementService.signUpFree');
		if (!response) {
			const retry = await this.onUnknownSignUpError(localize('signUpNoResponseError', "No response received."), '[chat entitlement] sign-up: no response');
			return retry ? this.doSignUpFree(sessions) : { errorCode: 1 };
		}

		if (response.res.statusCode && response.res.statusCode !== 200) {
			if (response.res.statusCode === 422) {
				try {
					const responseText = await asText(response);
					if (responseText) {
						const responseError: { message: string } = JSON.parse(responseText);
						if (typeof responseError.message === 'string' && responseError.message) {
							this.onUnprocessableSignUpError(`[chat entitlement] sign-up: unprocessable entity (${responseError.message})`, responseError.message);
							return { errorCode: response.res.statusCode };
						}
					}
				} catch (error) {
					// ignore - handled below
				}
			}
			const retry = await this.onUnknownSignUpError(localize('signUpUnexpectedStatusError', "Unexpected status code {0}.", response.res.statusCode), `[chat entitlement] sign-up: unexpected status code ${response.res.statusCode}`);
			return retry ? this.doSignUpFree(sessions) : { errorCode: response.res.statusCode };
		}

		let responseText: string | null = null;
		try {
			responseText = await asText(response);
		} catch (error) {
			// ignore - handled below
		}

		if (!responseText) {
			const retry = await this.onUnknownSignUpError(localize('signUpNoResponseContentsError', "Response has no contents."), '[chat entitlement] sign-up: response has no content');
			return retry ? this.doSignUpFree(sessions) : { errorCode: 2 };
		}

		let parsedResult: { subscribed: boolean } | undefined = undefined;
		try {
			parsedResult = JSON.parse(responseText);
			this.logService.trace(`[chat entitlement] sign-up: response is ${responseText}`);
		} catch (err) {
			const retry = await this.onUnknownSignUpError(localize('signUpInvalidResponseError', "Invalid response contents."), `[chat entitlement] sign-up: error parsing response (${err})`);
			return retry ? this.doSignUpFree(sessions) : { errorCode: 3 };
		}

		// We have made it this far, so the user either did sign-up or was signed-up already.
		// That is, because the endpoint throws in all other case according to Patrick.
		this.update({ entitlement: ChatEntitlement.Free });

		return Boolean(parsedResult?.subscribed);
	}

	private async getSessions(): Promise<AuthenticationSession[]> {
		const defaultAccount = await this.defaultAccountService.getDefaultAccount();
		if (defaultAccount) {
			const sessions = await this.authenticationService.getSessions(defaultAccount.authenticationProvider.id);
			const accountSessions = sessions.filter(s => s.id === defaultAccount.sessionId);
			if (accountSessions.length) {
				return accountSessions;
			}
		}
		return [...(await this.authenticationService.getSessions(this.defaultAccountService.getDefaultAccountAuthenticationProvider().id))];
	}

	private async onUnknownSignUpError(detail: string, logMessage: string): Promise<boolean> {
		this.logService.error(logMessage);

		if (!this.lifecycleService.willShutdown) {
			const { confirmed } = await this.dialogService.confirm({
				type: Severity.Error,
				message: localize('unknownSignUpError', "An error occurred while signing up for the GitHub Copilot Free plan. Would you like to try again?"),
				detail,
				primaryButton: localize('retry', "Retry")
			});

			return confirmed;
		}

		return false;
	}

	private onUnprocessableSignUpError(logMessage: string, logDetails: string): void {
		this.logService.error(logMessage);

		if (!this.lifecycleService.willShutdown) {
			this.dialogService.prompt({
				type: Severity.Error,
				message: localize('unprocessableSignUpError', "An error occurred while signing up for the GitHub Copilot Free plan."),
				detail: logDetails,
				buttons: [
					{
						label: localize('ok', "OK"),
						run: () => { /* noop */ }
					},
					{
						label: localize('learnMore', "Learn More"),
						run: () => this.openerService.open(URI.parse(defaultChatAgent.upgradePlanUrl))
					}
				]
			});
		}
	}

	async signIn(options?: { useSocialProvider?: string; additionalScopes?: readonly string[] }): Promise<{ defaultAccount?: IDefaultAccount; entitlements?: IEntitlements }> {
		const defaultAccount = await this.defaultAccountService.signIn({
			additionalScopes: options?.additionalScopes,
			extraAuthorizeParameters: { get_started_with: 'copilot-vscode' },
			provider: options?.useSocialProvider
		});
		if (!defaultAccount) {
			return {};
		}

		const entitlements = await this.doResolveEntitlement(defaultAccount, CancellationToken.None);
		return { defaultAccount, entitlements };
	}

	override dispose(): void {
		this.pendingResolveCts.dispose(true);

		super.dispose();
	}
}

//#endregion

//#region Context

export interface IChatEntitlementContextState extends IChatSentiment {

	/**
	 * Users last known or resolved entitlement.
	 */
	entitlement: ChatEntitlement;

	/**
	 * User's last known or resolved raw SKU type.
	 */
	sku: string | undefined;

	/**
	 * User's last known or resolved organisations.
	 */
	organisations: string[] | undefined;

	/**
	 * Whether the user's last known or resolved account is a GitHub or Microsoft staff account.
	 */
	isStaff: boolean | undefined;

	/**
	 * User's Copilot tracking ID from the entitlement API.
	 */
	copilotTrackingId: string | undefined;
}

export class ChatEntitlementContext extends Disposable {

	private static readonly CHAT_ENTITLEMENT_CONTEXT_STORAGE_KEY = 'chat.setupContext';
	private static readonly CHAT_ENTITLEMENT_CONTEXT_MIGRATED_STORAGE_KEY = 'chat.setupContext.migrated.v1';

	private readonly canSignUpContextKey: IContextKey<boolean>;
	private readonly signedOutContextKey: IContextKey<boolean>;

	private readonly freeContextKey: IContextKey<boolean>;
	private readonly eduContextKey: IContextKey<boolean>;
	private readonly proContextKey: IContextKey<boolean>;
	private readonly proPlusContextKey: IContextKey<boolean>;
	private readonly maxContextKey: IContextKey<boolean>;
	private readonly businessContextKey: IContextKey<boolean>;
	private readonly enterpriseContextKey: IContextKey<boolean>;

	private readonly organisationsContextKey: IContextKey<string[] | undefined>;
	private readonly isInternalContextKey: IContextKey<boolean>;
	private readonly skuContextKey: IContextKey<string | undefined>;

	private readonly completedContext: IContextKey<boolean>;
	private readonly hiddenContext: IContextKey<boolean>;
	private readonly disabledInWorkspaceContext: IContextKey<boolean>;
	private readonly laterContext: IContextKey<boolean>;
	private readonly installedContext: IContextKey<boolean>;
	private readonly disabledContext: IContextKey<boolean>;
	private readonly untrustedContext: IContextKey<boolean>;
	private readonly registeredContext: IContextKey<boolean>;

	private _state: IChatEntitlementContextState;
	private suspendedState: IChatEntitlementContextState | undefined = undefined;
	get state(): IChatEntitlementContextState { return this.withConfiguration(this.suspendedState ?? this._state); }

	private readonly _onDidChange = this._register(new Emitter<void>());
	readonly onDidChange = this._onDidChange.event;

	private updateBarrier: Barrier | undefined = undefined;

	constructor(
		@IContextKeyService contextKeyService: IContextKeyService,
		@IStorageService private readonly storageService: IStorageService,
		@ILogService private readonly logService: ILogService,
		@IConfigurationService private readonly configurationService: IConfigurationService,
		@ITelemetryService private readonly telemetryService: ITelemetryService
	) {
		super();

		this.canSignUpContextKey = ChatEntitlementContextKeys.Entitlement.canSignUp.bindTo(contextKeyService);
		this.signedOutContextKey = ChatEntitlementContextKeys.Entitlement.signedOut.bindTo(contextKeyService);

		this.freeContextKey = ChatEntitlementContextKeys.Entitlement.planFree.bindTo(contextKeyService);
		this.eduContextKey = ChatEntitlementContextKeys.Entitlement.planEdu.bindTo(contextKeyService);
		this.proContextKey = ChatEntitlementContextKeys.Entitlement.planPro.bindTo(contextKeyService);
		this.proPlusContextKey = ChatEntitlementContextKeys.Entitlement.planProPlus.bindTo(contextKeyService);
		this.maxContextKey = ChatEntitlementContextKeys.Entitlement.planMax.bindTo(contextKeyService);
		this.businessContextKey = ChatEntitlementContextKeys.Entitlement.planBusiness.bindTo(contextKeyService);
		this.enterpriseContextKey = ChatEntitlementContextKeys.Entitlement.planEnterprise.bindTo(contextKeyService);

		this.organisationsContextKey = ChatEntitlementContextKeys.Entitlement.organisations.bindTo(contextKeyService);
		this.isInternalContextKey = ChatEntitlementContextKeys.Entitlement.internal.bindTo(contextKeyService);
		this.skuContextKey = ChatEntitlementContextKeys.Entitlement.sku.bindTo(contextKeyService);

		this.completedContext = ChatEntitlementContextKeys.Setup.completed.bindTo(contextKeyService);
		this.hiddenContext = ChatEntitlementContextKeys.Setup.hidden.bindTo(contextKeyService);
		this.disabledInWorkspaceContext = ChatEntitlementContextKeys.Setup.disabledInWorkspace.bindTo(contextKeyService);
		this.laterContext = ChatEntitlementContextKeys.Setup.later.bindTo(contextKeyService);
		this.installedContext = ChatEntitlementContextKeys.Setup.installed.bindTo(contextKeyService);
		this.disabledContext = ChatEntitlementContextKeys.Setup.disabled.bindTo(contextKeyService);
		this.untrustedContext = ChatEntitlementContextKeys.Setup.untrusted.bindTo(contextKeyService);
		this.registeredContext = ChatEntitlementContextKeys.Setup.registered.bindTo(contextKeyService);

		this._state = this.storageService.getObject<IChatEntitlementContextState>(ChatEntitlementContext.CHAT_ENTITLEMENT_CONTEXT_STORAGE_KEY, StorageScope.PROFILE) ?? {
			entitlement: ChatEntitlement.Unknown,
			organisations: undefined,
			isStaff: undefined,
			sku: undefined,
			copilotTrackingId: undefined
		};

		const migrated = this.storageService.getBoolean(ChatEntitlementContext.CHAT_ENTITLEMENT_CONTEXT_MIGRATED_STORAGE_KEY, StorageScope.PROFILE) === true;
		if (!migrated) {
			this.storageService.store(ChatEntitlementContext.CHAT_ENTITLEMENT_CONTEXT_MIGRATED_STORAGE_KEY, true, StorageScope.PROFILE, StorageTarget.MACHINE);
			if (this._state.installed && !this._state.completed) {
				this._state.completed = true; // treat installation signal as completed signal once
				this.storageService.store(ChatEntitlementContext.CHAT_ENTITLEMENT_CONTEXT_STORAGE_KEY, this._state, StorageScope.PROFILE, StorageTarget.MACHINE);
			}
		}

		this.updateContextSync();

		this.registerListeners();
	}

	private registerListeners(): void {
		this._register(this.configurationService.onDidChangeConfiguration(e => {
			if (e.affectsConfiguration(ChatAIDisabledSettingId)) {
				this.updateContext();
			}
		}));
	}

	private _forceHidden = false;

	private withConfiguration(state: IChatEntitlementContextState): IChatEntitlementContextState {
		if (this._forceHidden || this.configurationService.getValue(ChatAIDisabledSettingId) === true) {
			return {
				...state,
				hidden: true
			};
		}

		return state;
	}

	setForceHidden(hidden: boolean): void {
		if (this._forceHidden !== hidden) {
			this._forceHidden = hidden;
			this.updateContext();
		}
	}

	update(context: { installed: boolean; disabled: boolean; untrusted: boolean; disabledInWorkspace: boolean }): Promise<void>;
	update(context: { completed: true }): Promise<void>;
	update(context: { hidden: false }): Promise<void>; // legacy UI state from before we had a setting to hide, keep around to still support users who used this
	update(context: { later: boolean }): Promise<void>;
	update(context: { entitlement: ChatEntitlement; organisations: string[] | undefined; isStaff: boolean | undefined; sku: string | undefined; copilotTrackingId: string | undefined }): Promise<void>;
	async update(context: { completed?: boolean; installed?: boolean; disabled?: boolean; untrusted?: boolean; disabledInWorkspace?: boolean; hidden?: false; later?: boolean; entitlement?: ChatEntitlement; organisations?: string[]; isStaff?: boolean; sku?: string; copilotTrackingId?: string }): Promise<void> {
		this.logService.trace(`[chat entitlement context] update(): ${JSON.stringify(context)}`);

		const oldState = JSON.stringify(this._state);

		if (typeof context.installed === 'boolean' && typeof context.disabled === 'boolean' && typeof context.untrusted === 'boolean') {
			this._state.installed = context.installed;
			this._state.disabled = context.disabled;
			this._state.untrusted = context.untrusted;
			this._state.disabledInWorkspace = context.disabledInWorkspace;

			if (context.installed && !context.disabled) {
				context.hidden = false; // treat this as a sign to make Chat visible again in case it is hidden
			}
		}

		if (typeof context.hidden === 'boolean') {
			this._state.hidden = context.hidden;
		}

		if (typeof context.later === 'boolean') {
			this._state.later = context.later;
		}

		if (typeof context.completed === 'boolean') {
			this._state.completed = context.completed;
		}

		if (typeof context.entitlement === 'number') {
			this._state.entitlement = context.entitlement;
			this._state.organisations = context.organisations;
			this._state.isStaff = context.isStaff;
			this._state.sku = context.sku;
			this._state.copilotTrackingId = context.copilotTrackingId;

			if (this._state.entitlement === ChatEntitlement.Free || isProUser(this._state.entitlement)) {
				this._state.registered = true;
			} else if (this._state.entitlement === ChatEntitlement.Available) {
				this._state.registered = false; // only reset when signed-in user can sign-up for free
			}
		}

		if (isAnonymous(this.configurationService, this._state.entitlement, this._state)) {
			this._state.sku = 'no_auth_limited_copilot'; // no-auth users have a fixed SKU
		}

		if (oldState === JSON.stringify(this._state)) {
			return; // state did not change
		}

		this.storageService.store(ChatEntitlementContext.CHAT_ENTITLEMENT_CONTEXT_STORAGE_KEY, {
			...this._state,
			later: undefined // do not persist this across restarts for now
		}, StorageScope.PROFILE, StorageTarget.MACHINE);

		return this.updateContext();
	}

	private async updateContext(): Promise<void> {
		await this.updateBarrier?.wait();

		this.updateContextSync();
	}

	private updateContextSync(): void {
		const state = this.withConfiguration(this._state);

		this.signedOutContextKey.set(state.entitlement === ChatEntitlement.Unknown);
		this.canSignUpContextKey.set(state.entitlement === ChatEntitlement.Available);

		this.freeContextKey.set(state.entitlement === ChatEntitlement.Free);
		this.eduContextKey.set(state.entitlement === ChatEntitlement.EDU);
		this.proContextKey.set(state.entitlement === ChatEntitlement.Pro);
		this.proPlusContextKey.set(state.entitlement === ChatEntitlement.ProPlus);
		this.maxContextKey.set(state.entitlement === ChatEntitlement.Max);
		this.businessContextKey.set(state.entitlement === ChatEntitlement.Business);
		this.enterpriseContextKey.set(state.entitlement === ChatEntitlement.Enterprise);

		this.organisationsContextKey.set(state.organisations);
		this.isInternalContextKey.set(isInternalAccount(state.isStaff, state.organisations));
		this.skuContextKey.set(state.sku);

		this.completedContext.set(!!state.completed);
		this.hiddenContext.set(!!state.hidden);
		this.disabledInWorkspaceContext.set(!!state.disabledInWorkspace);
		this.laterContext.set(!!state.later);
		this.installedContext.set(!!state.installed);
		this.disabledContext.set(!!state.disabled);
		this.untrustedContext.set(!!state.untrusted);
		this.registeredContext.set(!!state.registered);

		this.logService.trace(`[chat entitlement context] updateContext(): ${JSON.stringify(state)}`);
		logChatEntitlements(state, this.configurationService, this.telemetryService);

		this._onDidChange.fire();
	}

	suspend(): void {
		this.suspendedState = { ...this._state };
		this.updateBarrier = new Barrier();
	}

	resume(): void {
		this.suspendedState = undefined;
		this.updateBarrier?.open();
		this.updateBarrier = undefined;
	}
}

//#endregion

registerSingleton(IChatEntitlementService, ChatEntitlementService, InstantiationType.Eager /* To ensure context keys are set asap */);