package org.apache.kafka.connect.mongodb;

import com.mongodb.CursorType;
import com.mongodb.MongoClient;
import com.mongodb.MongoClientURI;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.Filters;
import com.mongodb.client.model.Projections;

import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.utils.StringUtils;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.bson.types.BSONTimestamp;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.ConcurrentLinkedQueue;

/**
 * Reads mutation from a mongodb database
 *
 * @author Andrea Patelli
 */
public class DatabaseReader implements Runnable {
    private final Logger log = LoggerFactory.getLogger(DatabaseReader.class);
    private final String host;
    private final Integer port;
    private final String uri;
    private final String db;
    private final Integer batchSize;
    private final String start;
    private MongoClient mongoClient;
    private int page = 0;

    private ConcurrentLinkedQueue<Document> messages;

    private MongoCollection<Document> oplog;
    private Bson query;
    
    public DatabaseReader(String uri, String db, String start, Integer batchSize, ConcurrentLinkedQueue<Document> messages) {
        this.uri = uri;
        this.host = null;
        this.port = null;
        this.batchSize = batchSize;
        this.db = db;
        this.start = start;
        this.messages = messages;
        try {
            init();
        } catch (ConnectException e) {
            throw e;
        }
    }

    public DatabaseReader(String host, Integer port, String db, String start, Integer batchSize, ConcurrentLinkedQueue<Document> messages) {
        this.uri = null;
        this.host = host;
        this.port = port;
        this.batchSize = batchSize;
        this.db = db;
        this.start = start;
        this.messages = messages;
        try {
            init();
        } catch (ConnectException e) {
            throw e;
        }
    }

    @Override
    public void run() {
    	while(true){
    		if(messages.isEmpty()){
		        log.debug("Query starting in page {}.", page);
		        final FindIterable<Document> documents = find(page);
		        try {
		            for (Document document : documents) {
		                log.trace(document.toString());
		                messages.add(document);
		            }
		        } catch(Exception e) {
		            log.error("Closed connection", e);
		        }
		        page++;
    		}
    		else{
    			try {
					Thread.sleep(1000);
				} catch (InterruptedException e) {
					log.error(e.getMessage(), e);
				}
    		}
    	}
    }
    
    private FindIterable<Document> find(int page){
        final FindIterable<Document> documents = oplog
                .find(query)
                .sort(new Document("$natural", 1))
                .skip(page * batchSize)
                .limit(batchSize)
                .projection(Projections.include("ts", "op", "ns", "o"))
                .cursorType(CursorType.TailableAwait);
        return documents;
    }
    
    @Override
    public void finalize(){
    	if(mongoClient != null){
    		mongoClient.close();
    	}
    }
    
    private MongoClient createMongoClient(){
    	MongoClient mongoClient;
    	if(uri != null){
    		final MongoClientURI mongoClientURI = new MongoClientURI(uri);
    		mongoClient = new MongoClient(mongoClientURI);
    	}
    	else{
    		mongoClient = new MongoClient(host, port);
    	}
		return mongoClient;
    }
    
    private void init() {
        oplog = readCollection();
        query = createQuery();
    }

    /**
     * Loads the oplog collection.
     *
     * @return the oplog collection
     */
    private MongoCollection<Document> readCollection() {
		mongoClient = createMongoClient();
		
        log.trace("Starting database reader with configuration: ");
		log.trace("addresses: {}", StringUtils.join(mongoClient.getAllAddress(), ","));
        log.trace("db: {}", db);
        log.trace("start: {}", start);
        
        final MongoDatabase db = mongoClient.getDatabase("local");
        return db.getCollection("oplog.rs");
    }

    /**
     * Creates the query to execute on the collection.
     *
     * @return the query
     */
    private Bson createQuery() {
        // timestamps are used as offsets, saved as a concatenation of seconds and order
    	Integer timestamp = 0;
        Integer order = 0;
    	if(!start.equals("0")){    	
	    	final String[] splitted = start.split("_");
	    	timestamp = Integer.valueOf(splitted[0]);
	        order = Integer.valueOf(splitted[1]);
    	}
    	
        Bson query = Filters.and(
                Filters.exists("fromMigrate", false),
                Filters.gt("ts", new BSONTimestamp(timestamp, order)),
                Filters.or(
                        Filters.eq("op", "i"),
                        Filters.eq("op", "u"),
                        Filters.eq("op", "d")
                ),
                Filters.eq("ns", db)
        );

        return query;
    }
}