# Copyright 2017 Google LLC
#
# 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

"""The YCSB client in Python.

Usage:

  # Set up instance and load data into database.

  # Set up environment variables. You should use your own credentials and gcloud
  # project.
  $ export GOOGLE_APPLICATION_CREDENTIALS=/path/to/credentials.json
  $ export GCLOUD_PROJECT=gcloud-project-name

  # Run the benchmark.
  $ python spanner/benchmark/ycsb.py run cloud_spanner -P pkb/workloada \
    -p table=usertable -p cloudspanner.instance=ycsb-542756a4 \
    -p recordcount=5000 -p operationcount=100 -p cloudspanner.database=ycsb \
    -p num_worker=1

  # To make a package so it can work with PerfKitBenchmarker.
  $ cd spanner; tar -cvzf ycsb-python.0.0.5.tar.gz benchmark/*

"""

from google.cloud import spanner

import argparse
import numpy
import random
import string
import threading
import timeit


OPERATIONS = ['readproportion', 'updateproportion', 'scanproportion',
              'insertproportion']
NUM_FIELD = 10


def parse_options():
    """Parses options."""
    parser = argparse.ArgumentParser()
    parser.add_argument('command', help='The YCSB command.')
    parser.add_argument('benchmark', help='The YCSB benchmark.')
    parser.add_argument('-P', '--workload', action='store', dest='workload',
                        default='', help='The path to a YCSB workload file.')
    parser.add_argument('-p', '--parameter', action='append', dest='parameters',
                        default=[], help='The key=value pair of parameter.')
    parser.add_argument('-b', '--num_bucket', action='store', type=int,
                        dest='num_bucket', default=1000,
                        help='The number of buckets in output.')

    args = parser.parse_args()

    parameters = {}
    parameters['command'] = args.command
    parameters['num_bucket'] = args.num_bucket

    for parameter in args.parameters:
        parts = parameter.strip().split('=')
        parameters[parts[0]] = parts[1]

    with open(args.workload, 'r') as f:
        for line in f.readlines():
            parts = line.split('=')
            key = parts[0].strip()
            if key in OPERATIONS:
                parameters[key] = parts[1].strip()

    return parameters


def open_database(parameters):
    """Opens a database specified by the parameters from parse_options()."""
    spanner_client = spanner.Client()
    instance_id = parameters['cloudspanner.instance']
    instance = spanner_client.instance(instance_id)
    database_id = parameters['cloudspanner.database']
    pool = spanner.BurstyPool(int(parameters['num_worker']))
    database = instance.database(database_id, pool=pool)

    return database


def load_keys(database, parameters):
    """Loads keys from database."""
    keys = []
    with database.snapshot() as snapshot:
        results = snapshot.execute_sql(
            'SELECT u.id FROM %s u' % parameters['table'])

    for row in results:
        keys.append(row[0])

    return keys


def read(database, table, key):
    """Does a single read operation."""
    with database.snapshot() as snapshot:
        result = snapshot.execute_sql('SELECT u.* FROM %s u WHERE u.id="%s"' %
                                      (table, key))
        for row in result:
            key = row[0]
        for i in range(NUM_FIELD):
            field = row[i + 1]


def update(database, table, key):
    """Does a single update operation."""
    field = random.randrange(10)
    value = ''.join(random.choice(string.printable) for i in range(100))
    with database.batch() as batch:
        batch.update(table=table, columns=('id', 'field%d' % field),
                     values=[(key, value)])


def do_operation(database, keys, table, operation, latencies_ms):
    """Does a single operation and records latency."""
    key = random.choice(keys)