package edu.caltech.cs2.datastructures;

import edu.caltech.cs2.interfaces.ICollection;
import edu.caltech.cs2.interfaces.IQueue;
import edu.caltech.cs2.interfaces.IDictionary;

import java.util.Iterator;

public class BSTDictionary<K extends Comparable<? super K>, V>
        implements IDictionary<K, V> {

    protected BSTNode<K, V> root;
    protected int size;

    /**
     * Class representing an individual node in the Binary Search Tree
     */
    protected static class BSTNode<K, V> {
        public final K key;
        public V value;

        public BSTNode<K, V> left;
        public BSTNode<K, V> right;

        /**
         * Constructor initializes this node's key, value, and children
         */
        public BSTNode(K key, V value) {
            this.key = key;
            this.value = value;
            this.left = null;
            this.right = null;
        }

        public BSTNode(BSTNode<K, V> o) {
            this.key = o.key;
            this.value = o.value;
            this.left = o.left;
            this.right = o.right;
        }

        public boolean isLeaf() {
            return this.left == null && this.right == null;
        }

        public boolean hasBothChildren() {
            return this.left != null && this.right != null;
        }
    }

    /**
     * Initializes an empty Binary Search Tree
     */
    public BSTDictionary() {
        this.root = null;
        this.size = 0;
    }

    public V getHelper(BSTNode<K, V> currNode, K key){
        if (currNode == null){
            return null;
        }else if (currNode.key.equals(key)){
            return currNode.value;
//        } else if (currNode.isLeaf()){
//            return null;
        } else if (currNode.key.compareTo(key) < 0){
            return getHelper(currNode.right, key);
        } else if (currNode.key.compareTo(key) > 0){
            return getHelper(currNode.left, key);
        }
        return null;
    }

    @Override
    public V get(K key) {

        if (root == null){
            return null;
        }
        return getHelper(root, key);

    }

    @Override
    public V remove(K key) {
        if (get(key) == null){
            return null;
        }
        size--;
        V returnVal = get(key);
        root = removeHelper(root, key);
        return returnVal;

    }

    private BSTNode<K, V> removeHelper(BSTNode<K, V> node, K key){
        if (node.key.equals(key)){
            if (node.isLeaf()){
                return null;
            } else if ((node.right != null) && (node.left == null)){
                return node.right;
            } else if ((node.right == null) && (node.left != null)){
                return node.left;
            } else {
                BSTNode<K, V> minNode = findMin(node);
                remove(minNode.key);
                this.size++;
                minNode.right = node.right;
                minNode.left = node.left;
                return minNode;
            }
        } else if (node.key.compareTo(key) < 0){
            node.right = removeHelper(node.right, key);
        } else if (node.key.compareTo(key) < 0){
            node.left = removeHelper(node.left, key);
        }
        return node;
    }

    private BSTNode<K, V> findMin(BSTNode<K, V> node){
        if (node.left == null){
            return node;
        } else {
            return findMin(node.left);
        }
    }


    @Override
    public V put(K key, V value) {
        if (root == null){
            root = new BSTNode<K, V>(key, value);
            size ++;
            return null;
        } else {

            V oldVal = get(key);
            root = putHelper(root, key, value);
            return oldVal;
        }

    }

    private BSTNode<K, V> putHelper(BSTNode<K, V> currNode, K key, V val){
        BSTNode<K, V> returnVal = null;
        if (currNode == null){
            size++;
            return new BSTNode<K, V>(key, val);

        }
        else if (currNode.key.compareTo(key) > 0){
            currNode.left = putHelper(currNode.left, key, val);
        }
        else if (currNode.key.compareTo(key) < 0){
            currNode.right = putHelper(currNode.right, key, val);
        }
        else if (currNode.key.equals(key)){
            returnVal = currNode;
            currNode.value = val;
            return returnVal;
        }

        return currNode;
    }

//    private BSTNode<K, V> putHelper(BSTNode<K, V> currNode, K key, V value){
//        if (currNode == null){
//            currNode = new BSTNode<K, V>(key, value);
//        } else if (currNode.key.compareTo(key) > 0){
//            currNode.left = putHelper(currNode.left, key, value);
//        } else if (currNode.key.compareTo(key) < 0){
//            currNode.right = putHelper(currNode.right, key, value);
//        } else if (currNode.key.equals(key)){
//            currNode.value = value;
//        }
//        return currNode;
//    }

    @Override
    public boolean containsKey(K key) {

        return (get(key) != null);
    }

    @Override
    public boolean containsValue(V value) {
        if (this.values() == null){
            return false;
        }
        return this.values().contains(value);

    }

    /**
     * @return number of nodes in the BST
     */
    @Override
    public int size() {
        return this.size;
    }

    @Override
    public ICollection<K> keys() {
        ICollection<K> keys = new LinkedDeque<K>();
        keysHelper(root, keys);
        return keys;
    }

    private void keysHelper(BSTNode<K, V> node, ICollection<K> keys){
        if (node != null){
            keys.add(node.key);
            if (node.left != null){
                keysHelper(node.left, keys);
            }
            if (node.right != null){
                keysHelper(node.right, keys);
            }
        }

    }


    @Override
    public ICollection<V> values() {
        ICollection<V> values = new LinkedDeque<V>();
//        return valuesHelper(root, values);
        if (keys() == null){
            return null;
        }
        for (K key : keys()){
            values.add(get(key));
        }
        return values;
    }

//    private ICollection<V> valuesHelper(BSTNode<K, V> node, ICollection<V> values){
//
//        if (node != null){
//            values.add(node.value);
//            if (node.left != null){
//                valuesHelper(node.left, values);
//            }
//            if (node.right != null){
//                valuesHelper(node.right, values);
//            }
//        }
//
//        return values;
//    }

    /**
     * Implementation of an iterator over the BST
     */

    @Override
    public Iterator<K> iterator() {
        return keys().iterator();
    }

    @Override
    public String toString() {
        if (this.root == null) {
            return "{}";
        }

        StringBuilder contents = new StringBuilder();

        IQueue<BSTNode<K, V>> nodes = new ArrayDeque<>();
        BSTNode<K, V> current = this.root;
        while (current != null) {
            contents.append(current.key + ": " + current.value + ", ");

            if (current.left != null) {
                nodes.enqueue(current.left);
            }
            if (current.right != null) {
                nodes.enqueue(current.right);
            }

            current = nodes.dequeue();
        }

        return "{" + contents.toString().substring(0, contents.length() - 2) + "}";
    }
}