import { ColumnMetadata } from '../metadata/ColumnMetadata';
import { QueryBuilder } from './QueryBuilder';
import { ObjectLiteral } from '../common/ObjectLiteral';
import { DataSource } from '../data-source/DataSource';
import { QueryRunner } from '../query-runner/QueryRunner';
import { WhereExpressionBuilder } from './WhereExpressionBuilder';
import { Brackets } from './Brackets';
import { UpdateResult } from './result/UpdateResult';
import { ReturningStatementNotSupportedError } from '../error/ReturningStatementNotSupportedError';
import { ReturningResultsEntityUpdator } from './ReturningResultsEntityUpdator';
import { OrderByCondition } from '../find-options/OrderByCondition';
import { LimitOnUpdateNotSupportedError } from '../error/LimitOnUpdateNotSupportedError';
import { UpdateValuesMissingError } from '../error/UpdateValuesMissingError';
import { QueryDeepPartialEntity } from './QueryPartialEntity';
import { TypeORMError } from '../error';
import { EntityPropertyNotFoundError } from '../error/EntityPropertyNotFoundError';
import { DriverUtils } from '../driver/DriverUtils';

/**
 * Allows to build complex sql queries in a fashion way and execute those queries.
 */
export class UpdateQueryBuilder<Entity extends ObjectLiteral>
	extends QueryBuilder<Entity>
	implements WhereExpressionBuilder
{
	readonly '@instanceof' = Symbol.for('UpdateQueryBuilder');

	// -------------------------------------------------------------------------
	// Constructor
	// -------------------------------------------------------------------------

	constructor(connectionOrQueryBuilder: DataSource | QueryBuilder<any>, queryRunner?: QueryRunner) {
		super(connectionOrQueryBuilder as any, queryRunner);
		this.expressionMap.aliasNamePrefixingEnabled = false;
	}

	// -------------------------------------------------------------------------
	// Public Implemented Methods
	// -------------------------------------------------------------------------

	/**
	 * Gets generated SQL query without parameters being replaced.
	 */
	getQuery(): string {
		let sql = this.createComment();
		sql += this.createCteExpression();
		sql += this.createUpdateExpression();
		sql += this.createOrderByExpression();
		sql += this.createLimitExpression();
		return this.replacePropertyNamesForTheWholeQuery(sql.trim());
	}

	/**
	 * Executes sql generated by query builder and returns raw database results.
	 */
	async execute(): Promise<UpdateResult> {
		const queryRunner = this.obtainQueryRunner();
		let transactionStartedByUs: boolean = false;

		try {
			// start transaction if it was enabled
			if (this.expressionMap.useTransaction === true && queryRunner.isTransactionActive === false) {
				await queryRunner.startTransaction();
				transactionStartedByUs = true;
			}

			// call before updation methods in listeners and subscribers
			if (this.expressionMap.callListeners === true && this.expressionMap.mainAlias!.hasMetadata) {
				await queryRunner.broadcaster.broadcast(
					'BeforeUpdate',
					this.expressionMap.mainAlias!.metadata,
					this.expressionMap.valuesSet,
				);
			}

			let declareSql: string | null = null;
			let selectOutputSql: string | null = null;

			// if update entity mode is enabled we may need extra columns for the returning statement
			const returningResultsEntityUpdator = new ReturningResultsEntityUpdator(
				queryRunner,
				this.expressionMap,
			);

			const returningColumns: ColumnMetadata[] = [];

			if (
				Array.isArray(this.expressionMap.returning) &&
				this.expressionMap.mainAlias!.hasMetadata
			) {
				for (const columnPath of this.expressionMap.returning) {
					returningColumns.push(
						...this.expressionMap.mainAlias!.metadata.findColumnsWithPropertyPath(columnPath),
					);
				}
			}

			if (
				this.expressionMap.updateEntity === true &&
				this.expressionMap.mainAlias!.hasMetadata &&
				this.expressionMap.whereEntities.length > 0
			) {
				this.expressionMap.extraReturningColumns =
					returningResultsEntityUpdator.getUpdationReturningColumns();

				returningColumns.push(
					...this.expressionMap.extraReturningColumns.filter((c) => !returningColumns.includes(c)),
				);
			}

			// execute update query
			const [updateSql, parameters] = this.getQueryAndParameters();

			const statements = [declareSql, updateSql, selectOutputSql];
			const queryResult = await queryRunner.query(
				statements.filter((sql) => sql != null).join(';\n\n'),
				parameters,
				true,
			);
			const updateResult = UpdateResult.from(queryResult);

			// if we are updating entities and entity updation is enabled we must update some of entity columns (like version, update date, etc.)
			if (
				this.expressionMap.updateEntity === true &&
				this.expressionMap.mainAlias!.hasMetadata &&
				this.expressionMap.whereEntities.length > 0
			) {
				await returningResultsEntityUpdator.update(updateResult, this.expressionMap.whereEntities);
			}

			// call after updation methods in listeners and subscribers
			if (this.expressionMap.callListeners === true && this.expressionMap.mainAlias!.hasMetadata) {
				await queryRunner.broadcaster.broadcast(
					'AfterUpdate',
					this.expressionMap.mainAlias!.metadata,
					this.expressionMap.valuesSet,
				);
			}

			// close transaction if we started it
			if (transactionStartedByUs) await queryRunner.commitTransaction();

			return updateResult;
		} catch (error) {
			// rollback transaction if we started it
			if (transactionStartedByUs) {
				try {
					await queryRunner.rollbackTransaction();
				} catch (rollbackError) {}
			}
			throw error;
		} finally {
			if (queryRunner !== this.queryRunner) {
				// means we created our own query runner
				await queryRunner.release();
			}
		}
	}

	// -------------------------------------------------------------------------
	// Public Methods
	// -------------------------------------------------------------------------

	/**
	 * Values needs to be updated.
	 */
	set(values: QueryDeepPartialEntity<Entity>): this {
		this.expressionMap.valuesSet = values;
		return this;
	}

	/**
	 * Sets WHERE condition in the query builder.
	 * If you had previously WHERE expression defined,
	 * calling this function will override previously set WHERE conditions.
	 * Additionally you can add parameters used in where expression.
	 */
	where(
		where: string | ((qb: this) => string) | Brackets | ObjectLiteral | ObjectLiteral[],
		parameters?: ObjectLiteral,
	): this {
		this.expressionMap.wheres = []; // don't move this block below since computeWhereParameter can add where expressions
		const condition = this.getWhereCondition(where);
		if (condition) this.expressionMap.wheres = [{ type: 'simple', condition: condition }];
		if (parameters) this.setParameters(parameters);
		return this;
	}

	/**
	 * Adds new AND WHERE condition in the query builder.
	 * Additionally you can add parameters used in where expression.
	 */
	andWhere(
		where: string | ((qb: this) => string) | Brackets | ObjectLiteral | ObjectLiteral[],
		parameters?: ObjectLiteral,
	): this {
		this.expressionMap.wheres.push({
			type: 'and',
			condition: this.getWhereCondition(where),
		});
		if (parameters) this.setParameters(parameters);
		return this;
	}

	/**
	 * Adds new OR WHERE condition in the query builder.
	 * Additionally you can add parameters used in where expression.
	 */
	orWhere(
		where: string | ((qb: this) => string) | Brackets | ObjectLiteral | ObjectLiteral[],
		parameters?: ObjectLiteral,
	): this {
		this.expressionMap.wheres.push({
			type: 'or',
			condition: this.getWhereCondition(where),
		});
		if (parameters) this.setParameters(parameters);
		return this;
	}

	/**
	 * Sets WHERE condition in the query builder with a condition for the given ids.
	 * If you had previously WHERE expression defined,
	 * calling this function will override previously set WHERE conditions.
	 */
	whereInIds(ids: any | any[]): this {
		return this.where(this.getWhereInIdsCondition(ids));
	}

	/**
	 * Adds new AND WHERE with conditions for the given ids.
	 */
	andWhereInIds(ids: any | any[]): this {
		return this.andWhere(this.getWhereInIdsCondition(ids));
	}

	/**
	 * Adds new OR WHERE with conditions for the given ids.
	 */
	orWhereInIds(ids: any | any[]): this {
		return this.orWhere(this.getWhereInIdsCondition(ids));
	}
	/**
	 * Optional returning/output clause.
	 * This will return given column values.
	 */
	output(columns: string[]): this;

	/**
	 * Optional returning/output clause.
	 * Returning is a SQL string containing returning statement.
	 */
	output(output: string): this;

	/**
	 * Optional returning/output clause.
	 */
	output(output: string | string[]): this;

	/**
	 * Optional returning/output clause.
	 */
	output(output: string | string[]): this {
		return this.returning(output);
	}

	/**
	 * Optional returning/output clause.
	 * This will return given column values.
	 */
	returning(columns: string[]): this;

	/**
	 * Optional returning/output clause.
	 * Returning is a SQL string containing returning statement.
	 */
	returning(returning: string): this;

	/**
	 * Optional returning/output clause.
	 */
	returning(returning: string | string[]): this;

	/**
	 * Optional returning/output clause.
	 */
	returning(returning: string | string[]): this {
		// not all databases support returning/output cause
		if (!this.connection.driver.isReturningSqlSupported('update')) {
			throw new ReturningStatementNotSupportedError();
		}

		this.expressionMap.returning = returning;
		return this;
	}

	/**
	 * Sets ORDER BY condition in the query builder.
	 * If you had previously ORDER BY expression defined,
	 * calling this function will override previously set ORDER BY conditions.
	 *
	 * Calling order by without order set will remove all previously set order bys.
	 */
	orderBy(): this;

	/**
	 * Sets ORDER BY condition in the query builder.
	 * If you had previously ORDER BY expression defined,
	 * calling this function will override previously set ORDER BY conditions.
	 */
	orderBy(sort: string, order?: 'ASC' | 'DESC', nulls?: 'NULLS FIRST' | 'NULLS LAST'): this;

	/**
	 * Sets ORDER BY condition in the query builder.
	 * If you had previously ORDER BY expression defined,
	 * calling this function will override previously set ORDER BY conditions.
	 */
	orderBy(order: OrderByCondition): this;

	/**
	 * Sets ORDER BY condition in the query builder.
	 * If you had previously ORDER BY expression defined,
	 * calling this function will override previously set ORDER BY conditions.
	 */
	orderBy(
		sort?: string | OrderByCondition,
		order: 'ASC' | 'DESC' = 'ASC',
		nulls?: 'NULLS FIRST' | 'NULLS LAST',
	): this {
		if (sort) {
			if (typeof sort === 'object') {
				this.expressionMap.orderBys = sort as OrderByCondition;
			} else {
				if (nulls) {
					this.expressionMap.orderBys = {
						[sort as string]: { order, nulls },
					};
				} else {
					this.expressionMap.orderBys = { [sort as string]: order };
				}
			}
		} else {
			this.expressionMap.orderBys = {};
		}
		return this;
	}

	/**
	 * Adds ORDER BY condition in the query builder.
	 */
	addOrderBy(
		sort: string,
		order: 'ASC' | 'DESC' = 'ASC',
		nulls?: 'NULLS FIRST' | 'NULLS LAST',
	): this {
		if (nulls) {
			this.expressionMap.orderBys[sort] = { order, nulls };
		} else {
			this.expressionMap.orderBys[sort] = order;
		}
		return this;
	}

	/**
	 * Sets LIMIT - maximum number of rows to be selected.
	 */
	limit(limit?: number): this {
		this.expressionMap.limit = limit;
		return this;
	}

	/**
	 * Indicates if entity must be updated after update operation.
	 * This may produce extra query or use RETURNING / OUTPUT statement (depend on database).
	 * Enabled by default.
	 */
	whereEntity(entity: Entity | Entity[]): this {
		if (!this.expressionMap.mainAlias!.hasMetadata)
			throw new TypeORMError(
				`.whereEntity method can only be used on queries which update real entity table.`,
			);

		this.expressionMap.wheres = [];
		const entities: Entity[] = Array.isArray(entity) ? entity : [entity];
		entities.forEach((entity) => {
			const entityIdMap = this.expressionMap.mainAlias!.metadata.getEntityIdMap(entity);
			if (!entityIdMap)
				throw new TypeORMError(`Provided entity does not have ids set, cannot perform operation.`);

			this.orWhereInIds(entityIdMap);
		});

		this.expressionMap.whereEntities = entities;
		return this;
	}

	/**
	 * Indicates if entity must be updated after update operation.
	 * This may produce extra query or use RETURNING / OUTPUT statement (depend on database).
	 * Enabled by default.
	 */
	updateEntity(enabled: boolean): this {
		this.expressionMap.updateEntity = enabled;
		return this;
	}

	// -------------------------------------------------------------------------
	// Protected Methods
	// -------------------------------------------------------------------------

	/**
	 * Creates UPDATE express used to perform insert query.
	 */
	protected createUpdateExpression() {
		const valuesSet = this.getValueSet();
		const metadata = this.expressionMap.mainAlias!.hasMetadata
			? this.expressionMap.mainAlias!.metadata
			: undefined;

		// it doesn't make sense to update undefined properties, so just skip them
		const valuesSetNormalized: ObjectLiteral = {};
		for (let key in valuesSet) {
			if (valuesSet[key] !== undefined) {
				valuesSetNormalized[key] = valuesSet[key];
			}
		}

		// prepare columns and values to be updated
		const updateColumnAndValues: string[] = [];
		const updatedColumns: ColumnMetadata[] = [];
		if (metadata) {
			this.createPropertyPath(metadata, valuesSetNormalized).forEach((propertyPath) => {
				// todo: make this and other query builder to work with properly with tables without metadata
				const columns = metadata.findColumnsWithPropertyPath(propertyPath);

				if (columns.length <= 0) {
					throw new EntityPropertyNotFoundError(propertyPath, metadata);
				}

				columns.forEach((column) => {
					if (!column.isUpdate || updatedColumns.includes(column)) {
						return;
					}

					updatedColumns.push(column);

					//
					let value = column.getEntityValue(valuesSetNormalized);
					if (
						column.referencedColumn &&
						typeof value === 'object' &&
						!(value instanceof Date) &&
						value !== null &&
						!Buffer.isBuffer(value)
					) {
						value = column.referencedColumn.getEntityValue(value);
					} else if (!(typeof value === 'function')) {
						value = this.connection.driver.preparePersistentValue(value, column);
					}

					// todo: duplication zone
					if (typeof value === 'function') {
						// support for SQL expressions in update query
						updateColumnAndValues.push(this.escape(column.databaseName) + ' = ' + value());
					} else {
						const paramName = this.createParameter(value);

						let expression = null;
						if (
							DriverUtils.isPostgresFamily(this.connection.driver) &&
							this.connection.driver.spatialTypes.indexOf(column.type) !== -1
						) {
							if (column.srid != null) {
								expression = `ST_SetSRID(ST_GeomFromGeoJSON(${paramName}), ${column.srid})::${column.type}`;
							} else {
								expression = `ST_GeomFromGeoJSON(${paramName})::${column.type}`;
							}
						} else {
							expression = paramName;
						}
						updateColumnAndValues.push(this.escape(column.databaseName) + ' = ' + expression);
					}
				});
			});

			// Don't allow calling update only with columns that are `update: false`
			if (updateColumnAndValues.length > 0 || Object.keys(valuesSetNormalized).length === 0) {
				if (metadata.versionColumn && updatedColumns.indexOf(metadata.versionColumn) === -1)
					updateColumnAndValues.push(
						this.escape(metadata.versionColumn.databaseName) +
							' = ' +
							this.escape(metadata.versionColumn.databaseName) +
							' + 1',
					);
				if (metadata.updateDateColumn && updatedColumns.indexOf(metadata.updateDateColumn) === -1)
					updateColumnAndValues.push(
						this.escape(metadata.updateDateColumn.databaseName) + ' = CURRENT_TIMESTAMP',
					); // todo: fix issue with CURRENT_TIMESTAMP(6) being used, can "DEFAULT" be used?!
			}
		} else {
			Object.keys(valuesSetNormalized).map((key) => {
				let value = valuesSetNormalized[key];

				// todo: duplication zone
				if (typeof value === 'function') {
					// support for SQL expressions in update query
					updateColumnAndValues.push(this.escape(key) + ' = ' + value());
				} else {
					// we need to store array values in a special class to make sure parameter replacement will work correctly
					// if (value instanceof Array)
					//     value = new ArrayParameter(value);

					const paramName = this.createParameter(value);
					updateColumnAndValues.push(this.escape(key) + ' = ' + paramName);
				}
			});
		}

		if (updateColumnAndValues.length <= 0) {
			throw new UpdateValuesMissingError();
		}

		// get a table name and all column database names
		const whereExpression = this.createWhereExpression();
		const returningExpression = this.createReturningExpression('update');

		if (returningExpression === '') {
			return `UPDATE ${this.getTableName(
				this.getMainTableName(),
			)} SET ${updateColumnAndValues.join(', ')}${whereExpression}`; // todo: how do we replace aliases in where to nothing?
		}

		return `UPDATE ${this.getTableName(this.getMainTableName())} SET ${updateColumnAndValues.join(
			', ',
		)}${whereExpression} RETURNING ${returningExpression}`;
	}

	/**
	 * Creates "ORDER BY" part of SQL query.
	 */
	protected createOrderByExpression() {
		const orderBys = this.expressionMap.orderBys;
		if (Object.keys(orderBys).length > 0)
			return (
				' ORDER BY ' +
				Object.keys(orderBys)
					.map((columnName) => {
						if (typeof orderBys[columnName] === 'string') {
							return this.replacePropertyNames(columnName) + ' ' + orderBys[columnName];
						} else {
							return (
								this.replacePropertyNames(columnName) +
								' ' +
								(orderBys[columnName] as any).order +
								' ' +
								(orderBys[columnName] as any).nulls
							);
						}
					})
					.join(', ')
			);

		return '';
	}

	/**
	 * Creates "LIMIT" parts of SQL query.
	 */
	protected createLimitExpression(): string {
		let limit: number | undefined = this.expressionMap.limit;

		if (limit) {
			throw new LimitOnUpdateNotSupportedError();
		}

		return '';
	}

	/**
	 * Gets array of values need to be inserted into the target table.
	 */
	protected getValueSet(): ObjectLiteral {
		if (typeof this.expressionMap.valuesSet === 'object') return this.expressionMap.valuesSet;

		throw new UpdateValuesMissingError();
	}
}