/**
 * Copyright (C) 2014 Stratio (http://stratio.com)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.stratio.ingestion.source.redis;

import static com.stratio.ingestion.source.redis.RedisConstants.*;

import java.nio.charset.Charset;
import java.util.Map;

import org.apache.flume.Context;
import org.apache.flume.Event;
import org.apache.flume.EventDrivenSource;
import org.apache.flume.channel.ChannelProcessor;
import org.apache.flume.conf.Configurable;
import org.apache.flume.event.EventBuilder;
import org.apache.flume.source.AbstractSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Maps;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;
import redis.clients.jedis.JedisPubSub;
import redis.clients.jedis.exceptions.JedisConnectionException;

/**
 *
 * Redis Subscribe Source with connection Pool.
 *
 * Configuration parameters are:
 *
 * <p>
 * <ul>
 * <li><tt>host</tt> <em>(string)</em>: Redist host. Default: localhost.</li>
 * <li><tt>port</tt> <em>(integer)</em>: Redis port. Default: 6379.</li>
 * <li><tt>subscribe</tt> <em>(string)</em>: Channels to subscribe. Comma separated channels values.</li>
 * <li><tt>psubscribe</tt> <em>(string)</em>: Channels to subscribe with given pattern./li>
 * <li><tt></tt>charset</tt> <em>(string)</em>: Charset. Default: utf-8. </li>
 * <li><tt>pool.<property></tt></li>: Prefix for pool properties. Set whatever property you want to Jedis Pool.
 * </ul>
 * </p>
 *
 */
public class RedisSource extends AbstractSource implements Configurable, EventDrivenSource {
	
	private static final Logger log = LoggerFactory.getLogger(RedisSource.class);

    private ChannelProcessor channelProcessor;
	
    private JedisPool jedisPool;
    private String host;
    private Integer port;
    private String charset;
    private String [] channels;
    private String [] patterns;
    private Map<String, String> poolProps;

    boolean pattern = false;

	@Override
	public void configure(Context context) {

        host = context.getString(RedisConstants.CONF_HOST, DEFAULT_HOST);
        port = context.getInteger(RedisConstants.CONF_PORT, DEFAULT_PORT);
        charset = context.getString(RedisConstants.CONF_CHARSET, DEFAULT_CHARSET);
        String rawChannels = context.getString(CONF_CHANNELS);
        String rawPatterns = context.getString(CONF_PCHANNELS);
        if(null != rawChannels){
            channels = rawChannels.trim().split(",");
            pattern = false;
        } else if (null != rawPatterns){
            patterns = rawPatterns.trim().split(",");
            pattern = true;
        } else {
            throw new RuntimeException("You must set " + CONF_CHANNELS  + " or " + CONF_PCHANNELS + " property.");
        }

        poolProps = context.getSubProperties("pool.");

        log.info("Redis Source Configured");
	}
	
	@Override
	public synchronized void start() {
        super.start();

        channelProcessor = getChannelProcessor();

        init();
        log.info("Redis Connected. (host: " + host + ", port: " + String.valueOf(port) + ")");

        new SubscribeManager().run();
	}
	
	@Override
	public synchronized void stop() {
        super.stop();
        jedisPool.destroy();
	}

    private class SubscribeManager implements Runnable {

        Jedis subscriber;

        @Override
        public void run() {
            log.info("Subscribe Manager Thread is started.");
            subscriber = jedisPool.getResource();

            JedisPubSub jedisPubSub = new JedisPubSub() {
                @Override
                public void onMessage(String channel, String message) {
                    Event event = EventBuilder.withBody(message, Charset.forName(charset));
                    channelProcessor.processEvent(event);
                }

                @Override
                public void onPMessage(String pattern, String channel, String message) {
                    Map<String, String> headers = Maps.newHashMap();
                    headers.put("channel", channel);
                    Event event = EventBuilder.withBody(message, Charset.forName(charset), headers);
                    channelProcessor.processEvent(event);
                }

                @Override
                public void onSubscribe(String channel, int subscribedChannels) {
                    log.info("onSubscribe (Channel: " + channel + ")");
                }

                @Override
                public void onUnsubscribe(String channel, int subscribedChannels) {
                    log.info("onUnsubscribe (Channel: " + channel + ")");
                }

                @Override
                public void onPUnsubscribe(String Pattern, int subscribedChannels) {
                    log.info("onPUnSubscribe (Pattern: " + Pattern + ")");
                }

                @Override
                public void onPSubscribe(String pattern, int subscribedChannels) {
                    log.info("onPSubscribe (Pattern: " + pattern + ")");
                }
            };

            while (true) {
                if (pattern) {
                    for (String pattern : patterns) {
                        log.info("Jedis is going to subscribe to pattern: " + pattern);
                    }
                    try {
                        subscriber.psubscribe(jedisPubSub, patterns);
                    } catch (JedisConnectionException ex) {
                        jedisPool.returnBrokenResource(subscriber);
                        subscriber = jedisPool.getResource();
                    }
                } else {
                    for (String channel : channels) {
                        log.info("Jedis is going to subscribe to channel: " + channel);
                    }
                    try {
                        subscriber.subscribe(jedisPubSub, channels);
                    } catch (JedisConnectionException ex) {
                        jedisPool.returnBrokenResource(subscriber);
                        subscriber = jedisPool.getResource();
                        log.info("Active connections: " + jedisPool.getNumActive());
                    }
                }
            }
        }
    }

    private void init() {
        try {
            JedisPoolConfig poolConfig = new JedisPoolConfig();
            String prop;

            if((prop = poolProps.get(CONF_TESTONBORROW)) == null){
                poolConfig.setTestOnBorrow(DEFAULT_TESTONBORROW);
            } else {
                log.info("Setting testOnBorrow property to " + prop);
                poolConfig.setTestOnBorrow(Boolean.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_MAXTOTAL)) == null){
                poolConfig.setMaxTotal(DEFAULT_MAXTOTAL);
            } else {
                log.info("Setting maxTotal property to " + prop);
                poolConfig.setMaxTotal(Integer.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_MAXIDLE)) == null){
                poolConfig.setMaxTotal(DEFAULT_MAXIDLE);
            } else {
                log.info("Setting maxIdle property to " + prop);
                poolConfig.setMaxIdle(Integer.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_MINIDLE)) == null){
                poolConfig.setMinIdle(DEFAULT_MINIDLE);
            } else {
                log.info("Setting minIdle property to " + prop);
                poolConfig.setMinIdle(Integer.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_MAXWAITINMILLIS)) == null){
                poolConfig.setMaxWaitMillis(DEFAULT_MAXWAITINMILLIS);
            } else {
                log.info("Setting maxWaitInMillis property to " + prop);
                poolConfig.setMaxWaitMillis(Integer.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_TESTWHILEIDLE)) == null){
                poolConfig.setTestWhileIdle(DEFAULT_TESTWHILEIDLE);
            } else {
                log.info("Setting testWhileIdle property to " + prop);
                poolConfig.setTestWhileIdle(Boolean.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_TESTONRETURN)) == null){
                poolConfig.setTestOnReturn(DEFAULT_TESTONRETURN);
            } else {
                log.info("Setting testOnReturn property to " + prop);
                poolConfig.setTestOnReturn(Boolean.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_MINEVICTABLEIDLETIMEINMILLIS)) == null){
                poolConfig.setMinEvictableIdleTimeMillis(DEFAULT_MINEVICTABLEIDLETIMEINMILLIS);
            } else {
                log.info("Setting minEvictableIdleTimeInMillis property to " + prop);
                poolConfig.setMinEvictableIdleTimeMillis(Integer.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_TIMEBETWEETNEVICTIONRUNSMILLIS)) == null){
                poolConfig.setTimeBetweenEvictionRunsMillis(DEFAULT_TIMEBETWEETNEVICTIONRUNSMILLIS);
            } else {
                log.info("Setting timeBetweenEvictionRunMillis property to " + prop);
                poolConfig.setTimeBetweenEvictionRunsMillis(Integer.valueOf(prop));
            }
            if((prop = poolProps.get(CONF_NUMTESTSPEREVICTIONRUN)) == null){
                poolConfig.setNumTestsPerEvictionRun(DEFAULT_NUMTESTSPEREVICTIONRUN);
            } else {
                log.info("Setting numTestsPerEvictionRun property to " + prop);
                poolConfig.setNumTestsPerEvictionRun(Integer.valueOf(prop));
            }

            // create JEDIS pool
            this.jedisPool = new JedisPool(poolConfig, host, port);

        } catch (Exception e) {
            e.printStackTrace();
            System.exit(-1);
        }
    }
}
