package org.vps;

import java.lang.ref.WeakReference;
import java.util.LinkedList;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.TreeMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * Provides the service for storing / retrieving thread
 * local variables. You can override this class to override
 * the {@link #initVar()} method, to create custom code that
 * instantiates a new variable when it's not created for a
 * thread yet.
 * @since 1.5
 */
public class ThreadVarStore<K> {

    /**
     * Used to store all the instances of the
     * thread var stores, used during clean up
     */
    private static final List<WeakReference<ThreadVarStore<?>>> allStores =
        new LinkedList<WeakReference<ThreadVarStore<?>>>();

    /**
     * Synchronization object for the all stores list.
     */
    private static final Object allStoresLock = new Object();

    private static final Timer timer = new Timer(true);

    /**
     * Clean up internval, in milliseconds.
     * Default value is 10 min.
     */
    private static final long CLEAN_UP_INTERVAL = 600000L;

    static {

        timer.schedule(new TimerTask() {

            @Override
            public void run() {

                synchronized (allStoresLock) {

                    ListIterator<WeakReference<ThreadVarStore<?>>> li =
                        allStores.listIterator();

                    while (li.hasNext()) {

                        WeakReference<ThreadVarStore<?>> wr =
                            li.next();

                        ThreadVarStore<?> tvs = wr.get();
                        if (tvs == null) {
                            li.remove();
                            continue;
                        }

                        List<ThreadKey> forClear = null;
                        try {

                            tvs.readLock.lock();

                            for (ThreadKey tk : tvs.myMap.keySet()) {
                                Thread t = tk.threadRef.get();
                                if (t == null) {
                                    if (forClear == null) {
                                        forClear = new LinkedList<ThreadKey>();
                                    }
                                    forClear.add(tk);
                                }
                            }

                        } finally {
                            tvs.readLock.unlock();
                        }

                        if (forClear == null) {
                            continue;
                        }


                        try {
                            tvs.writeLock.lock();

                            for (ThreadKey tk : forClear) {
                                tvs.myMap.remove(tk);
                            }

                        } finally {
                            tvs.writeLock.unlock();
                        }

                    }

                }

            }

                },
                CLEAN_UP_INTERVAL, CLEAN_UP_INTERVAL);
    }

    /**
     * Internal map
     */
    private final Map<ThreadKey,K> myMap =
        new TreeMap<ThreadKey,K>();

    /**
     * Create an unfair read write lock, to allow writes
     * go through as soon as possible.
     */
    private final ReadWriteLock rwl =
        new ReentrantReadWriteLock(false);

    /**
     * Read lock for table access
     */
    private final Lock readLock = rwl.readLock();

    /**
     * Write lock for table access
     */
    private final Lock writeLock = rwl.writeLock();

    /**
     * Class used as the container of varaibles.
     */
    private Class<K> containerClass;

    /**
     * Creates a new isntance of this class, with the
     * specified variable type.
     * @param containerClass class to use when creating
     * new instances. This class should match the parameter
     * type, but can be <code>null</code> to prevent variables
     * from being created automatically.
     */
    public ThreadVarStore(Class<K> containerClass) {
        this.containerClass = containerClass;
        synchronized (allStoresLock) {
            allStores.add(new WeakReference<ThreadVarStore<?>>(this));
        }
    }

    /**
     * Returns the variable for the current thread.
     * If no variable has been created for this thread,
     * it is created.
     * @return the variable for the current thread.
     * @see #get(Thread)
     */
    public final K get() {
        return get(Thread.currentThread());
    }

    /**
     * Returns the variable for the specified thread.
     * If no container has been created for this thread,
     * it is created.
     * Note, if you are planning to modifying variables
     * that belong to a thread that is not current, you may need
     * to implement a customer synchronzation.
     * If you have extended this class, and the
     * {@link #initVar()} method returns <code>null</code>
     * value, it will not be stored in the map. If you have
     * passed <code>null</code> as the variable class, then
     * <code>null</code> will be returned as the variable,
     * and it will not be inserted into the internal map.
     * @param t thread to get the variable from
     * @return the variable for the specified thread
     */
    public final K get(Thread t) {

        ThreadKey id = new ThreadKey(t, false);

        try {
            readLock.lock();
            K val = myMap.get(id);
            if (val != null) {
                return val;
            }
        } finally {
            readLock.unlock();
        }

        try {
            writeLock.lock();
            // we have to do the look up again now,
            // as we have no guarantee that something
            // wasn't jammed into this map before we
            // got the write lock.
            K val = myMap.get(id);
            if (val != null) {
                return val;
            }
            val = initVar();

            id.setThread(t);

            if (val != null) {
                myMap.put(id, val);
            }

            return val;

        } finally {
            writeLock.unlock();
        }
    }

    /**
     * Removes the variable from the current thread.
     * This method should only be called to expedite the
     * release of the reference to the variable.
     */
    public final void remove() {
        remove(Thread.currentThread());
    }

    /**
     * Removes the variable from the specified thread.
     * This method should only be called to expedite the
     * release of the reference to the variable.
     * @param t thread to remove the variable from
     */
    public final void remove(Thread t) {
        ThreadKey id = new ThreadKey(t, false);
        try {
            writeLock.lock();
            myMap.remove(id);
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * Replaces the value of the variable for the current thread.
     * @param var new variable value
     * @see #set(Thread, Object)
     */
    public final void set(K var) {
        set(Thread.currentThread(), var);
    }

    /**
     * Replaces the value of the variable for the specified thread.
     * You should only call this method once. After assocaiting
     * the variable with the thread, you should then rather
     * modify the value of the variable internally rather than
     * resetting the variable value.<br>
     * If you are accessing the variables of other than current
     * threads, make sure you employ external synchronization to
     * not overwrite n already associated variable.
     * @param t thread to set the variable for
     * @param var variable value
     */
    public final void set(Thread t, K var) {

        ThreadKey id = new ThreadKey(t, true);

        try {
            writeLock.lock();
            myMap.put(id, var);
        } finally {
            writeLock.unlock();
        }

    }

    /**
     * Initializes the variable for a thread that doesn't
     * have it yet.
     * This implemntation creates a new instance of the
     * class that was passed to the contructor. Thus that
     * class must have a public no-arg constructor.
     * @return new instance of the variable, or <code>null</code>
     * if variable can not be instantiated.
     */
    public K initVar() {
        if (containerClass == null) {
            return null;
        }
        try {
            return containerClass.newInstance();
        } catch (InstantiationException e) {
            throw new RuntimeException(e);
        } catch (IllegalAccessException ignore) {
            return null;
        }
    }

    static class ThreadKey implements Comparable<ThreadKey> {

        private final long id;
        private WeakReference<Thread> threadRef;

        ThreadKey(Thread thread, boolean full) {
            id = thread.getId();
            if (full) {
                threadRef = new WeakReference<Thread>(thread);
            }
        }

        private void setThread(Thread thread) {
            if (threadRef == null) {
                threadRef = new WeakReference<Thread>(thread);
            }
        }

        @Override
        public int hashCode() {
            // stolen from java.lang.Long
            return (int)(id ^ (id >>> 32));
        }

        @Override
        public boolean equals(Object other) {
            return other instanceof ThreadKey &&
                    id == ((ThreadKey) other).id;
        }

        public int compareTo(ThreadKey other) {

            if (other.id == id) { return 0; }
            if (id > other.id) { return 1; }
            return -1;

        }


    }

}