/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import type { Column } from 'druid-query-toolkit';
import { C, F, SqlColumn, SqlExpression, SqlQuery, SqlStar } from 'druid-query-toolkit';

import { filterMap, filterOrReturn, mapRecordOrReturn } from '../../../utils';

import { ExpressionMeta } from './expression-meta';
import { Measure } from './measure';
import type { OptionValue, ParameterDefinition, Parameters, ParameterValues } from './parameter';
import { evaluateFunctor } from './parameter';

function expressionWithinColumns(ex: SqlExpression, columns: readonly Column[]): boolean {
  const usedColumns = ex.getUsedColumnNames();
  return usedColumns.every(columnName => columns.some(c => c.name === columnName));
}

interface QuerySourceValue {
  query: SqlQuery;
  baseColumns: readonly Column[];
  columns: readonly Column[];
  measures: Measure[];
}

export class QuerySource {
  static isSimpleSelect(query: SqlQuery): boolean {
    return Boolean(
      !query.hasGroupBy() && !query.unionQuery && query.getFromExpressions().length === 1,
    );
  }

  static isSingleStarQuery(query: SqlQuery): boolean {
    const selectExpressions = query.getSelectExpressionsArray();
    return selectExpressions.length === 1 && selectExpressions[0] instanceof SqlStar;
  }

  static makeLimitZeroIntrospectionQuery(query: SqlQuery): SqlQuery {
    return SqlQuery.selectStarFrom(query).changeLimitValue(0);
  }

  static stripToBaseSource(query: SqlQuery): SqlQuery {
    if (query.hasGroupBy()) {
      if (!query.fromClause) throw new Error('must have FROM clause');
      return SqlQuery.selectStarFrom(query.fromClause);
    } else {
      return query
        .changeSelectExpressions([SqlStar.PLAIN])
        .changeLimitValue(undefined)
        .changeContextStatements([]);
    }
  }

  static fromIntrospectResult(
    query: SqlQuery,
    baseColumns: readonly Column[],
    columns: readonly Column[],
  ): QuerySource {
    let effectiveColumns = columns;
    if (query.getSelectExpressionsArray().some(ex => ex instanceof SqlStar)) {
      // The query has a star so carefully pick the columns that make sense
      effectiveColumns = columns.filter(
        c => c.sqlType !== 'OTHER' || c.nativeType === 'COMPLEX<json>',
      );
    }

    let measures = Measure.extractQueryMeasures(query);
    if (!measures.length) {
      const countColumn = columns.find(c => c.name === 'count' || c.name === '__count');
      measures = [
        countColumn
          ? new Measure({
              expression: F.sum(C(countColumn.name)),
              as: 'Count',
            })
          : Measure.COUNT,
      ];

      measures = [
        ...measures,
        ...filterMap(columns, column =>
          column.nativeType?.startsWith('COMPLEX<')
            ? Measure.getPossibleMeasuresForColumn(column)[0]
            : undefined,
        ),
      ];
    }

    return new QuerySource({
      query,
      baseColumns,
      columns: effectiveColumns,
      measures,
    });
  }

  public readonly query: SqlQuery;
  public readonly baseColumns: readonly Column[];
  public readonly columns: readonly Column[];
  public readonly measures: Measure[];

  constructor(value: QuerySourceValue) {
    this.query = value.query;
    this.baseColumns = value.baseColumns;
    this.columns = value.columns;
    this.measures = value.measures;
  }

  public valueOf(): QuerySourceValue {
    return {
      query: this.query,
      baseColumns: this.baseColumns,
      columns: this.columns,
      measures: this.measures,
    };
  }

  public getInitQuery(where?: SqlExpression): SqlQuery {
    return SqlQuery.from(this.query.as('t')).changeWhereExpression(where);
  }

  public getInitBaseQuery(): SqlQuery {
    return SqlQuery.from(QuerySource.stripToBaseSource(this.query).as('t'));
  }

  private materializeStarIfNeeded(): SqlQuery {
    const { query, columns, measures } = this;
    let columnsToExpand = columns.map(c => c.name);
    const selectExpressions = query.getSelectExpressionsArray();
    let starCount = 0;
    for (const selectExpression of selectExpressions) {
      if (selectExpression instanceof SqlStar) {
        starCount++;
        continue;
      }
      const outputName = selectExpression.getOutputName();
      if (!outputName) continue;
      columnsToExpand = columnsToExpand.filter(c => c !== outputName);
    }
    if (starCount === 0) return query;
    if (starCount > 1) throw new Error('can not handle multiple stars');

    return Measure.addMeasuresToQuery(
      query
        .changeSelectExpressions(
          selectExpressions.flatMap(selectExpression =>
            selectExpression instanceof SqlStar ? columnsToExpand.map(c => C(c)) : selectExpression,
          ),
        )
        .prettify(),
      measures,
    );
  }

  public getFirstAggregateMeasure(): Measure | undefined {
    return this.measures[0]?.toAggregateBasedMeasure();
  }

  public getFirstAggregateMeasureArray(): Measure[] {
    return this.measures.length ? [this.measures[0].toAggregateBasedMeasure()] : [];
  }

  public getAvailableName(prefix: string): string {
    let columnName = prefix;
    let counter = 1;
    while (this.nameInUse(columnName)) {
      counter++;
      columnName = `${prefix}_${counter}`;
    }
    return columnName;
  }

  public nameInUse(nameToCheck: string): boolean {
    return this.hasColumnByName(nameToCheck) || this.hasMeasureByName(nameToCheck);
  }

  public hasColumnByName(name: string): boolean {
    return this.columns.some(c => c.name === name);
  }

  public hasMeasureByName(name: string): boolean {
    return this.measures.some(m => m.name === name);
  }

  public hasBaseTimeColumn(): boolean {
    return this.baseColumns.some(column => column.isTimeColumn());
  }

  public getSourceExpressionForColumn(outputName: string): SqlExpression {
    const selectExpressionsArray = this.query.getSelectExpressionsArray();

    const sourceExpression = selectExpressionsArray.find(ex => ex.getOutputName() === outputName);
    if (sourceExpression) return sourceExpression;

    const m = /^EXPR\$(\d+)$/.exec(outputName);
    if (m) {
      const index = parseInt(m[1], 10);
      if (selectExpressionsArray[index]) return selectExpressionsArray[index];
    }

    return C(outputName);
  }

  private getSourceToBaseSubstitutions(): Map<string, SqlExpression> {
    return new Map<string, SqlExpression>(
      filterMap(this.query.getSelectExpressionsArray(), ex => {
        const outputName = ex.getOutputName();
        const underlyingExpression = ex.getUnderlyingExpression();
        if (!outputName || underlyingExpression.getOutputName() === outputName) return;
        return [outputName, underlyingExpression];
      }),
    );
  }

  public transformExpressionToBaseColumns(expression: SqlExpression): SqlExpression {
    const sourceToBaseSubstitutions = this.getSourceToBaseSubstitutions();
    return expression.walk(ex => {
      if (ex instanceof SqlColumn) {
        return sourceToBaseSubstitutions.get(ex.getName()) || ex;
      }
      return ex;
    }) as SqlExpression;
  }

  public addWhereClause(clause: SqlExpression): SqlQuery {
    return this.query.addWhere(clause);
  }

  public addColumn(newExpression: SqlExpression): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return noStarQuery.addSelect(newExpression);
  }

  public addColumnAfter(neighborName: string, ...newExpressions: SqlExpression[]): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return noStarQuery.changeSelectExpressions(
      noStarQuery
        .getSelectExpressionsArray()
        .flatMap(ex => (ex.getOutputName() === neighborName ? [ex, ...newExpressions] : ex)),
    );
  }

  public changeColumn(oldName: string, newExpression: SqlExpression): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return noStarQuery.changeSelectExpressions(
      noStarQuery
        .getSelectExpressionsArray()
        .map(ex => (ex.getOutputName() === oldName ? newExpression : ex)),
    );
  }

  public deleteColumn(outputName: string): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return noStarQuery.changeSelectExpressions(
      noStarQuery.getSelectExpressionsArray().filter(ex => ex.getOutputName() !== outputName),
    );
  }

  public getColumnNameMap(nameTransform: (columnName: string) => string): Map<string, string> {
    return new Map(this.columns.map(column => [column.name, nameTransform(column.name)]));
  }

  public applyColumnNameMap(columnNameMap: Map<string, string>): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return noStarQuery.changeSelectExpressions(
      noStarQuery.getSelectExpressionsArray().map(ex => {
        const outputName = ex.getOutputName();
        if (!outputName) return ex;
        const newOutputName = columnNameMap.get(outputName);
        if (!newOutputName || newOutputName === outputName) return ex;
        return ex.as(newOutputName);
      }),
    );
  }

  // ------------------------------------

  public addMeasure(measure: Measure): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return Measure.addMeasuresToQuery(noStarQuery, this.measures.concat(measure));
  }

  public addMeasureAfter(neighborName: string, newMeasure: Measure): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return Measure.addMeasuresToQuery(
      noStarQuery,
      this.measures.flatMap(m => (m.name === neighborName ? [m, newMeasure] : m)),
    );
  }

  public changeMeasure(oldName: string, newMeasure: Measure): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return Measure.addMeasuresToQuery(
      noStarQuery,
      this.measures.map(m => (m.name === oldName ? newMeasure : m)),
    );
  }

  public deleteMeasure(measureName: string): SqlQuery {
    const noStarQuery = this.materializeStarIfNeeded();
    return Measure.addMeasuresToQuery(
      noStarQuery,
      this.measures.filter(m => m.name !== measureName),
    );
  }

  // --------------------------------

  public restrictWhere(where: SqlExpression): SqlExpression {
    const { columns } = this;
    const parts = where.decomposeViaAnd();
    const filterParts = parts.filter(ex => expressionWithinColumns(ex, columns));
    if (parts.length === filterParts.length) return where;
    return SqlExpression.and(...filterParts);
  }

  public restrictParameterValues(
    parameterValues: ParameterValues,
    parameters: Parameters,
    where: SqlExpression,
  ): ParameterValues {
    return mapRecordOrReturn(parameterValues, (parameterValue, k) => {
      const parameter = parameters[k];
      if (!parameter) return;
      return this.restrictParameterValue(parameterValue, parameter, where, parameterValues);
    });
  }

  private restrictParameterValue(
    parameterValue: any,
    parameter: ParameterDefinition,
    where: SqlExpression,
    parameterValues: ParameterValues,
  ): any {
    if (typeof parameterValue !== 'undefined') {
      switch (parameter.type) {
        case 'number': {
          if (typeof parameter.min === 'number') {
            parameterValue = Math.max(parameterValue, parameter.min);
          }
          if (typeof parameter.max === 'number') {
            parameterValue = Math.min(parameterValue, parameter.max);
          }
          return parameterValue;
        }

        case 'option': {
          const options = evaluateFunctor(parameter.options, parameterValues, this, where);
          if (!options || !options.includes(parameterValue as OptionValue)) return;
          return parameterValue as OptionValue;
        }

        case 'options': {
          const options = evaluateFunctor(parameter.options, {}, this, where) || [];
          return filterOrReturn<OptionValue>(parameterValue, v => options.includes(v));
        }

        case 'expression':
          if (!this.validateExpressionMeta(parameterValue)) return;
          break;

        case 'measure':
          if (!this.validateMeasure(parameterValue)) return;
          break;

        case 'expressions':
          return filterOrReturn<ExpressionMeta>(parameterValue, v =>
            this.validateExpressionMeta(v),
          );

        case 'measures':
          return filterOrReturn<Measure>(parameterValue, v => this.validateMeasure(v));

        default:
          break;
      }
    }
    return parameterValue;
  }

  public validateExpressionMeta(e: ExpressionMeta | undefined): e is ExpressionMeta {
    if (!(e instanceof ExpressionMeta)) return false;
    return expressionWithinColumns(e.expression, this.columns);
  }

  public validateMeasure(m: Measure | undefined): m is Measure {
    if (!(m instanceof Measure)) return false;

    const usedAggregates = m.getUsedAggregates();
    if (usedAggregates.some(usedAggregate => !this.hasMeasureByName(usedAggregate))) return false;

    return expressionWithinColumns(m.expression, this.columns);
  }
}