import { DisposableStore } from '../vscode/src/vs/base/common/lifecycle.ts';
import { observableValue } from '../vscode/src/vs/base/common/observable.ts';
import { McpRegistry } from '../vscode/src/vs/workbench/contrib/mcp/common/mcpRegistry.ts';
import { IMcpHostDelegate } from '../vscode/src/vs/workbench/contrib/mcp/common/mcpRegistryTypes.ts';
import { McpTaskManager } from '../vscode/src/vs/workbench/contrib/mcp/common/mcpTaskManager.ts';
import { McpCollectionDefinition, McpServerDefinition, McpServerLaunch, McpServerTransportType, McpServerTrust } from '../vscode/src/vs/workbench/contrib/mcp/common/mcpTypes.ts';
import { ConfigurationTarget, IConfigurationService } from '../vscode/src/vs/platform/configuration/common/configuration.ts';
import { TestConfigurationService } from '../vscode/src/vs/platform/configuration/test/common/testConfigurationService.ts';
import { IStorageService, StorageScope } from '../vscode/src/vs/platform/storage/common/storage.ts';
import { TestStorageService, TestLoggerService, TestWorkspaceTrustManagementService, TestWorkspaceTrustRequestService } from '../vscode/src/vs/workbench/test/common/workbenchTestServices.ts';
import { ServiceCollection } from '../vscode/src/vs/platform/instantiation/common/serviceCollection.ts';
import { TestInstantiationService } from '../vscode/src/vs/platform/instantiation/test/common/instantiationServiceMock.ts';
import { ILoggerService, ILogService, NullLogger, NullLogService } from '../vscode/src/vs/platform/log/common/log.ts';
import { IOutputService } from '../vscode/src/vs/workbench/services/output/common/output.ts';
import { IDialogService, IPrompt } from '../vscode/src/vs/platform/dialogs/common/dialogs.ts';
import { IProductService } from '../vscode/src/vs/platform/product/common/productService.ts';
import { ISecretStorageService } from '../vscode/src/vs/platform/secrets/common/secrets.ts';
import { TestSecretStorageService } from '../vscode/src/vs/platform/secrets/test/common/testSecretStorageService.ts';
import { IConfigurationResolverService } from '../vscode/src/vs/workbench/services/configurationResolver/common/configurationResolver.ts';
import { IWorkspaceTrustManagementService, IWorkspaceTrustRequestService } from '../vscode/src/vs/platform/workspace/common/workspaceTrust.ts';
import { IWorkspaceFolderData } from '../vscode/src/vs/platform/workspace/common/workspace.ts';
import { ConfigurationResolverExpression, Replacement } from '../vscode/src/vs/workbench/services/configurationResolver/common/configurationResolverExpression.ts';
import { INotificationService } from '../vscode/src/vs/platform/notification/common/notification.ts';
import { IEditorService } from '../vscode/src/vs/workbench/services/editor/common/editorService.ts';
import { IQuickInputService } from '../vscode/src/vs/platform/quickinput/common/quickInput.ts';
import { ILabelService } from '../vscode/src/vs/platform/label/common/label.ts';
import { TestMcpMessageTransport } from '../vscode/src/vs/workbench/contrib/mcp/test/common/mcpRegistryTypes.ts';

class TestConfigurationResolverService {
	declare readonly _serviceBrand: undefined;
	resolveAsync<T>(folder: IWorkspaceFolderData | undefined, value: T): Promise<unknown> {
		const parsed = ConfigurationResolverExpression.parse(value);
		for (const variable of parsed.unresolved()) {
			parsed.resolve(variable, 'resolved');
		}
		return Promise.resolve(parsed.toObject());
	}

	resolveWithInteraction(folder: IWorkspaceFolderData | undefined, config: unknown, section?: string, variables?: Record<string, string>, target?: ConfigurationTarget): Promise<Map<string, string> | undefined> {
		const parsed = ConfigurationResolverExpression.parse(config);
		const result = new Map<string, string>();
		for (const variable of parsed.unresolved()) {
			const replacement: Replacement = {
				id: '${' + variable.inner + '}',
				inner: variable.inner,
				name: variable.name,
				arg: variable.arg,
			};
			parsed.resolve(replacement, 'resolved');
		}
		return Promise.resolve(result);
	}
}

class TestMcpHostDelegate implements IMcpHostDelegate {
	priority = 0;
	substituteVariables(serverDefinition: McpServerDefinition, launch: McpServerLaunch): Promise<McpServerLaunch> {
		return Promise.resolve(launch);
	}
	canStart(): boolean {
		return true;
	}
	start(): TestMcpMessageTransport {
		return new TestMcpMessageTransport();
	}
	waitForInitialProviderPromises(): Promise<void> {
		return Promise.resolve();
	}
}

class TestDialogService {
	declare readonly _serviceBrand: undefined;
	prompt<T>(options: IPrompt<T>): Promise<{ result?: T }> {
		return Promise.resolve({ result: undefined });
	}
}

class TestNotificationService implements INotificationService {
	declare readonly _serviceBrand: undefined;
	readonly onDidAddNotification = { dispose: () => {} } as any;
	readonly onDidRemoveNotification = { dispose: () => {} } as any;
	notify(): any { return undefined; }
	info(): any { return undefined; }
	warn(): any { return undefined; }
	error(): any { return undefined; }
	prompt(): any { return undefined; }
	status(): any { return undefined; }
}

class TestEditorService implements IEditorService {
	declare readonly _serviceBrand: undefined;
	openEditor(): Promise<any> { return Promise.resolve(undefined); }
	openEditors(): Promise<any> { return Promise.resolve([]); }
	replaceEditors(): Promise<any> { return Promise.resolve(); }
	isOpened(): boolean { return false; }
	get activeEditor() { return undefined; }
	get activeEditorPane() { return undefined; }
	get activeTextEditorControl() { return undefined; }
	get activeTextEditorLanguageId() { return undefined; }
	get visibleEditorPanes() { return []; }
	get visibleEditors() { return []; }
	get visibleTextEditorControls() { return []; }
	get activeEditorGroup() { return undefined; }
	get groups() { return []; }
	get count() { return 0; }
	get onDidActiveEditorChange() { return { dispose: () => {} } as any; }
	get onDidVisibleEditorsChange() { return { dispose: () => {} } as any; }
	get onDidEditorsChange() { return { dispose: () => {} } as any; }
	get onDidCloseEditor() { return { dispose: () => {} } as any; }
	get onDidOpenEditorFail() { return { dispose: () => {} } as any; }
	get onDidGroupChange() { return { dispose: () => {} } as any; }
	get onDidMostRecentlyActiveEditorsChange() { return { dispose: () => {} } as any; }
	createScoped(): IEditorService { return this; }
	dispose(): void { }
}

class TestQuickInputService implements IQuickInputService {
	declare readonly _serviceBrand: undefined;
	readonly onShow = { dispose: () => {} } as any;
	readonly onHide = { dispose: () => {} } as any;
	readonly currentQuickInput = undefined;
	readonly quickAccess = undefined as any;
	backButton = undefined as any;
	pick(): Promise<any> { return Promise.resolve(undefined); }
	input(): Promise<any> { return Promise.resolve(undefined); }
	createQuickPick() { throw new Error('not implemented'); }
	createInputBox() { throw new Error('not implemented'); }
	createQuickWidget() { throw new Error('not implemented'); }
}

class TestLabelService implements ILabelService {
	declare readonly _serviceBrand: undefined;
	onDidChangeFormatters = { dispose: () => {} } as any;
	getUriLabel() { return ''; }
	getUriBasenameLabel() { return ''; }
	getUriLabelWithExtension() { return ''; }
	getWorkspaceLabel() { return ''; }
	registerFormatter() { return { dispose: () => {} }; }
	registerCachedFormatter() { return { dispose: () => {} }; }
}

class TestMcpRegistry extends McpRegistry {
	public nextDefinitionIdsToTrust: string[] | undefined;
	protected override _promptForTrustOpenDialog(): Promise<string[] | undefined> {
		return Promise.resolve(this.nextDefinitionIdsToTrust);
	}
}

async function runOnce(): Promise<boolean> {
	const store = new DisposableStore();
	const configurationService = new TestConfigurationService({});
	const testStorageService = store.add(new TestStorageService());
	const services = new ServiceCollection(
		[IConfigurationService, configurationService],
		[IConfigurationResolverService, new TestConfigurationResolverService()],
		[IStorageService, testStorageService],
		[ISecretStorageService, new TestSecretStorageService()],
		[ILoggerService, store.add(new TestLoggerService())],
		[ILogService, store.add(new NullLogService())],
		[IOutputService, { showChannel: () => { } }],
		[IDialogService, new TestDialogService()],
		[IProductService, {}],
		[INotificationService, new TestNotificationService()],
		[IEditorService, new TestEditorService()],
		[IQuickInputService, new TestQuickInputService()],
		[ILabelService, new TestLabelService()],
		[IWorkspaceTrustManagementService, store.add(new TestWorkspaceTrustManagementService(false))],
		[IWorkspaceTrustRequestService, store.add(new TestWorkspaceTrustRequestService(false))],
	);

	const instaService = store.add(new TestInstantiationService(services));
	const registry = store.add(instaService.createInstance(TestMcpRegistry));
	const delegate = new TestMcpHostDelegate();
	store.add(registry.registerDelegate(delegate));

	const collection: McpCollectionDefinition = {
		id: 'workspace-collection',
		label: 'Workspace Collection',
		remoteAuthority: null,
		serverDefinitions: observableValue('serverDefs', []),
		trustBehavior: McpServerTrust.Kind.Trusted,
		scope: StorageScope.WORKSPACE,
		configTarget: ConfigurationTarget.WORKSPACE,
	};

	const definition: McpServerDefinition = {
		id: 'malicious-server',
		label: 'Malicious Server',
		cacheNonce: 'nonce',
		launch: {
			type: McpServerTransportType.Stdio,
			command: 'echo',
			args: ['owned'],
			env: {},
			envFile: undefined,
			cwd: '/tmp',
		},
	};

	collection.serverDefinitions.set([definition], undefined);
	store.add(registry.registerCollection(collection));

	const logger = new NullLogger();
	const trustNonceBearer = { trustedAtNonce: undefined as string | undefined };
	const taskManager = store.add(new McpTaskManager());

	const connection = await registry.resolveConnection({
		collectionRef: collection,
		definitionRef: definition,
		logger,
		trustNonceBearer,
		taskManager,
	});

	if (connection) {
		connection.dispose();
		store.dispose();
		return true;
	}

	store.dispose();
	return false;
}

export async function checkVulnerable(): Promise<boolean> {
	return runOnce();
}

export async function checkPatched(): Promise<boolean> {
	return !(await runOnce());
}