/*
 * 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.omid.benchmarks.tso;

import org.apache.phoenix.thirdparty.com.google.common.util.concurrent.RateLimiter;
import org.apache.phoenix.thirdparty.com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.inject.Binder;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.Module;
import org.apache.omid.benchmarks.utils.IntegerGenerator;
import org.apache.omid.committable.CommitTable;
import org.apache.omid.metrics.Counter;
import org.apache.omid.metrics.MetricsRegistry;
import org.apache.omid.metrics.Timer;
import org.apache.omid.tso.util.DummyCellIdImpl;
import org.apache.omid.tso.client.AbortException;
import org.apache.omid.tso.client.CellId;
import org.apache.omid.tso.client.OmidClientConfiguration;
import org.apache.omid.tso.client.TSOClient;
import org.apache.omid.tso.client.TSOFuture;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetAddress;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import static com.codahale.metrics.MetricRegistry.name;

class RawTxRunner implements Runnable {

    private static final Logger LOG = LoggerFactory.getLogger(RawTxRunner.class);

    private static volatile int txRunnerCounter = 0;
    private int txRunnerId = txRunnerCounter++;

    // Config params
    private final int writesetSize;
    private final boolean fixedWriteSetSize;
    private final long commitDelayInMs;
    private final int percentageOfReadOnlyTxs;
    private final IntegerGenerator cellIdGenerator;
    private final Random randomGen;

    // Main elements
    private final TSOClient tsoClient;
    private final CommitTable.Client commitTableClient;

    // Asynchronous executor for tx post begin sequence: TimestampListener -> Committer -> CommitListener
    private final ScheduledExecutorService callbackExec =
            Executors.newScheduledThreadPool(1, new ThreadFactoryBuilder()
                    .setNameFormat("tx-runner-" + txRunnerId + "-callback")
                    .setUncaughtExceptionHandler(new Thread.UncaughtExceptionHandler() {
                        @Override
                        public void uncaughtException(Thread t, Throwable e) {
                            LOG.error("Thread {} threw exception", t, e);
                        }
                    }).build());

    // Statistics to save
    private final Timer timestampTimer;
    private final Timer commitTimer;
    private final Timer abortTimer;
    private final Counter errorCounter;

    // Allows to setup a maximum rate for the client in req/sec
    private final RateLimiter rateLimiter;

    // Is this TxRunner still running?
    private volatile boolean isRunning = false;

    RawTxRunner(final TSOServerBenchmarkConfig expConfig) throws IOException, InterruptedException {

        // Injector configuration
        List<Module> guiceModules = new ArrayList<>();
        guiceModules.add(new Module() {
            @Override
            public void configure(Binder binder) {
                binder.bind(MetricsRegistry.class).toInstance(expConfig.getMetrics());
            }
        });
        guiceModules.add(expConfig.getCommitTableStoreModule());
        Injector injector = Guice.createInjector(guiceModules);

        // Tx Runner config
        this.writesetSize = expConfig.getWritesetSize();
        this.fixedWriteSetSize = expConfig.isFixedWritesetSize();
        this.commitDelayInMs = expConfig.getCommitDelayInMs();
        this.percentageOfReadOnlyTxs = expConfig.getPercentageOfReadOnlyTxs();
        this.cellIdGenerator = expConfig.getCellIdGenerator();
        this.randomGen = new Random(System.currentTimeMillis() * txRunnerId); // to make it channel dependent

        int txRateInReqPerSec = expConfig.getTxRateInRequestPerSecond();
        long warmUpPeriodInSecs = expConfig.getWarmUpPeriodInSecs();

        LOG.info("TxRunner-{} [ Tx Rate (Req per Sec) -> {} ]", txRunnerId, txRateInReqPerSec);
        LOG.info("TxRunner-{} [ Warm Up Period -> {} Secs ]", txRunnerId, warmUpPeriodInSecs);
        LOG.info("TxRunner-{} [ Cell Id Distribution Generator -> {} ]", txRunnerId, expConfig.getCellIdGenerator().getClass());
        LOG.info("TxRunner-{} [ Max Tx Size -> {} Fixed: {} ]", txRunnerId, writesetSize, fixedWriteSetSize);
        LOG.info("TxRunner-{} [ Commit delay -> {} Ms ]", txRunnerId, commitDelayInMs);
        LOG.info("TxRunner-{} [ % of Read-Only Tx -> {} % ]", txRunnerId, percentageOfReadOnlyTxs);

        // Commit table client initialization
        CommitTable commitTable = injector.getInstance(CommitTable.class);
        this.commitTableClient = commitTable.getClient();

        // Stat initialization
        MetricsRegistry metrics = injector.getInstance(MetricsRegistry.class);
        String hostName = InetAddress.getLocalHost().getHostName();
        this.timestampTimer = metrics.timer(name("tx_runner", Integer.toString(txRunnerId), hostName, "timestamp"));
        this.commitTimer = metrics.timer(name("tx_runner", Integer.toString(txRunnerId), hostName, "commit"));
        this.abortTimer = metrics.timer(name("tx_runner", Integer.toString(txRunnerId), hostName, "abort"));
        this.errorCounter = metrics.counter(name("tx_runner", Integer.toString(txRunnerId), hostName, "errors"));
        LOG.info("TxRunner-{} [ Metrics provider module -> {} ]", txRunnerId, expConfig.getMetrics().getClass());

        // TSO Client initialization
        OmidClientConfiguration tsoClientConf = expConfig.getOmidClientConfiguration();
        this.tsoClient = TSOClient.newInstance(tsoClientConf);
        LOG.info("TxRunner-{} [ Connection Type {}/Connection String {} ]", txRunnerId,
                 tsoClientConf.getConnectionType(), tsoClientConf.getConnectionString());

        // Limiter for configured request per second
        this.rateLimiter = RateLimiter.create((double) txRateInReqPerSec, warmUpPeriodInSecs, TimeUnit.SECONDS);
    }

    @Override
    public void run() {

        isRunning = true;

        while (isRunning) {
            rateLimiter.acquire();
            long tsRequestTime = System.nanoTime();
            final TSOFuture<Long> tsFuture = tsoClient.getNewStartTimestamp();
            tsFuture.addListener(new TimestampListener(tsFuture, tsRequestTime), callbackExec);
        }

        shutdown();

    }

    void stop() {
        isRunning = false;
    }

    private void shutdown() {

        try {
            LOG.info("Finishing TxRunner in 3 secs", txRunnerId);
            boolean wasSuccess = callbackExec.awaitTermination(3, TimeUnit.SECONDS);
            if (!wasSuccess) {
                callbackExec.shutdownNow();
            }
            tsoClient.close().get();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            // ignore
        } catch (ExecutionException e) {
            // ignore
        } finally {
            LOG.info("TxRunner {} finished", txRunnerId);
        }

    }

    private class TimestampListener implements Runnable {

        final TSOFuture<Long> tsFuture;
        final long tsRequestTime;

        TimestampListener(TSOFuture<Long> tsFuture, long tsRequestTime) {
            this.tsFuture = tsFuture;
            this.tsRequestTime = tsRequestTime;
        }

        @Override
        public void run() {

            try {
                long txId = tsFuture.get();
                timestampTimer.update(System.nanoTime() - tsRequestTime);
                if (commitDelayInMs <= 0) {
                    callbackExec.execute(new Committer(txId));
                } else {
                    callbackExec.schedule(new Committer(txId), commitDelayInMs, TimeUnit.MILLISECONDS);
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                errorCounter.inc();
            } catch (ExecutionException e) {
                errorCounter.inc();
            }

        }

    }

    private class Committer implements Runnable {

        final long txId;

        Committer(long txId) {
            this.txId = txId;
        }

        @Override
        public void run() {

            int txWritesetSize = calculateTxWritesetSize();

            if (txWritesetSize == 0) {
                return; // Read only tx, no need to commit
            }
            // Otherwise, we create the writeset...
            final Set<CellId> cells = new HashSet<>();
            for (byte i = 0; i < txWritesetSize; i++) {
                long cellId = cellIdGenerator.nextInt();
                cells.add(new DummyCellIdImpl(cellId));
            }
            // ... and we commit the transaction
            long startCommitTimeInNs = System.nanoTime();
            final TSOFuture<Long> commitFuture = tsoClient.commit(txId, cells);
            commitFuture.addListener(new CommitListener(txId, commitFuture, startCommitTimeInNs), callbackExec);

        }

        private int calculateTxWritesetSize() {
            int txSize = 0;
            boolean readOnly = (randomGen.nextFloat() * 100) < percentageOfReadOnlyTxs;
            if (!readOnly) {
                if (fixedWriteSetSize) {
                    txSize = writesetSize;
                } else {
                    txSize = randomGen.nextInt(writesetSize) + 1;
                }
            }
            return txSize;
        }

    }

    private class CommitListener implements Runnable {

        final long txId;
        final long commitRequestTime;
        final TSOFuture<Long> commitFuture;

        CommitListener(long txId, TSOFuture<Long> commitFuture, long commitRequestTime) {
            this.txId = txId;
            this.commitFuture = commitFuture;
            this.commitRequestTime = commitRequestTime;
        }

        @Override
        public void run() {

            try {
                commitFuture.get();
                commitTableClient.deleteCommitEntry(txId).get();
                commitTimer.update(System.nanoTime() - commitRequestTime);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                errorCounter.inc();
            } catch (ExecutionException e) {
                if (e.getCause() instanceof AbortException) {
                    abortTimer.update(System.nanoTime() - commitRequestTime);
                } else {
                    errorCounter.inc();
                }
            }

        }

    }

}