/**
 * Copyright (c) 2012, The University of Southampton and the individual contributors.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 *   *  Redistributions of source code must retain the above copyright notice,
 *      this list of conditions and the following disclaimer.
 *
 *   *  Redistributions in binary form must reproduce the above copyright notice,
 *      this list of conditions and the following disclaimer in the documentation
 *      and/or other materials provided with the distribution.
 *
 *   *  Neither the name of the University of Southampton nor the names of its
 *      contributors may be used to endorse or promote products derived from this
 *      software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
package backtype.storm.spout;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;

import net.lag.kestrel.thrift.Item;

import org.apache.log4j.Logger;
import org.apache.thrift7.TException;

import backtype.storm.Config;
import backtype.storm.task.TopologyContext;
import backtype.storm.topology.OutputFieldsDeclarer;
import backtype.storm.topology.base.BaseRichSpout;
import backtype.storm.tuple.Fields;
import backtype.storm.utils.Utils;

/**
 * This spout can be used to consume messages in a reliable way from a cluster
 * of Kestrel servers. It is recommended that you set the parallelism hint to a
 * multiple of the number of Kestrel servers, otherwise the read load will be
 * higher on some Kestrel servers than others.
 */
@SuppressWarnings("serial")
public class KestrelThriftSpout extends BaseRichSpout {
        public static Logger LOG = Logger.getLogger(KestrelThriftSpout.class);

        public static final long BLACKLIST_TIME_MS = 1000 * 60;
        public static final int BATCH_SIZE = 4000;

        private List<String> _hosts = null;
        private int _port = -1;
        private String _queueName = null;
        private SpoutOutputCollector _collector;
        private Scheme _scheme;

        private List<KestrelClientInfo> _kestrels;
        private int _emitIndex;

        private Queue<EmitItem> _emitBuffer = new LinkedList<EmitItem>();

        private class EmitItem {
                public KestrelSourceId sourceId;
                public List<Object> tuple;

                public EmitItem(List<Object> tuple, KestrelSourceId sourceId) {
                        this.tuple = tuple;
                        this.sourceId = sourceId;
                }
        }

        private static class KestrelSourceId {
                public KestrelSourceId(int index, long id) {
                        this.index = index;
                        this.id = id;
                }

                int index;
                long id;
        }

        private static class KestrelClientInfo {
                public Long blacklistTillTimeMs;
                public String host;
                public int port;

                private KestrelThriftClient client;

                public KestrelClientInfo(String host, int port) {
                        this.host = host;
                        this.port = port;
                        this.blacklistTillTimeMs = 0L;
                        this.client = null;
                }

                public KestrelThriftClient getValidClient() throws TException {
                        if (this.client == null) { // If client was blacklisted, remake it.
                                LOG.info("Attempting reconnect to kestrel " + this.host + ":" + this.port);
                                this.client = new KestrelThriftClient(this.host, this.port);
                        }
                        return this.client;
                }

                public void closeClient() {
                        if (this.client != null) {
                                this.client.close();
                                this.client = null;
                        }
                }
        }

        public KestrelThriftSpout(List<String> hosts, int port, String queueName, Scheme scheme) {
                if (hosts.isEmpty()) {
                        throw new IllegalArgumentException("Must configure at least one host");
                }
                _port = port;
                _hosts = hosts;
                _queueName = queueName;
                _scheme = scheme;
        }

        public KestrelThriftSpout(String hostname, int port, String queueName, Scheme scheme) {
                this(Arrays.asList(hostname), port, queueName, scheme);
        }

        public KestrelThriftSpout(String hostname, int port, String queueName) {
                this(hostname, port, queueName, new RawScheme());
        }

        public KestrelThriftSpout(List<String> hosts, int port, String queueName) {
                this(hosts, port, queueName, new RawScheme());
        }

        public Fields getOutputFields() {
                return _scheme.getOutputFields();
        }

        int _messageTimeoutMillis;

        @SuppressWarnings("rawtypes")
        @Override
        public void open(Map conf, TopologyContext context, SpoutOutputCollector collector) {
                // TODO: should switch this to maxTopologyMessageTimeout
                final Number timeout = (Number) conf.get(Config.TOPOLOGY_MESSAGE_TIMEOUT_SECS);
                _messageTimeoutMillis = 1000 * timeout.intValue();
                _collector = collector;
                _emitIndex = 0;
                _kestrels = new ArrayList<KestrelClientInfo>();
                final int numTasks = context.getComponentTasks(context.getThisComponentId()).size();
                final int myIndex = context.getThisTaskIndex();
                final int numHosts = _hosts.size();
                if (numTasks < numHosts) {
                        for (final String host : _hosts) {
                                _kestrels.add(new KestrelClientInfo(host, _port));
                        }
                } else {
                        final String host = _hosts.get(myIndex % numHosts);
                        _kestrels.add(new KestrelClientInfo(host, _port));
                }
        }

        @Override
        public void close() {
                for (final KestrelClientInfo info : _kestrels)
                        info.closeClient();

                // Closing the client connection causes all the open reliable reads to
                // be aborted.
                // Thus, clear our local buffer of these reliable reads.
                _emitBuffer.clear();

                _kestrels.clear();
        }

        public boolean bufferKestrelGet(int index) {
                assert _emitBuffer.size() == 0; // JTODO

                final KestrelClientInfo info = _kestrels.get(index);

                final long now = System.currentTimeMillis();
                if (now > info.blacklistTillTimeMs) {
                        List<Item> items = null;
                        try {
                                items = info.getValidClient().get(_queueName, BATCH_SIZE, 0, _messageTimeoutMillis);
                                // if(items.size()!=0)LOG.debug("Got this many items: " +
                                // items.size());
                        } catch (final TException e) {
                                LOG.error("Error reading from client: " + e.getMessage());
                                blacklist(info, e);
                                return false;
                        }

                        assert items.size() <= BATCH_SIZE;
                        // LOG.info("Kestrel batch get fetched " + items.size() +
                        // " items. (batchSize= " + BATCH_SIZE +
                        // " queueName=" + _queueName + ", index=" + index + ", host=" +
                        // info.host + ")");

                        final HashSet<Long> toAck = new HashSet<Long>();

                        for (final Item item : items) {
                                final List<Object> retItems = _scheme.deserialize(item.get_data());

                                if (retItems != null) {
                                        final EmitItem emitItem = new EmitItem(retItems, new KestrelSourceId(index, item.get_id()));

                                        if (!_emitBuffer.offer(emitItem)) {
                                                throw new RuntimeException("KestrelThriftSpout's Internal Buffer Enqeueue Failed.");
                                        }
                                } else {
                                        toAck.add(item.get_id());
                                }
                        }

                        if (toAck.size() > 0) {
                                try {
                                        info.client.confirm(_queueName, toAck);
                                } catch (final TException e) {
                                        blacklist(info, e);
                                }
                        }

                        if (items.size() > 0)
                                return true;
                }
                return false;
        }

        public void tryEachKestrelUntilBufferFilled() {
                for (int i = 0; i < _kestrels.size(); i++) {
                        final int index = (_emitIndex + i) % _kestrels.size();
                        if (bufferKestrelGet(index)) {
                                _emitIndex = index;
                                break;
                        }
                }
                _emitIndex = (_emitIndex + 1) % _kestrels.size();
        }

        int countTriples = 1;
        int lastEmit = countTriples;
        int emptyIterations = 0;

        @Override
        public void nextTuple() {
                if (_emitBuffer.isEmpty())
                        tryEachKestrelUntilBufferFilled();
                if (countTriples % 1000 == 0 && countTriples != lastEmit) {
                        LOG.debug("Number of triples emitted: " + countTriples);
                        LOG.debug("Number of empty iterations: " + emptyIterations);
                        emptyIterations = 0;
                        lastEmit = countTriples;
                }
                final EmitItem item = _emitBuffer.poll();
                if (item != null) {
                        countTriples += 1;
                        _collector.emit(item.tuple, item.sourceId);
                } else { // If buffer is still empty here, then every kestrel Q is also
                                        // empty.
                        emptyIterations++;
                        Utils.sleep(10);
                }
        }

        private void blacklist(KestrelClientInfo info, Throwable t) {

                // this case can happen when it fails to connect to Kestrel (and so
                // never stores the connection)
                info.closeClient();
                info.blacklistTillTimeMs = System.currentTimeMillis() + BLACKLIST_TIME_MS;

                final int index = _kestrels.indexOf(info);

                // we just closed the connection, so all open reliable reads will be
                // aborted. empty buffers.
                for (final Iterator<EmitItem> i = _emitBuffer.iterator(); i.hasNext();) {
                        final EmitItem item = i.next();
                        if (item.sourceId.index == index)
                                i.remove();
                }
        }

        @Override
        public void ack(Object msgId) {
                final KestrelSourceId sourceId = (KestrelSourceId) msgId;
                final KestrelClientInfo info = _kestrels.get(sourceId.index);

                // if the transaction didn't exist, it just returns false. so this code
                // works
                // even if client gets blacklisted, disconnects, and kestrel puts the
                // item
                // back on the queue
                try {
                        if (info.client != null) {
                                final HashSet<Long> xids = new HashSet<Long>();
                                xids.add(sourceId.id);
                                info.client.confirm(_queueName, xids);
                        }
                } catch (final TException e) {
                        blacklist(info, e);
                }
        }

        @Override
        public void fail(Object msgId) {
                final KestrelSourceId sourceId = (KestrelSourceId) msgId;
                final KestrelClientInfo info = _kestrels.get(sourceId.index);

                // see not above about why this works with blacklisting strategy
                try {
                        if (info.client != null) {
                                final HashSet<Long> xids = new HashSet<Long>();
                                xids.add(sourceId.id);
                                info.client.abort(_queueName, xids);
                        }
                } catch (final TException e) {
                        blacklist(info, e);
                }
        }

        @Override
        public void declareOutputFields(OutputFieldsDeclarer declarer) {
                declarer.declare(getOutputFields());
        }
}