/*
 * Copyright © 2020 by Data Virtuality GmbH
 */
package org.teiid.query.optimizer.relational.rules;

import org.teiid.core.util.StringUtil;
import org.teiid.query.analysis.AnalysisRecord;
import org.teiid.query.metadata.QueryMetadataInterface;
import org.teiid.query.optimizer.capabilities.CapabilitiesFinder;
import org.teiid.query.optimizer.relational.OptimizerRule;
import org.teiid.query.optimizer.relational.RuleStack;
import org.teiid.query.optimizer.relational.plantree.NodeConstants;
import org.teiid.query.optimizer.relational.plantree.NodeConstants.Info;
import org.teiid.query.optimizer.relational.plantree.NodeEditor;
import org.teiid.query.optimizer.relational.plantree.PlanNode;
import org.teiid.query.sql.lang.Unpivot;
import org.teiid.query.sql.symbol.Constant;
import org.teiid.query.sql.symbol.ElementSymbol;
import org.teiid.query.util.CommandContext;

import java.util.*;

/**
 * @author Roman Pedchenko
 * @date 2020.02.14
 */
public final class RuleReprojectUnpivot implements OptimizerRule {
    public PlanNode execute(PlanNode plan,
                            QueryMetadataInterface metadata,
                            CapabilitiesFinder capFinder,
                            RuleStack rules,
                            AnalysisRecord analysisRecord,
                            CommandContext context) {
        PlanNode unpivotNode = NodeEditor.findNodePreOrder(plan, NodeConstants.Types.UNPIVOT);
        if (unpivotNode == null)
            return plan;
        PlanNode outerProjNode = unpivotNode.getParent();
        if (outerProjNode.getType() != NodeConstants.Types.PROJECT)
            return plan;
        if (unpivotNode.getChildCount() != 1)
            return plan;
        PlanNode innerProjNode = unpivotNode.getFirstChild();
        if (innerProjNode.getType() != NodeConstants.Types.PROJECT)
            return plan;
        //noinspection unchecked
        List<ElementSymbol> outerProj = (List<ElementSymbol>) outerProjNode.getProperty(Info.PROJECT_COLS);
        //noinspection unchecked
        List<ElementSymbol> innerProj = (List<ElementSymbol>) innerProjNode.getProperty(Info.PROJECT_COLS);
        if (outerProj == null || outerProj.isEmpty() || innerProj == null || innerProj.isEmpty())
            return plan;
        final ElementSymbol col = (ElementSymbol) unpivotNode.getProperty(Info.UNPIVOT_COL);
        boolean hasFoldCol = Unpivot.colIndex(outerProj, col) >= 0;
        //noinspection unchecked
        final Map<Integer, String> foldCols = (Map<Integer, String>) unpivotNode.getProperty(Info.UNPIVOT_FOLD_COLS);
        final Map<Integer, String> newFoldCols = new HashMap<>();
        final Collection<String> values = new TreeSet<>(StringUtil.CASE_INSENSITIVE_ORDER);
        values.addAll(foldCols.values());
        int cnt = 0;
        for (Iterator<ElementSymbol> it = innerProj.iterator(); it.hasNext(); ) {
            ElementSymbol el = it.next();
            if (Unpivot.colIndex(outerProj, el) >= 0) {
                cnt++;
                continue;
            }
            final String foldCol = el.getShortName();
            if (hasFoldCol && values.contains(foldCol)) {
                newFoldCols.put(cnt, foldCol);
                cnt++;
                continue;
            }
            it.remove();
        }
        if (hasFoldCol)
            unpivotNode.setProperty(Info.UNPIVOT_FOLD_COLS, newFoldCols);
        if (innerProj.isEmpty())
            unpivotNode.setProperty(Info.PROJECT_COLS, Collections.singletonList(new Constant(1)));
        return plan;
    }

    public String toString() {
        return "ReprojectUnpivot";
    }
}
