/**
 * Copyright (c) 2020-present, Goldman Sachs
 *
 * Licensed 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 { action, flow, makeObservable, observable } from 'mobx';
import { assertErrorThrown, LogEvent, guaranteeNonNullable, ActionState, StopWatch, getContentTypeFileExtension, } from '@finos/legend-shared';
import { GRAPH_MANAGER_EVENT, buildRawLambdaFromLambdaFunction, reportGraphAnalytics, TDSExecutionResult, V1_ZIPKIN_TRACE_HEADER, ExecutionError, V1_DELEGATED_EXPORT_HEADER, } from '@finos/legend-graph';
import { buildLambdaFunction } from './QueryBuilderValueSpecificationBuilder.js';
import { buildExecutionParameterValues, getExecutionQueryFromRawLambda, } from './shared/LambdaParameterState.js';
import { QueryBuilderTelemetryHelper } from '../__lib__/QueryBuilderTelemetryHelper.js';
import { QUERY_BUILDER_EVENT } from '../__lib__/QueryBuilderEvent.js';
import { ExecutionPlanState } from './execution-plan/ExecutionPlanState.js';
import { downloadStream } from '@finos/legend-application';
import { QueryBuilderDataGridCustomAggregationFunction } from '../components/result/tds/QueryBuilderTDSGridResult.js';
import { QueryBuilderTDSState } from './fetch-structure/tds/QueryBuilderTDSState.js';
import { LineageState } from './lineage/LineageState.js';
export const DEFAULT_LIMIT = 1000;
export class QueryBuilderResultWavgAggregationState {
    isApplyingWavg = false;
    weightedColumnIdPairs;
    constructor() {
        makeObservable(this, {
            isApplyingWavg: observable,
            weightedColumnIdPairs: observable,
            setIsApplyingWavg: action,
            addWeightedColumnIdPair: action,
            removeWeightedColumnIdPair: action,
        });
        this.weightedColumnIdPairs = new Map();
    }
    setIsApplyingWavg(val) {
        this.isApplyingWavg = val;
    }
    addWeightedColumnIdPair(col, weightedColumnId) {
        this.weightedColumnIdPairs.set(col, weightedColumnId);
    }
    removeWeightedColumnIdPair(col) {
        this.weightedColumnIdPairs.delete(col);
    }
}
export class QueryBuilderResultState {
    queryBuilderState;
    executionPlanState;
    exportState = ActionState.create();
    previewLimit = DEFAULT_LIMIT;
    pressedRunQuery = ActionState.create();
    isRunningQuery = false;
    isGeneratingPlan = false;
    executionResult;
    isExecutionResultOverflowing = false;
    executionDuration;
    executionTraceId;
    latestRunHashCode;
    queryRunPromise = undefined;
    isQueryUsageViewerOpened = false;
    executionError;
    selectedCells;
    mousedOverCell = null;
    isSelectingCells;
    gridConfig;
    wavgAggregationState;
    lineageState;
    isGeneratingLineage = false;
    constructor(queryBuilderState) {
        makeObservable(this, {
            executionResult: observable,
            executionTraceId: observable,
            previewLimit: observable,
            executionDuration: observable,
            latestRunHashCode: observable,
            queryRunPromise: observable,
            isGeneratingPlan: observable,
            selectedCells: observable,
            mousedOverCell: observable,
            isRunningQuery: observable,
            isSelectingCells: observable,
            isQueryUsageViewerOpened: observable,
            isExecutionResultOverflowing: observable,
            gridConfig: observable,
            wavgAggregationState: observable,
            executionError: observable,
            setGridConfig: action,
            setWavgAggregationState: action,
            setIsSelectingCells: action,
            setIsRunningQuery: action,
            setExecutionResult: action,
            setExecutionTraceId: action,
            setExecutionDuration: action,
            setPreviewLimit: action,
            addSelectedCell: action,
            setSelectedCells: action,
            setMouseOverCell: action,
            setQueryRunPromise: action,
            setIsQueryUsageViewerOpened: action,
            setIsExecutionResultOverflowing: action,
            handlePreConfiguredGridConfig: action,
            updatePreviewLimitInConfig: action,
            setExecutionError: action,
            exportData: flow,
            runQuery: flow,
            cancelQuery: flow,
            generatePlan: flow,
            generateLineage: flow,
        });
        this.isSelectingCells = false;
        this.selectedCells = [];
        this.gridConfig = undefined;
        this.queryBuilderState = queryBuilderState;
        this.executionPlanState = new ExecutionPlanState(this.queryBuilderState.applicationStore, this.queryBuilderState.graphManagerState);
        this.lineageState = new LineageState(this.queryBuilderState.applicationStore);
    }
    setGridConfig(val) {
        this.gridConfig = val;
    }
    setWavgAggregationState(val) {
        this.wavgAggregationState = val;
    }
    setIsSelectingCells(val) {
        this.isSelectingCells = val;
    }
    setIsRunningQuery(val) {
        this.isRunningQuery = val;
    }
    setExecutionResult(val) {
        this.executionResult = val;
    }
    setExecutionTraceId(val) {
        this.executionTraceId = val;
    }
    setExecutionDuration(val) {
        this.executionDuration = val;
    }
    setPreviewLimit(val) {
        this.previewLimit = Math.max(1, val);
    }
    addSelectedCell(val) {
        // Deduplicate by coordinates — Ctrl+clicking an already-selected cell
        // must not inflate the count or corrupt statistics.
        const alreadyPresent = this.selectedCells.some((c) => c.coordinates.rowIndex === val.coordinates.rowIndex &&
            c.coordinates.colIndex === val.coordinates.colIndex);
        if (!alreadyPresent) {
            this.selectedCells.push(val);
        }
    }
    setSelectedCells(val) {
        this.selectedCells = val;
    }
    setMouseOverCell(val) {
        this.mousedOverCell = val;
    }
    setQueryRunPromise(promise) {
        this.queryRunPromise = promise;
    }
    setIsQueryUsageViewerOpened(val) {
        this.isQueryUsageViewerOpened = val;
    }
    setExecutionError(val) {
        this.executionError = val;
    }
    setIsExecutionResultOverflowing(val) {
        this.isExecutionResultOverflowing = val;
    }
    updatePreviewLimitInConfig() {
        if (this.gridConfig) {
            this.gridConfig.previewLimit = this.previewLimit;
        }
    }
    getExecutionResultLimit = () => Math.min(this.queryBuilderState.fetchStructureState.implementation instanceof
        QueryBuilderTDSState &&
        this.queryBuilderState.fetchStructureState.implementation
            .resultSetModifierState.limit
        ? this.queryBuilderState.fetchStructureState.implementation
            .resultSetModifierState.limit
        : Number.MAX_SAFE_INTEGER, this.previewLimit);
    processExecutionResult = (result) => {
        this.setIsExecutionResultOverflowing(false);
        if (result instanceof TDSExecutionResult &&
            this.queryBuilderState.isQuerySupported) {
            const resultLimit = this.getExecutionResultLimit();
            if (result.result.rows.length > resultLimit) {
                this.setIsExecutionResultOverflowing(true);
                result.result.rows = result.result.rows.slice(0, resultLimit);
            }
        }
        this.setExecutionResult(result);
    };
    processWeightedColumnPairsMap(config) {
        if (config.weightedColumnPairs) {
            const wavgColumns = config.columns
                .filter((col) => col.aggFunc ===
                QueryBuilderDataGridCustomAggregationFunction.WAVG)
                .map((col) => col.colId);
            const weightedColumnPairsMap = new Map();
            config.weightedColumnPairs.forEach((wc) => {
                if (wc[0] && wc[1]) {
                    weightedColumnPairsMap.set(wc[0], wc[1]);
                }
            });
            for (const wavgCol of weightedColumnPairsMap.keys()) {
                if (!wavgColumns.includes(wavgCol)) {
                    weightedColumnPairsMap.delete(wavgCol);
                }
            }
            return weightedColumnPairsMap;
        }
        return undefined;
    }
    handlePreConfiguredGridConfig(config) {
        let newConfig;
        const weightedColumnPairsMap = this.processWeightedColumnPairsMap(config);
        if (weightedColumnPairsMap) {
            this.wavgAggregationState = new QueryBuilderResultWavgAggregationState();
            this.wavgAggregationState.weightedColumnIdPairs = weightedColumnPairsMap;
            newConfig = {
                ...config,
                weightedColumnPairs: weightedColumnPairsMap,
                columns: config.columns,
            };
        }
        else {
            newConfig = {
                ...config,
                columns: config.columns,
            };
        }
        if (config.previewLimit) {
            this.setPreviewLimit(config.previewLimit);
        }
        this.setGridConfig(newConfig);
    }
    getQueryGridConfig() {
        if (this.gridConfig) {
            return {
                ...this.gridConfig,
                columns: this.gridConfig.columns,
            };
        }
        return undefined;
    }
    get checkForStaleResults() {
        if (this.latestRunHashCode !== this.queryBuilderState.hashCode) {
            return true;
        }
        return false;
    }
    buildExecutionRawLambda(executionOptions) {
        let query;
        if (this.queryBuilderState.isQuerySupported) {
            const lambdaFunction = buildLambdaFunction(this.queryBuilderState, {
                isBuildingExecutionQuery: true,
                useTypedRelationFunctions: this.queryBuilderState.isFetchStructureTyped,
                ...executionOptions,
            });
            query = buildRawLambdaFromLambdaFunction(lambdaFunction, this.queryBuilderState.graphManagerState);
        }
        else {
            query = guaranteeNonNullable(this.queryBuilderState.unsupportedQueryState.rawLambda, 'Lambda is required to execute query');
            if (!this.queryBuilderState.isParameterSupportDisabled) {
                return getExecutionQueryFromRawLambda(query, this.queryBuilderState.parametersState.parameterStates, this.queryBuilderState.graphManagerState);
            }
        }
        return query;
    }
    *exportData(format) {
        try {
            this.exportState.inProgress();
            this.queryBuilderState.applicationStore.notificationService.notifySuccess(`Export ${format} will run in background`);
            const exportData = this.queryBuilderState.fetchStructureState.implementation.getExportDataInfo(format);
            const contentType = exportData.contentType;
            const serializationFormat = exportData.serializationFormat;
            const query = this.buildExecutionRawLambda({
                isExportingResult: true,
            });
            QueryBuilderTelemetryHelper.logEvent_ExportQueryDataLaunched(this.queryBuilderState.applicationStore.telemetryService);
            const result = (yield this.queryBuilderState.graphManagerState.graphManager.exportData(query, this.queryBuilderState.executionContextState.explicitMappingValue, this.queryBuilderState.executionContextState.explicitRuntimeValue, this.queryBuilderState.graphManagerState.graph, {
                serializationFormat,
                parameterValues: buildExecutionParameterValues(this.queryBuilderState.parametersState.parameterStates, this.queryBuilderState.graphManagerState),
                forceFromExpression: this.queryBuilderState.forceFromExpressionForExec,
                floatingExecutionElements: this.queryBuilderState.floatingExecutionElements,
            }, undefined, contentType));
            if (result.headers.get(V1_DELEGATED_EXPORT_HEADER) === 'true') {
                if (result.status === 200) {
                    this.exportState.pass();
                }
                else {
                    this.exportState.fail();
                }
                return;
            }
            const report = reportGraphAnalytics(this.queryBuilderState.graphManagerState.graph);
            downloadStream(result, `result.${getContentTypeFileExtension(contentType)}`, exportData.contentType)
                .then(() => {
                const reportWithState = Object.assign({}, report, this.queryBuilderState.getStateInfo());
                QueryBuilderTelemetryHelper.logEvent_ExportQueryDataSucceeded(this.queryBuilderState.applicationStore.telemetryService, reportWithState);
                this.exportState.pass();
            })
                .catch((error) => {
                assertErrorThrown(error);
                this.queryBuilderState.applicationStore.logService.error(LogEvent.create(GRAPH_MANAGER_EVENT.EXECUTION_FAILURE), error);
            });
        }
        catch (error) {
            this.exportState.fail();
            assertErrorThrown(error);
            this.queryBuilderState.applicationStore.logService.error(LogEvent.create(GRAPH_MANAGER_EVENT.EXECUTION_FAILURE), error);
            this.queryBuilderState.applicationStore.notificationService.notifyError(error);
            this.exportState.complete();
        }
    }
    *runQuery() {
        let promise;
        try {
            this.setIsRunningQuery(true);
            const currentHashCode = this.queryBuilderState.hashCode;
            if (this.queryBuilderState.requiresMappingForExecution) {
                guaranteeNonNullable(this.queryBuilderState.executionContextState.mapping, 'Mapping is required to execute query');
            }
            guaranteeNonNullable(this.queryBuilderState.executionContextState.runtimeValue, `Runtime is required to execute query`);
            const query = this.buildExecutionRawLambda({
                withDataOverflowCheck: true,
            });
            const parameterValues = buildExecutionParameterValues(this.queryBuilderState.parametersState.parameterStates, this.queryBuilderState.graphManagerState);
            QueryBuilderTelemetryHelper.logEvent_QueryRunLaunched(this.queryBuilderState.applicationStore.telemetryService);
            const stopWatch = new StopWatch();
            const report = reportGraphAnalytics(this.queryBuilderState.graphManagerState.graph);
            const contextstate = this.queryBuilderState.executionContextState;
            promise = this.queryBuilderState.graphManagerState.graphManager.runQuery(query, contextstate.explicitMappingValue, contextstate.explicitRuntimeValue, this.queryBuilderState.graphManagerState.graph, {
                parameterValues,
                convertUnsafeNumbersToString: true,
                preservedResponseHeadersList: [V1_ZIPKIN_TRACE_HEADER],
                tracingtags: { ...this.queryBuilderState.sourceInfo },
                forceFromExpression: this.queryBuilderState.forceFromExpressionForExec,
                floatingExecutionElements: this.queryBuilderState.floatingExecutionElements,
            });
            this.setQueryRunPromise(promise);
            const result = (yield promise);
            if (this.queryRunPromise === promise) {
                this.processExecutionResult(result.executionResult);
                if (result.executionTraceId) {
                    this.setExecutionTraceId(result.executionTraceId);
                }
                this.latestRunHashCode = currentHashCode;
                this.setExecutionDuration(stopWatch.elapsed);
                report.timings =
                    this.queryBuilderState.applicationStore.timeService.finalizeTimingsRecord(stopWatch, report.timings);
                const reportWithState = Object.assign({}, report, this.queryBuilderState.getStateInfo());
                QueryBuilderTelemetryHelper.logEvent_QueryRunSucceeded(this.queryBuilderState.applicationStore.telemetryService, reportWithState);
            }
        }
        catch (error) {
            // When user cancels the query by calling the cancelQuery api, it will throw an execution failure error.
            // For now, we don't want to notify users about this failure. Therefore we check to ensure the promise is still the same one.
            // When cancelled the query, we set the queryRunPromise as undefined.
            if (this.queryRunPromise === promise) {
                assertErrorThrown(error);
                this.queryBuilderState.applicationStore.logService.error(LogEvent.create(GRAPH_MANAGER_EVENT.EXECUTION_FAILURE), error);
                this.setExecutionError(error);
                if (error instanceof ExecutionError && error.executionTraceId) {
                    this.setExecutionTraceId(error.executionTraceId);
                }
            }
        }
        finally {
            this.setIsRunningQuery(false);
            this.pressedRunQuery.complete();
        }
    }
    *cancelQuery() {
        this.pressedRunQuery.complete();
        this.setIsRunningQuery(false);
        this.setQueryRunPromise(undefined);
        try {
            yield this.queryBuilderState.graphManagerState.graphManager.cancelUserExecutions(true);
        }
        catch (error) {
            // Don't notify users about success or failure
            this.queryBuilderState.applicationStore.logService.error(LogEvent.create(GRAPH_MANAGER_EVENT.EXECUTION_FAILURE), error);
        }
    }
    *generatePlan(debug) {
        try {
            this.isGeneratingPlan = true;
            if (this.queryBuilderState.requiresMappingForExecution) {
                guaranteeNonNullable(this.queryBuilderState.executionContextState.mapping, 'Mapping is required to execute query');
            }
            guaranteeNonNullable(this.queryBuilderState.executionContextState.runtimeValue, `Runtime is required to execute query`);
            const query = this.queryBuilderState.buildQuery();
            let rawPlan;
            const stopWatch = new StopWatch();
            const report = reportGraphAnalytics(this.queryBuilderState.graphManagerState.graph);
            const contextstate = this.queryBuilderState.executionContextState;
            if (debug) {
                QueryBuilderTelemetryHelper.logEvent_ExecutionPlanDebugLaunched(this.queryBuilderState.applicationStore.telemetryService);
                const debugResult = (yield this.queryBuilderState.graphManagerState.graphManager.debugExecutionPlanGeneration(query, contextstate.explicitMappingValue, contextstate.explicitRuntimeValue, this.queryBuilderState.graphManagerState.graph, {
                    forceFromExpression: this.queryBuilderState.forceFromExpressionForExec,
                    floatingExecutionElements: this.queryBuilderState.floatingExecutionElements,
                }, report));
                rawPlan = debugResult.plan;
                this.executionPlanState.setDebugText(debugResult.debug);
            }
            else {
                QueryBuilderTelemetryHelper.logEvent_ExecutionPlanGenerationLaunched(this.queryBuilderState.applicationStore.telemetryService);
                rawPlan =
                    (yield this.queryBuilderState.graphManagerState.graphManager.generateExecutionPlan(query, contextstate.explicitMappingValue, contextstate.explicitRuntimeValue, this.queryBuilderState.graphManagerState.graph, {
                        forceFromExpression: this.queryBuilderState.forceFromExpressionForExec,
                        floatingExecutionElements: this.queryBuilderState.floatingExecutionElements,
                    }, report));
            }
            stopWatch.record();
            try {
                this.executionPlanState.setRawPlan(rawPlan);
                const plan = this.queryBuilderState.graphManagerState.graphManager.buildExecutionPlan(rawPlan, this.queryBuilderState.graphManagerState.graph);
                this.executionPlanState.initialize(plan);
            }
            catch {
                // do nothing
            }
            stopWatch.record(QUERY_BUILDER_EVENT.BUILD_EXECUTION_PLAN__SUCCESS);
            // report
            report.timings =
                this.queryBuilderState.applicationStore.timeService.finalizeTimingsRecord(stopWatch, report.timings);
            const reportWithState = Object.assign({}, report, this.queryBuilderState.getStateInfo());
            if (debug) {
                QueryBuilderTelemetryHelper.logEvent_ExecutionPlanDebugSucceeded(this.queryBuilderState.applicationStore.telemetryService, reportWithState);
            }
            else {
                QueryBuilderTelemetryHelper.logEvent_ExecutionPlanGenerationSucceeded(this.queryBuilderState.applicationStore.telemetryService, reportWithState);
            }
        }
        catch (error) {
            assertErrorThrown(error);
            this.queryBuilderState.applicationStore.logService.error(LogEvent.create(GRAPH_MANAGER_EVENT.EXECUTION_FAILURE), error);
            this.queryBuilderState.applicationStore.notificationService.notifyError(error);
        }
        finally {
            this.isGeneratingPlan = false;
        }
    }
    *generateLineage() {
        if (this.isGeneratingLineage) {
            return;
        }
        try {
            this.isGeneratingLineage = true;
            const lambda = this.buildExecutionRawLambda();
            const lineageRawData = (yield this.queryBuilderState.graphManagerState.graphManager.generateLineage(lambda, this.queryBuilderState.executionContextState.explicitMappingValue, undefined, this.queryBuilderState.graphManagerState.graph, undefined));
            const lineageData = this.queryBuilderState.graphManagerState.graphManager.buildLineage(lineageRawData);
            this.lineageState.setLineageData(lineageData);
        }
        catch (error) {
            assertErrorThrown(error);
            this.queryBuilderState.applicationStore.logService.error(LogEvent.create(GRAPH_MANAGER_EVENT.LINEAGE_GENERATION_FAILURE), error);
            this.queryBuilderState.applicationStore.notificationService.notifyError(error);
        }
        finally {
            this.isGeneratingLineage = false;
        }
    }
}
//# sourceMappingURL=QueryBuilderResultState.js.map