/**
 * 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.activemq.store.kahadb.disk.index;

import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

import junit.framework.TestCase;
import org.apache.activemq.store.kahadb.disk.page.PageFile;
import org.apache.activemq.store.kahadb.disk.page.Transaction;
import org.apache.activemq.util.IOHelper;

/**
 * @author chirino
 */
public abstract class IndexBenchmark extends TestCase {

    // Slower machines might need to make this bigger.
    private static final long SAMPLE_DURATION = Integer.parseInt(System.getProperty("SAMPLES_DURATION", "" + 1000 * 5));
    // How many times do we sample?
    private static final long SAMPLES = Integer.parseInt(System.getProperty("SAMPLES", "" + 60 * 1000 / SAMPLE_DURATION));
    // How many indexes will we be benchmarking concurrently?
    private static final int INDEX_COUNT = Integer.parseInt(System.getProperty("INDEX_COUNT", "" + 1));
    // Indexes tend to perform worse when they get big.. so how many items
    // should we put into the index before we start sampling.
    private static final int INDEX_PRE_LOAD_COUNT = Integer.parseInt(System.getProperty("INDEX_PRE_LOAD_COUNT", "" + 10000 / INDEX_COUNT));

    protected File ROOT_DIR;
    protected final HashMap<String, Index<String, Long>> indexes = new HashMap<String, Index<String, Long>>();
    protected PageFile pf;

    public void setUp() throws Exception {
        ROOT_DIR = new File(IOHelper.getDefaultDataDirectory());
        IOHelper.delete(ROOT_DIR);

        pf = new PageFile(ROOT_DIR, getClass().getName());
        pf.load();
    }

    protected void tearDown() throws Exception {
        Transaction tx = pf.tx();
        for (Index<?, ?> i : indexes.values()) {
            try {
                i.unload(tx);
            } catch (Throwable ignore) {
            }
        }
        tx.commit();
    }

    abstract protected Index<String, Long> createIndex() throws Exception;

    synchronized private Index<String, Long> openIndex(String name) throws Exception {
        Transaction tx = pf.tx();
        Index<String, Long> index = indexes.get(name);
        if (index == null) {
            index = createIndex();
            index.load(tx);
            indexes.put(name, index);
        }
        tx.commit();
        return index;
    }

    class Producer extends Thread {
        private final String name;
        AtomicBoolean shutdown = new AtomicBoolean();

        public Producer(String name) {
            super("Producer: " + name);
            this.name = name;
        }

        public void shutdown() {
            shutdown.set(true);
        }

        @Override
        public void run() {
            try {

                Transaction tx = pf.tx();

                Index<String,Long> index = openIndex(name);
                long counter = 0;
                while (!shutdown.get()) {
                    long c = counter;

                    String key = key(c);
                    index.put(tx, key, c);
                    tx.commit();
                    Thread.yield(); // This avoids consumer starvation..

                    onProduced(counter++);
                }

            } catch (Throwable e) {
                e.printStackTrace();
            }
        }

        public void onProduced(long counter) {
        }
    }

    protected String key(long c) {
        return "a-long-message-id-like-key-" + c;
    }


    class Consumer extends Thread {
        private final String name;
        AtomicBoolean shutdown = new AtomicBoolean();

        public Consumer(String name) {
            super("Consumer: " + name);
            this.name = name;
        }

        public void shutdown() {
            shutdown.set(true);
        }

        @Override
        public void run() {
            try {
                Transaction tx = pf.tx();

                Index<String,Long> index = openIndex(name);
                long counter = 0;
                while (!shutdown.get()) {
                    long c = counter;
                    String key = key(c);

                    Long record = index.get(tx, key);
                    if (record != null) {
                        if( index.remove(tx, key) == null ) {
                            System.out.print("Remove failed...");
                        }
                        tx.commit();
                        onConsumed(counter++);
                    }
                }
            } catch (Throwable e) {
                e.printStackTrace();
            }
        }

        public void onConsumed(long counter) {
        }
    }

    protected void dumpIndex(Index<String, Long> index) throws IOException {
    }

    public void testLoad() throws Exception {

        final Producer producers[] = new Producer[INDEX_COUNT];
        final Consumer consumers[] = new Consumer[INDEX_COUNT];
        final CountDownLatch preloadCountDown = new CountDownLatch(INDEX_COUNT);
        final AtomicLong producedRecords = new AtomicLong();
        final AtomicLong consumedRecords = new AtomicLong();

        System.out.println("Starting: " + INDEX_COUNT + " producers");
        for (int i = 0; i < INDEX_COUNT; i++) {
            producers[i] = new Producer("test-" + i) {
                private boolean prelaodDone;

                public void onProduced(long counter) {
                    if (!prelaodDone && counter >= INDEX_PRE_LOAD_COUNT) {
                        prelaodDone = true;
                        preloadCountDown.countDown();
                    }
                    producedRecords.incrementAndGet();
                }
            };
            producers[i].start();
        }

        long start = System.currentTimeMillis();
        System.out.println("Waiting for each producer create " + INDEX_PRE_LOAD_COUNT + " records before starting the consumers.");
        preloadCountDown.await();
        long end = System.currentTimeMillis();
        System.out.println("Preloaded " + INDEX_PRE_LOAD_COUNT * INDEX_COUNT + " records at " + (INDEX_PRE_LOAD_COUNT * INDEX_COUNT * 1000f / (end - start)) + " records/sec");

        System.out.println("Starting: " + INDEX_COUNT + " consumers");
        for (int i = 0; i < INDEX_COUNT; i++) {
            consumers[i] = new Consumer("test-" + i) {
                public void onConsumed(long counter) {
                    consumedRecords.incrementAndGet();
                }
            };
            consumers[i].start();
        }

        long sample_start = System.currentTimeMillis();
        System.out.println("Taking " + SAMPLES + " performance samples every " + SAMPLE_DURATION + " ms");
        System.out.println("time (s), produced, produce rate (r/s), consumed, consume rate (r/s), used memory (k)");
        producedRecords.set(0);
        consumedRecords.set(0);
        for (int i = 0; i < SAMPLES; i++) {
            start = System.currentTimeMillis();
            Thread.sleep(SAMPLE_DURATION);
            end = System.currentTimeMillis();
            long p = producedRecords.getAndSet(0);
            long c = consumedRecords.getAndSet(0);

            long usedMemory = Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory();

            System.out.println(((end-sample_start)/1000f)+", "+p+", "+(p * 1000f / (end - start)) + ", "+ c+", " + (c * 1000f / (end - start))+", "+(usedMemory/(1024)) );
        }
        System.out.println("Samples done... Shutting down the producers and consumers...");
        for (int i = 0; i < INDEX_COUNT; i++) {
            producers[i].shutdown();
            consumers[i].shutdown();
        }
        for (int i = 0; i < INDEX_COUNT; i++) {
            producers[i].join(1000 * 5);
            consumers[i].join(1000 * 5);
        }
        System.out.println("Shutdown.");
    }

}