#!/usr/bin/env python
#
# ryw_benchmark.py
#
# This source file is part of the FoundationDB open source project
#
# Copyright 2013-2018 Apple Inc. and the FoundationDB project authors
# 
# 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.
#


import argparse
import os
import sys
import time
import traceback

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from python_tests import PythonTest

import fdb
fdb.api_version(400)

class RYWBenchmark(PythonTest):
    tests = {
        'get_single' : "RYW: get single cached value throughput",
        'get_many_sequential' : "RYW: get sequential cached values throughput",
        'get_range_basic' : "RYW: get range cached values throughput",
        'single_clear_get_range' : "RYW: get range cached values with clears throughput",
        'clear_range_get_range' : "RYW: get range cached values with clear ranges throughput",
        'interleaved_sets_gets' : "RYW: interleaved sets and gets on a single key throughput",
    }

    def __init__(self):
        super(RYWBenchmark, self).__init__()

    def run_test(self):
        try:
            db = fdb.open(None, 'DB')
        except KeyboardInterrupt:
            raise
        except:
            self.result.add_error(self.get_error('fdb.open failed'))
            return

        try:
            self.test_performance(db)
        except KeyboardInterrupt:
            raise
        except:
            self.result.add_error(self.get_error('Failed to complete all tests'))

    #Adds the stack trace to an error message
    def get_error(self, message):
        errorMessage = message + "\n" + traceback.format_exc()
        print(errorMessage)
        return errorMessage

    def test_performance(self, db):
        tr = db.create_transaction()
        self.insert_data(tr)

        if not self.args.tests_to_run:
            self.args.tests_to_run = RYWBenchmark.tests.keys()
        else:
            for t in self.args.tests_to_run:
                if not t in RYWBenchmark.tests:
                    raise Exception("Unknown RYW benchmark test '%s'" % t)

        num_runs = 25

        for test in self.args.tests_to_run:
            time.sleep(5)
            print('Running test %s' % test)
            results = []

            fxn_name = 'run_%s' % test
            assert hasattr(self, fxn_name), 'Test function %s is not implemented' % fxn_name

            for x in range(0, num_runs):
                try:
                    results.append(getattr(self, fxn_name)(tr))
                except KeyboardInterrupt:
                    raise
                except:
                    self.result.add_error(self.get_error('Performance test failed: ' + RYWBenchmark.tests[test]))
                    break

            if len(results) == num_runs:
                median = sorted(results)[num_runs / 2]
                self.result.add_kpi(RYWBenchmark.tests[test], int(median), 'keys/s')

    def insert_data(self, tr):
        for i in range(0, 10000):
            tr[str(i)] = 'foo'

    def run_get_single(self, tr, count=10000):
      start = time.time()
      for i in range(count):
        tr.get("5001").wait()
      return count / (time.time() - start)

    def run_get_many_sequential(self, tr, count=10000):
      start = time.time()
      for j in range(count):
        tr.get(str(j)).wait()
      return count / (time.time() - start)

    def run_get_range_basic(self, tr, count=1000):
      start = time.time()
      for i in range(count):
        list(tr.get_range(str(0), str(10000)))
      return 10000 * count / (time.time() - start)

    def run_single_clear_get_range(self, tr, count=1000):
      for i in range(0, 10000, 2):
        tr.clear(str(i))
      start = time.time()
      for i in range(0, count):
        list(tr.get_range(str(0), str(10000)))
      kpi = 10000 * count / (time.time() - start)
      self.insert_data(tr)
      return kpi

    def run_clear_range_get_range(self, tr, count=1000):
      for i in range(0, 10000, 4):
        tr.clear_range(str(i), str(i+1))
      start = time.time()
      for i in range(0, count):
        list(tr.get_range(str(0), str(10000)))
      kpi = 10000 * count / (time.time() - start)
      self.insert_data(tr)
      return kpi

    def run_interleaved_sets_gets(self, tr, count=10000):
      start = time.time()
      tr['foo'] = str(1)
      for i in range(count):
        old = int(tr.get('foo').wait())
        tr.set('foo', str(old + 1))
      return count / (time.time() - start)

if __name__ == '__main__':
    print("Running RYW Benchmark test on Python version %d.%d.%d%s%d" % (sys.version_info[0], sys.version_info[1], sys.version_info[2], sys.version_info[3][0], sys.version_info[4]))

    parser = argparse.ArgumentParser()

    tests = sorted(RYWBenchmark.tests.keys())
    assert len(tests) > 0, 'RYW benchmark test has no test_functions'
    test_string = ', '.join(tests[:-1])
    if len(tests) > 1:
        test_string += ', and '

    test_string += tests[-1]

    parser.add_argument('--tests-to-run', nargs='*', help='Names of tests to run. Can be any of %s. By default, all tests are run.' % test_string)
    RYWBenchmark().run(parser=parser)