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

import org.teiid.common.buffer.TupleBatch;
import org.teiid.common.buffer.TupleBuffer;
import org.teiid.core.TeiidComponentException;
import org.teiid.core.TeiidProcessingException;
import org.teiid.query.sql.LanguageObject;
import org.teiid.query.sql.lang.Unpivot;
import org.teiid.query.sql.symbol.ElementSymbol;
import org.teiid.query.sql.symbol.Expression;

import java.util.*;

/**
 * @author Roman Pedchenko
 * @date 2020.01.25
 */
public class UnpivotNode extends SubqueryAwareRelationalNode {
    private static final Object POS_CODE = new Object();
    private static final Object POS_VALUE = new Object();

    private Map<Integer, String> foldColumns;
    private ElementSymbol col;
    private ElementSymbol forCol;
    private List<ElementSymbol> innerProjection;
    private List<Object> projection;
    private TupleBatch currentBatch;
    private int currentRow = 1;

    protected UnpivotNode() {
        super();
    }

    public UnpivotNode(int nodeID, Map<Integer, String> foldColumns, ElementSymbol col, ElementSymbol forCol, List<ElementSymbol> innerProjection) {
        super(nodeID);
        this.foldColumns = foldColumns;
        this.col = col;
        this.forCol = forCol;
        this.innerProjection = innerProjection;
    }

    public void reset() {
        super.reset();

        currentBatch = null;
        currentRow = 1;
    }

    @Override
    public void addChild(RelationalNode child) {
        super.addChild(child);
        init();
    }

    void init() {
        final List<? extends Expression> output = getOutputElements();
        projection = new ArrayList<>(output.size());
        List<? extends Expression> workProjection = innerProjection;
        if (this.getChildren().length > 0)
            workProjection = this.getChildren()[0].getElements();
        for (Expression el : output) {
            final ElementSymbol el1 = (ElementSymbol) el;
            if (Unpivot.sameCol(forCol, el1))
                projection.add(POS_CODE);
            else if (Unpivot.sameCol(col, el1))
                projection.add(POS_VALUE);
            else
                projection.add(new Position(Unpivot.colIndex(workProjection, el1)));
        }
    }

    public TupleBatch nextBatchDirect() throws TeiidComponentException, TeiidProcessingException {
        if (currentBatch == null) {
            if (this.getChildren()[0] == null) {
                currentBatch = new TupleBatch(1, new List[]{Collections.emptyList()});
                currentBatch.setTerminationFlag(true);
            } else {
                currentBatch = this.getChildren()[0].nextBatch();
            }
        }

        while (currentRow <= currentBatch.getEndRow() && !isBatchFull()) {
            List<?> tuple = currentBatch.getTuple(currentRow);

            for (Map.Entry<Integer, String> foldEntry : foldColumns.entrySet()) {
                List<Object> resultTuple = new ArrayList<>(projection.size());

                for (Object proj : projection) {
                    Object value = null;
                    if (proj instanceof Position)
                        value = ((Position) proj).get(tuple);
                    else if (proj == POS_CODE)
                        value = foldEntry.getValue();
                    else if (proj == POS_VALUE)
                        value = tuple.get(foldEntry.getKey());
                    resultTuple.add(value);
                }

                addBatchRow(resultTuple);
            }

            currentRow++;
        }

        if (currentRow > currentBatch.getEndRow()) {
            if (currentBatch.getTerminationFlag()) {
                terminateBatches();
            }
            currentBatch = null;
        }

        return pullBatch();
    }

    public Object clone() {
        UnpivotNode clonedNode = new UnpivotNode();
        this.copyTo(clonedNode);
        return clonedNode;
    }

    protected void copyTo(UnpivotNode target) {
        super.copyTo(target);
        target.foldColumns = this.foldColumns;
        target.col = this.col;
        target.forCol = this.forCol;
        target.innerProjection = this.innerProjection;
        target.projection = this.projection;
    }

    @Override
    public Collection<? extends LanguageObject> getObjects() {
        return Collections.emptyList();
    }

    @Override
    public boolean hasBuffer(boolean requireFinal) {
        return false;
    }

    @Override
    public TupleBuffer getBuffer(int maxRows) throws TeiidComponentException, TeiidProcessingException {
        return this.getChildren()[0].getBuffer(maxRows);
    }

    private static class Position {
        private final int pos;

        public Position(int pos) {
            this.pos = pos;
        }

        public Object get(List<?> tuple) {
            return tuple.get(pos);
        }
    }
}
