package simpledb;
import java.util.*;
import java.util.Queue;

/**
 * The Join operator implements the relational join operation.
 */
public class Join extends AbstractDbIterator {
    private JoinPredicate jp;
    private DbIterator child1;
    private DbIterator child2;
    private Tuple currOuterTuple; // used only by brute-force join
    HashMap<Field, ArrayList<Tuple>> fieldTwoToTuples; // used only by equality join
    Queue<Tuple> tupleQueue; // used only by equality join

    /**
     * Constructor.  Accepts to children to join and the predicate
     * to join them on
     *
     * @param p The predicate to use to join the children
     * @param child1 Iterator for the left(outer) relation to join
     * @param child2 Iterator for the right(inner) relation to join
     */
    public Join(JoinPredicate p, DbIterator child1, DbIterator child2) {
        this.jp = p;
        this.child1 = child1;
        this.child2 = child2;

        if(p.getOp() == Predicate.Op.EQUALS) {
            fieldTwoToTuples = new HashMap<Field, ArrayList<Tuple>>();
            tupleQueue = new LinkedList<Tuple>();
        }
    }

    /**
     * @see simpledb.TupleDesc#combine(TupleDesc, TupleDesc) for possible implementation logic.
     */
    public TupleDesc getTupleDesc() {
        return TupleDesc.combine(child1.getTupleDesc(), child2.getTupleDesc());
    }

    public void open()
        throws DbException, NoSuchElementException, TransactionAbortedException {
        child1.open();
        child2.open();

        // for equality joins, set up hash table to do hash-joins
        if(jp.getOp() == Predicate.Op.EQUALS) {
            int fieldTwoIndex = jp.getFieldTwo();
            while(child2.hasNext()) {
                Tuple t = child2.next();
                Field fieldTwo = t.getField(fieldTwoIndex);
                if(!fieldTwoToTuples.containsKey(fieldTwo))
                    fieldTwoToTuples.put(fieldTwo, new ArrayList<Tuple>());
                fieldTwoToTuples.get(fieldTwo).add(t);
            }
        } else {
            if(child1.hasNext())
                currOuterTuple = child1.next();
        }
    }

    public void close() {
        currOuterTuple = null;
        child1.close();
        child2.close();
        super.close();
    }

    public void rewind() throws DbException, TransactionAbortedException {
        currOuterTuple = null;
        child1.rewind();
        child2.rewind();
        if(jp.getOp() != Predicate.Op.EQUALS && child1.hasNext())
            currOuterTuple = child1.next();
    }


    private Tuple readNextBruteForce() throws TransactionAbortedException, DbException {
        while(currOuterTuple != null) {
            while(child2.hasNext()) {
                Tuple rightTuple = child2.next();
                if(jp.filter(currOuterTuple, rightTuple))
                    return Tuple.combine(currOuterTuple, rightTuple);
            }
            if(!child1.hasNext())
                return null;
            currOuterTuple = child1.next();
            child2.rewind();
        }
        return null;
    }

    private Tuple readNextHashJoin() throws TransactionAbortedException, DbException {
       if(!tupleQueue.isEmpty())
           return tupleQueue.poll();

       // keep looking until you find a scenario where field1 == field2
       int fieldOneIndex = jp.getFieldOne();
       while(child1.hasNext()) {
           Tuple childOneTuple = child1.next();
           Field fieldOne = childOneTuple.getField(fieldOneIndex);
           if(!fieldTwoToTuples.containsKey(fieldOne))
               continue;

           // add all matches that pass filter into the queue and then return first one if possible
           ArrayList<Tuple> childTwoTuples = fieldTwoToTuples.get(fieldOne);
           for(Tuple childTwoTuple : childTwoTuples) {
               if(jp.filter(childOneTuple, childTwoTuple))
                   tupleQueue.add(Tuple.combine(childOneTuple, childTwoTuple));
           }

           if(!tupleQueue.isEmpty())
               return tupleQueue.poll();
       }

       if(!tupleQueue.isEmpty())
           return tupleQueue.poll();

       return null;
    }

    /**
     * Returns the next tuple generated by the join, or null if there are no more tuples.
     * Logically, this is the next tuple in r1 cross r2 that satisfies the join
     * predicate.  There are many possible implementations; the simplest is a
     * nested loops join.
     * <p>
     * Note that the tuples returned from this particular implementation of
     * Join are simply the concatenation of joining tuples from the left and
     * right relation. Therefore, if an equality predicate is used
     * there will be two copies of the join attribute
     * in the results.  (Removing such duplicate columns can be done with an
     * additional projection operator if needed.)
     * <p>
     * For example, if one tuple is {1,2,3} and the other tuple is {1,5,6},
     * joined on equality of the first column, then this returns {1,2,3,1,5,6}.
     *
     * @return The next matching tuple.
     * @see JoinPredicate#filter
     */
    protected Tuple readNext() throws TransactionAbortedException, DbException {
        if(jp.getOp() == Predicate.Op.EQUALS)
            return readNextHashJoin();
        else
            return readNextBruteForce();
    }
}
