From 60dd761211bb3f744eb5f591e5643f2fddb55bf6 Mon Sep 17 00:00:00 2001
From: rebornix <penn.lv@gmail.com>
Date: Wed, 2 Jun 2021 14:56:35 -0700
Subject: [PATCH] update the layout based on treatment.

---
 .../contrib/profile/notebookProfile.ts        | 57 +++++++++++++++++--
 1 file changed, 52 insertions(+), 5 deletions(-)

diff --git a/src/vs/workbench/contrib/notebook/browser/contrib/profile/notebookProfile.ts b/src/vs/workbench/contrib/notebook/browser/contrib/profile/notebookProfile.ts
index 1ea0765f964..a13b215651f 100644
--- a/src/vs/workbench/contrib/notebook/browser/contrib/profile/notebookProfile.ts
+++ b/src/vs/workbench/contrib/notebook/browser/contrib/profile/notebookProfile.ts
@@ -3,11 +3,17 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { Disposable } from 'vs/base/common/lifecycle';
+import { Registry } from 'vs/platform/registry/common/platform';
 import { ServicesAccessor } from 'vs/editor/browser/editorExtensions';
 import { localize } from 'vs/nls';
 import { Action2, registerAction2 } from 'vs/platform/actions/common/actions';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { optional } from 'vs/platform/instantiation/common/instantiation';
 import { CellToolbarLocation, CompactView, ConsolidatedRunButton, FocusIndicator, GlobalToolbar, InsertToolbarPosition, ShowCellStatusBarAfterExecuteKey, ShowCellStatusBarKey, UndoRedoPerCell } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { ITASExperimentService } from 'vs/workbench/services/experiment/common/experimentService';
+import { Extensions as WorkbenchExtensions, IWorkbenchContributionsRegistry } from 'vs/workbench/common/contributions';
+import { LifecyclePhase } from 'vs/workbench/services/lifecycle/common/lifecycle';
 
 export enum NotebookProfileType {
 	default = 'default',
@@ -24,7 +30,8 @@ const profiles = {
 		[CompactView]: true,
 		[ShowCellStatusBarKey]: true,
 		[ShowCellStatusBarAfterExecuteKey]: false,
-		[ConsolidatedRunButton]: true
+		[ConsolidatedRunButton]: true,
+		[UndoRedoPerCell]: false
 	},
 	[NotebookProfileType.jupyter]: {
 		[FocusIndicator]: 'gutter',
@@ -43,12 +50,12 @@ const profiles = {
 		[CellToolbarLocation]: { default: 'right' },
 		[CompactView]: false,
 		[ShowCellStatusBarKey]: false,
-		[ConsolidatedRunButton]: true
+		[ConsolidatedRunButton]: true,
+		[UndoRedoPerCell]: false
 	}
 };
 
-async function applyProfile(accessor: ServicesAccessor, profile: Record<string, any>): Promise<void> {
-	const configService = accessor.get(IConfigurationService);
+async function applyProfile(configService: IConfigurationService, profile: Record<string, any>): Promise<void> {
 	const promises = [];
 	for (let settingKey in profile) {
 		promises.push(configService.updateValue(settingKey, profile[settingKey]));
@@ -74,7 +81,8 @@ registerAction2(class extends Action2 {
 			return;
 		}
 
-		return applyProfile(accessor, profiles[args.profile]);
+		const configService = accessor.get(IConfigurationService);
+		return applyProfile(configService, profiles[args.profile]);
 	}
 });
 
@@ -84,3 +92,42 @@ function isSetProfileArgs(args: unknown): args is ISetProfileArgs {
 		setProfileArgs.profile === NotebookProfileType.default ||
 		setProfileArgs.profile === NotebookProfileType.jupyter;
 }
+
+export class NotebookProfileContribution extends Disposable {
+	constructor(@IConfigurationService configService: IConfigurationService, @optional(ITASExperimentService) private readonly experimentService: ITASExperimentService) {
+		super();
+
+		if (this.experimentService) {
+			this.experimentService.getTreatment<NotebookProfileType.default | NotebookProfileType.jupyter | NotebookProfileType.colab>('notebookprofile').then(treatment => {
+				if (treatment === undefined) {
+					return;
+				} else {
+					// check if settings are already modified
+					const focusIndicator = configService.getValue(FocusIndicator);
+					const insertToolbarPosition = configService.getValue(InsertToolbarPosition);
+					const globalToolbar = configService.getValue(GlobalToolbar);
+					// const cellToolbarLocation = configService.getValue(CellToolbarLocation);
+					const compactView = configService.getValue(CompactView);
+					const showCellStatusBarKey = configService.getValue(ShowCellStatusBarKey);
+					const showCellStatusBarAfterExecuteKey = configService.getValue(ShowCellStatusBarAfterExecuteKey);
+					const consolidatedRunButton = configService.getValue(ConsolidatedRunButton);
+					if (focusIndicator === 'border'
+						&& insertToolbarPosition === 'both'
+						&& globalToolbar === false
+						// && cellToolbarLocation === undefined
+						&& compactView === true
+						&& showCellStatusBarKey === true
+						&& showCellStatusBarAfterExecuteKey === false
+						&& consolidatedRunButton === true
+					) {
+						applyProfile(configService, profiles[treatment] ?? profiles[NotebookProfileType.default]);
+					}
+				}
+			});
+		}
+	}
+}
+
+const workbenchContributionsRegistry = Registry.as<IWorkbenchContributionsRegistry>(WorkbenchExtensions.Workbench);
+workbenchContributionsRegistry.registerWorkbenchContribution(NotebookProfileContribution, LifecyclePhase.Ready);
+