/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.cassandra.fqltool;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Predicate;
import java.util.stream.Collectors;

import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.MoreExecutors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.Timer;
import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.Statement;
import org.apache.cassandra.utils.FBUtilities;

public class QueryReplayer implements Closeable
{
    private static final Logger logger = LoggerFactory.getLogger(QueryReplayer.class);
    private static final int PRINT_RATE = 5000;
    private final ExecutorService es = Executors.newFixedThreadPool(1);
    private final Iterator<List<FQLQuery>> queryIterator;
    private final List<Predicate<FQLQuery>> filters;
    private final List<Session> sessions;
    private final ResultHandler resultHandler;
    private final MetricRegistry metrics = new MetricRegistry();
    private final SessionProvider sessionProvider;

    /**
     * @param queryIterator the queries to be replayed
     * @param targetHosts hosts to connect to, in the format "<user>:<password>@<host>:<port>" where only <host> is mandatory, port defaults to 9042
     * @param resultPaths where to write the results of the queries, for later comparisons, size should be the same as the number of iterators
     * @param filters query filters
     * @param queryFilePathString where to store the queries executed
     */
    public QueryReplayer(Iterator<List<FQLQuery>> queryIterator,
                         List<String> targetHosts,
                         List<File> resultPaths,
                         List<Predicate<FQLQuery>> filters,
                         String queryFilePathString)
    {
        this(queryIterator, targetHosts, resultPaths, filters, queryFilePathString, new DefaultSessionProvider(), null);
    }

    /**
     * Constructor public to allow external users to build their own session provider
     *
     * sessionProvider takes the hosts in targetHosts and creates one session per entry
     */
    public QueryReplayer(Iterator<List<FQLQuery>> queryIterator,
                         List<String> targetHosts,
                         List<File> resultPaths,
                         List<Predicate<FQLQuery>> filters,
                         String queryFilePathString,
                         SessionProvider sessionProvider,
                         MismatchListener mismatchListener)
    {
        this.sessionProvider = sessionProvider;
        this.queryIterator = queryIterator;
        this.filters = filters;
        sessions = targetHosts.stream().map(sessionProvider::connect).collect(Collectors.toList());
        File queryFilePath = queryFilePathString != null ? new File(queryFilePathString) : null;
        resultHandler = new ResultHandler(targetHosts, resultPaths, queryFilePath, mismatchListener);
    }

    public void replay()
    {
        while (queryIterator.hasNext())
        {
            List<FQLQuery> queries = queryIterator.next();
            for (FQLQuery query : queries)
            {
                if (filters.stream().anyMatch(f -> !f.test(query)))
                    continue;
                try (Timer.Context ctx = metrics.timer("queries").time())
                {
                    List<ListenableFuture<ResultHandler.ComparableResultSet>> results = new ArrayList<>(sessions.size());
                    Statement statement = query.toStatement();
                    for (Session session : sessions)
                    {
                        maybeSetKeyspace(session, query);
                        if (logger.isDebugEnabled())
                            logger.debug("Executing query: {}", query);
                        ListenableFuture<ResultSet> future = session.executeAsync(statement);
                        results.add(handleErrors(future));
                    }

                    ListenableFuture<List<ResultHandler.ComparableResultSet>> resultList = Futures.allAsList(results);

                    Futures.addCallback(resultList, new FutureCallback<List<ResultHandler.ComparableResultSet>>()
                    {
                        public void onSuccess(List<ResultHandler.ComparableResultSet> resultSets)
                        {
                            // note that the order of resultSets is signifcant here - resultSets.get(x) should
                            // be the result from a query against targetHosts.get(x)
                            resultHandler.handleResults(query, resultSets);
                        }

                        public void onFailure(Throwable throwable)
                        {
                            throw new AssertionError("Errors should be handled in FQLQuery.execute", throwable);
                        }
                    }, es);

                    FBUtilities.waitOnFuture(resultList);
                }
                catch (Throwable t)
                {
                    logger.error("QUERY %s got exception: %s", query, t.getMessage());
                }

                Timer timer = metrics.timer("queries");
                if (timer.getCount() % PRINT_RATE == 0)
                    logger.info(String.format("%d queries, rate = %.2f", timer.getCount(), timer.getOneMinuteRate()));
            }
        }
    }

    private void maybeSetKeyspace(Session session, FQLQuery query)
    {
        try
        {
            if (query.keyspace() != null && !query.keyspace().equals(session.getLoggedKeyspace()))
            {
                if (logger.isDebugEnabled())
                    logger.debug("Switching keyspace from {} to {}", session.getLoggedKeyspace(), query.keyspace());
                session.execute("USE " + query.keyspace());
            }
        }
        catch (Throwable t)
        {
            logger.error("USE {} failed: {}", query.keyspace(), t.getMessage());
        }
    }

    /**
     * Make sure we catch any query errors
     *
     * On error, this creates a failed ComparableResultSet with the exception set to be able to store
     * this fact in the result file and handle comparison of failed result sets.
     */
    private static ListenableFuture<ResultHandler.ComparableResultSet> handleErrors(ListenableFuture<ResultSet> result)
    {
        ListenableFuture<ResultHandler.ComparableResultSet> res = Futures.transform(result, DriverResultSet::new, MoreExecutors.directExecutor());
        return Futures.catching(res, Throwable.class, DriverResultSet::failed, MoreExecutors.directExecutor());
    }

    public void close() throws IOException
    {
        es.shutdown();
        sessionProvider.close();
        resultHandler.close();
    }

    static class ParsedTargetHost
    {
        final int port;
        final String user;
        final String password;
        final String host;

        ParsedTargetHost(String host, int port, String user, String password)
        {
            this.host = host;
            this.port = port;
            this.user = user;
            this.password = password;
        }

        static ParsedTargetHost fromString(String s)
        {
            String [] userInfoHostPort = s.split("@");

            String hostPort = null;
            String user = null;
            String password = null;
            if (userInfoHostPort.length == 2)
            {
                String [] userPassword = userInfoHostPort[0].split(":");
                if (userPassword.length != 2)
                    throw new RuntimeException("Username provided but no password");
                hostPort = userInfoHostPort[1];
                user = userPassword[0];
                password = userPassword[1];
            }
            else if (userInfoHostPort.length == 1)
                hostPort = userInfoHostPort[0];
            else
                throw new RuntimeException("Malformed target host: "+s);

            String[] splitHostPort = hostPort.split(":");
            int port = 9042;
            if (splitHostPort.length == 2)
                port = Integer.parseInt(splitHostPort[1]);

            return new ParsedTargetHost(splitHostPort[0], port, user, password);
        }
    }

    public static interface SessionProvider extends Closeable
    {
        Session connect(String connectionString);
        void close();
    }

    private static final class DefaultSessionProvider implements SessionProvider
    {
        private final static Map<String, Session> sessionCache = new HashMap<>();

        public synchronized Session connect(String connectionString)
        {
            if (sessionCache.containsKey(connectionString))
                return sessionCache.get(connectionString);
            Cluster.Builder builder = Cluster.builder();
            ParsedTargetHost pth = ParsedTargetHost.fromString(connectionString);
            builder.addContactPoint(pth.host);
            builder.withPort(pth.port);
            if (pth.user != null)
                builder.withCredentials(pth.user, pth.password);
            Cluster c = builder.build();
            sessionCache.put(connectionString, c.connect());
            return sessionCache.get(connectionString);
        }

        public void close()
        {
            sessionCache.entrySet().removeIf(entry -> {
                try (Session s = entry.getValue())
                {
                    s.getCluster().close();
                    return true;
                }
                catch (Throwable t)
                {
                    logger.error("Could not close connection", t);
                    return false;
                }
            });
        }
    }
}