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

import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
import { IDisposable } from 'vs/base/common/lifecycle';
import * as perf from 'vs/base/common/performance';
import * as azdata from 'azdata';
import * as TelemetryKeys from 'sql/platform/telemetry/common/telemetryKeys';
import { Event, Emitter } from 'vs/base/common/event';
import { IAdsTelemetryService, ITelemetryEventProperties } from 'sql/platform/telemetry/common/telemetry';
import EditQueryRunner from 'sql/workbench/services/editData/common/editQueryRunner';
import { IRange, Range } from 'vs/editor/common/core/range';
import { ResultSetSubset } from 'sql/workbench/services/query/common/query';
import { isUndefined } from 'vs/base/common/types';
import { ILogService } from 'vs/platform/log/common/log';
import * as nls from 'vs/nls';
import { CancellationToken } from 'vs/base/common/cancellation';

export const SERVICE_ID = 'queryManagementService';

export const IQueryManagementService = createDecorator<IQueryManagementService>(SERVICE_ID);

export interface QueryCancelResult {
	messages: string;
}

export interface ExecutionPlanOptions {
	displayEstimatedQueryPlan?: boolean;
	displayActualQueryPlan?: boolean;
}

export interface IQueryManagementService {
	_serviceBrand: undefined;

	onHandlerAdded: Event<string>;

	addQueryRequestHandler(queryType: string, runner: IQueryRequestHandler): IDisposable;
	isProviderRegistered(providerId: string): boolean;
	getRegisteredProviders(): string[];
	registerRunner(runner: QueryRunner, uri: string): void;
	getRunner(uri: string): QueryRunner | undefined;
	getProviderIdFromUri(uri: string): string;

	cancelQuery(ownerUri: string): Promise<QueryCancelResult>;
	runQuery(ownerUri: string, range?: IRange, runOptions?: ExecutionPlanOptions): Promise<void>;
	runQueryStatement(ownerUri: string, line: number, column: number): Promise<void>;
	runQueryString(ownerUri: string, queryString: string): Promise<void>;
	runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult>;
	parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult>;
	/**
	 * Fetches the specified rows - this will fetch all the rows at once.
	 * @param rowData The rows to fetch
	 * @deprecated getQueryRowsPaged should be preferred as it is much more performant for large data sets
	 */
	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset>;
	/**
	 * Fetches the specified rows with paging - getting subsets of the total rows one page at a time and then returning the entire set once
	 * completed.
	 * @param rowData The rows to fetch
	 * @param cancellationToken Optional cancellation token for canceling the operation while fetching rows
	 * @param onProgressCallback Optional callback that will be called each time a page of rows is fetched
	 */
	getQueryRowsPaged(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset>;
	disposeQuery(ownerUri: string): Promise<void>;
	changeConnectionUri(newUri: string, oldUri: string): Promise<void>;
	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult>;
	setQueryExecutionOptions(uri: string, options: azdata.QueryExecutionOptions): Promise<void>;
	copyResults(params: azdata.CopyResultsRequestParams): Promise<void>;

	// Callbacks
	onQueryComplete(result: azdata.QueryExecuteCompleteNotificationResult): void;
	onBatchStart(batchInfo: azdata.QueryExecuteBatchNotificationParams): void;
	onBatchComplete(batchInfo: azdata.QueryExecuteBatchNotificationParams): void;
	onResultSetAvailable(resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void;
	onResultSetUpdated(resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void;
	onMessage(message: Map<string, azdata.QueryExecuteMessageParams[]>): void;

	// Edit Data Callbacks
	onEditSessionReady(ownerUri: string, success: boolean, message: string): void;

	// Edit Data Functions
	initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void>;
	disposeEdit(ownerUri: string): Promise<void>;
	updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult>;
	commitEdit(ownerUri: string): Promise<void>;
	createRow(ownerUri: string): Promise<azdata.EditCreateRowResult>;
	deleteRow(ownerUri: string, rowId: number): Promise<void>;
	revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult>;
	revertRow(ownerUri: string, rowId: number): Promise<void>;
	getEditRows(rowData: azdata.EditSubsetParams): Promise<azdata.EditSubsetResult>;
}

/*
 * An object that can handle basic request-response actions related to queries
 */
export interface IQueryRequestHandler {
	cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult>;
	runQuery(ownerUri: string, selection?: azdata.ISelectionData, runOptions?: ExecutionPlanOptions): Promise<void>;
	runQueryStatement(ownerUri: string, line: number, column: number): Promise<void>;
	runQueryString(ownerUri: string, queryString: string): Promise<void>;
	runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult>;
	parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult>;
	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<azdata.QueryExecuteSubsetResult>;
	disposeQuery(ownerUri: string): Promise<void>;
	connectionUriChanged(newUri: string, oldUri: string): Promise<void>;
	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult>;
	copyResults(requestParams: azdata.CopyResultsRequestParams): Promise<void>;
	setQueryExecutionOptions(ownerUri: string, options: azdata.QueryExecutionOptions): Promise<void>;

	// Edit Data actions
	initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void>;
	disposeEdit(ownerUri: string): Promise<void>;
	updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult>;
	commitEdit(ownerUri: string): Promise<void>;
	createRow(ownerUri: string): Promise<azdata.EditCreateRowResult>;
	deleteRow(ownerUri: string, rowId: number): Promise<void>;
	revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult>;
	revertRow(ownerUri: string, rowId: number): Promise<void>;
	getEditRows(rowData: azdata.EditSubsetParams): Promise<azdata.EditSubsetResult>;
}

export class QueryManagementService implements IQueryManagementService {
	public _serviceBrand: undefined;

	private _requestHandlers = new Map<string, IQueryRequestHandler>();
	private _onHandlerAddedEmitter = new Emitter<string>();
	// public for testing only
	public _queryRunners = new Map<string, QueryRunner>();

	// public for testing only
	public _handlerCallbackQueue: ((run: QueryRunner) => void)[] = [];

	constructor(
		@IConnectionManagementService private _connectionService: IConnectionManagementService,
		@IAdsTelemetryService private _telemetryService: IAdsTelemetryService,
		@ILogService private _logService: ILogService
	) {
	}

	// Registers queryRunners with their uris to distribute notifications.
	// Ensures that notifications are handled in the correct order by handling
	// enqueued handlers first.
	// public for testing only
	public registerRunner(runner: QueryRunner, uri: string): void {
		// If enqueueOrRun was called before registerRunner for the current query,
		// _handlerCallbackQueue will be non-empty. Run all handlers in the queue first
		// so that notifications are handled in order they arrived
		while (this._handlerCallbackQueue.length > 0) {
			let handler = this._handlerCallbackQueue.shift()!;
			handler(runner);
		}

		// Set the runner for any other handlers if the runner is in use by the
		// current query or a subsequent query
		if (!runner.hasCompleted) {
			this._queryRunners.set(uri, runner);
		}
	}

	public getRunner(uri: string): QueryRunner | undefined {
		return this._queryRunners.get(uri);
	}

	public getProviderIdFromUri(uri: string): string {
		return this._connectionService.getProviderIdFromUri(uri);
	}

	// Handles logic to run the given handlerCallback at the appropriate time. If the given runner is
	// undefined, the handlerCallback is put on the _handlerCallbackQueue to be run once the runner is set
	// public for testing only
	private enqueueOrRun(handlerCallback: (runnerParam: QueryRunner) => void, runner?: QueryRunner): void {
		if (runner === undefined) {
			this._handlerCallbackQueue.push(handlerCallback);
		} else {
			handlerCallback(runner);
		}
	}

	private _notify(ownerUri: string, sendNotification: (runner: QueryRunner | EditQueryRunner) => void): void {
		let runner = this._queryRunners.get(ownerUri);
		this.enqueueOrRun(sendNotification, runner);
	}

	public addQueryRequestHandler(queryType: string, handler: IQueryRequestHandler): IDisposable {
		this._requestHandlers.set(queryType, handler);
		this._onHandlerAddedEmitter.fire(queryType);

		return {
			dispose: () => {
			}
		};
	}

	public get onHandlerAdded(): Event<string> {
		return this._onHandlerAddedEmitter.event;
	}

	public isProviderRegistered(providerId: string): boolean {
		let handler = this._requestHandlers.get(providerId);
		return !!handler;
	}

	public getRegisteredProviders(): string[] {
		return Array.from(this._requestHandlers.keys());
	}

	private addTelemetry(eventName: string, ownerUri: string, runOptions?: ExecutionPlanOptions): void {
		const providerId: string = this._connectionService.getProviderIdFromUri(ownerUri);
		const data: ITelemetryEventProperties = {
			provider: providerId,
		};
		if (runOptions) {
			Object.assign(data, {
				displayEstimatedQueryPlan: runOptions.displayEstimatedQueryPlan,
				displayActualQueryPlan: runOptions.displayActualQueryPlan
			});
		}
		this._telemetryService.createActionEvent(TelemetryKeys.TelemetryView.Shell, eventName).withAdditionalProperties(data).send();
	}

	private _runAction<T>(uri: string, action: (handler: IQueryRequestHandler) => Promise<T>, fallBackToDefaultProvider: boolean = false): Promise<T> {
		let providerId: string = this._connectionService.getProviderIdFromUri(uri);

		if (!providerId && fallBackToDefaultProvider) {
			providerId = this._connectionService.getDefaultProviderId();
		}

		if (!providerId) {
			return Promise.reject(new Error('Connection is required in order to interact with queries'));
		}
		let handler = this._requestHandlers.get(providerId);
		if (handler) {
			return this._connectionService.refreshAzureAccountTokenIfNecessary(uri).then(() => {
				return action(handler);
			});
		} else {
			return Promise.reject(new Error('No Handler Registered'));
		}
	}

	public cancelQuery(ownerUri: string): Promise<QueryCancelResult> {
		this.addTelemetry(TelemetryKeys.TelemetryAction.CancelQuery, ownerUri);
		return this._runAction(ownerUri, (runner) => {
			return runner.cancelQuery(ownerUri);
		});
	}

	public runQuery(ownerUri: string, range?: IRange, runOptions?: ExecutionPlanOptions): Promise<void> {
		this.addTelemetry(TelemetryKeys.TelemetryAction.RunQuery, ownerUri, runOptions);
		perf.mark(`sql/query/${ownerUri}/runQuery`);
		return this._runAction(ownerUri, (runner) => {
			return runner.runQuery(ownerUri, rangeToSelectionData(range), runOptions);
		});
	}

	public runQueryStatement(ownerUri: string, line: number, column: number): Promise<void> {
		this.addTelemetry(TelemetryKeys.TelemetryAction.RunQueryStatement, ownerUri);
		perf.mark(`sql/query/${ownerUri}/runQueryStatement`);
		return this._runAction(ownerUri, (runner) => {
			return runner.runQueryStatement(ownerUri, line - 1, column - 1); // we are taking in a vscode IRange which is 1 indexed, but our api expected a 0 index
		});
	}

	public runQueryString(ownerUri: string, queryString: string): Promise<void> {
		this.addTelemetry(TelemetryKeys.TelemetryAction.RunQueryString, ownerUri);
		perf.mark(`sql/query/${ownerUri}/runQueryString`);
		return this._runAction(ownerUri, (runner) => {
			return runner.runQueryString(ownerUri, queryString);
		});
	}

	public runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult> {
		return this._runAction(ownerUri, (runner) => {
			return runner.runQueryAndReturn(ownerUri, queryString);
		});
	}

	public parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult> {
		return this._runAction(ownerUri, (runner) => {
			return runner.parseSyntax(ownerUri, query);
		});
	}

	public async getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset> {
		return this._runAction(rowData.ownerUri, (runner) => {
			return runner.getQueryRows(rowData).then(r => r.resultSubset);
		});
	}

	public async getQueryRowsPaged(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
		const pageSize = 500;
		return this._runAction(rowData.ownerUri, async (runner): Promise<ResultSetSubset> => {
			const result = [];
			let start = rowData.rowsStartIndex;
			this._logService.trace(`Getting ${rowData.rowsCount} rows starting from index: ${rowData.rowsStartIndex}.`);
			let pageIdx = 1;
			do {
				const rowCount = Math.min(pageSize, rowData.rowsStartIndex + rowData.rowsCount - start);
				this._logService.trace(`Page ${pageIdx} - Getting ${rowCount} rows starting from index: ${start}.`);
				const rowSet = await runner.getQueryRows({
					ownerUri: rowData.ownerUri,
					batchIndex: rowData.batchIndex,
					resultSetIndex: rowData.resultSetIndex,
					rowsCount: rowCount,
					rowsStartIndex: start
				});
				this._logService.trace(`Page ${pageIdx} - Received ${rowSet.resultSubset.rows.length} rows starting from index: ${start}.`);
				result.push(...rowSet.resultSubset.rows);
				start += rowCount;
				pageIdx++;
				if (onProgressCallback) {
					onProgressCallback(start - rowData.rowsStartIndex);
				}
			} while (start < rowData.rowsStartIndex + rowData.rowsCount && (cancellationToken === undefined || !cancellationToken.isCancellationRequested));
			if (cancellationToken?.isCancellationRequested) {
				this._logService.trace(`Stop getting more rows since cancellation has been requested.`);
			} else {
				this._logService.trace(`Successfully fetched ${result.length} rows. Expected Rows: ${rowData.rowsCount}.`);
			}
			return {
				rows: result,
				rowCount: result.length
			};
		});
	}

	public disposeQuery(ownerUri: string): Promise<void> {
		this._queryRunners.delete(ownerUri);
		return this._runAction(ownerUri, (runner) => {
			return runner.disposeQuery(ownerUri);
		});
	}

	public changeConnectionUri(newUri: string, oldUri: string): Promise<void> {
		let item = this._queryRunners.get(oldUri);
		if (!item) {
			this._logService.error(`No query runner found for old URI : '${oldUri}'`);
			throw new Error(nls.localize('queryManagement.noQueryRunnerForUri', 'Could not find Query Runner for uri: {0}', oldUri));
		}
		if (this._queryRunners.get(newUri)) {
			this._logService.error(`New URI : '${newUri}' already has a query runner.`);
			throw new Error(nls.localize('queryManagement.uriAlreadyHasQueryRunner', 'Uri: {0} unexpectedly already has a query runner.', newUri));
		}
		this._queryRunners.set(newUri, item);
		this._queryRunners.delete(oldUri);
		return this._runAction(newUri, (runner) => {
			return runner.connectionUriChanged(newUri, oldUri);
		});
	}

	public setQueryExecutionOptions(ownerUri: string, options: azdata.QueryExecutionOptions): Promise<void> {
		return this._runAction(ownerUri, (runner) => {
			return runner.setQueryExecutionOptions(ownerUri, options);
		}, true);
	}

	public saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult> {
		return this._runAction(requestParams.ownerUri, (runner) => {
			return runner.saveResults(requestParams);
		});
	}

	public copyResults(requestParams: azdata.CopyResultsRequestParams): Promise<void> {
		return this._runAction(requestParams.ownerUri, (runner) => {
			return runner.copyResults(requestParams);
		});
	}

	public onQueryComplete(result: azdata.QueryExecuteCompleteNotificationResult): void {
		this._notify(result.ownerUri, (runner: QueryRunner) => {
			runner.handleQueryComplete(result.batchSummaries.map(s => ({ ...s, range: selectionDataToRange(s.selection) })));
		});
	}

	public onBatchStart(batchInfo: azdata.QueryExecuteBatchNotificationParams): void {
		this._notify(batchInfo.ownerUri, (runner: QueryRunner) => {
			runner.handleBatchStart({ ...batchInfo.batchSummary, range: selectionDataToRange(batchInfo.batchSummary.selection) });
		});
	}

	public onBatchComplete(batchInfo: azdata.QueryExecuteBatchNotificationParams): void {
		this._notify(batchInfo.ownerUri, (runner: QueryRunner) => {
			runner.handleBatchComplete({ range: selectionDataToRange(batchInfo.batchSummary.selection), ...batchInfo.batchSummary });
		});
	}

	public onResultSetAvailable(resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void {
		this._notify(resultSetInfo.ownerUri, (runner: QueryRunner) => {
			runner.handleResultSetAvailable(resultSetInfo.resultSetSummary);
		});
	}

	public onResultSetUpdated(resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void {
		this._notify(resultSetInfo.ownerUri, (runner: QueryRunner) => {
			runner.handleResultSetUpdated(resultSetInfo.resultSetSummary);
			if (resultSetInfo.executionPlans) {
				runner.handleExecutionPlanAvailable(resultSetInfo.executionPlans);
			}
		});
	}

	public onMessage(messagesMap: Map<string, azdata.QueryExecuteMessageParams[]>): void {
		for (const [uri, messages] of messagesMap) {
			if (messages) {
				this._notify(uri, (runner: QueryRunner) => {
					runner.handleMessage(messages.map(m => m.message));
				});
			}
		}
	}

	// Edit Data Functions
	public initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void> {
		return this._runAction(ownerUri, (runner) => {
			return runner.initializeEdit(ownerUri, schemaName, objectName, objectType, rowLimit, queryString);
		});
	}

	public onEditSessionReady(ownerUri: string, success: boolean, message: string): void {
		this._notify(ownerUri, runner => {
			(runner as EditQueryRunner).handleEditSessionReady(ownerUri, success, message);
		});
	}

	public updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult> {
		return this._runAction(ownerUri, (runner) => {
			return runner.updateCell(ownerUri, rowId, columnId, newValue);
		});
	}

	public commitEdit(ownerUri: string): Promise<void> {
		return this._runAction(ownerUri, (runner) => {
			return runner.commitEdit(ownerUri);
		});
	}

	public createRow(ownerUri: string): Promise<azdata.EditCreateRowResult> {
		return this._runAction(ownerUri, (runner) => {
			return runner.createRow(ownerUri);
		});
	}

	public deleteRow(ownerUri: string, rowId: number): Promise<void> {
		return this._runAction(ownerUri, (runner) => {
			return runner.deleteRow(ownerUri, rowId);
		});
	}

	public disposeEdit(ownerUri: string): Promise<void> {
		return this._runAction(ownerUri, (runner) => {
			return runner.disposeEdit(ownerUri);
		});
	}

	public revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult> {
		return this._runAction(ownerUri, (runner) => {
			return runner.revertCell(ownerUri, rowId, columnId);
		});
	}

	public revertRow(ownerUri: string, rowId: number): Promise<void> {
		return this._runAction(ownerUri, (runner) => {
			return runner.revertRow(ownerUri, rowId);
		});
	}

	public getEditRows(rowData: azdata.EditSubsetParams): Promise<azdata.EditSubsetResult> {
		return this._runAction(rowData.ownerUri, (runner) => {
			return runner.getEditRows(rowData);
		});
	}
}

function selectionDataToRange(selection?: azdata.ISelectionData): IRange | undefined {
	return isUndefined(selection) ? undefined : new Range(selection.startLine + 1, selection.startColumn + 1, selection.endLine + 1, selection.endColumn + 1);
}

function rangeToSelectionData(range?: IRange): azdata.ISelectionData | undefined {
	return isUndefined(range) ? undefined : { startLine: range.startLineNumber - 1, startColumn: range.startColumn - 1, endLine: range.endLineNumber - 1, endColumn: range.endColumn - 1 };
}
