#!/usr/bin/env python
#
# Copyright 2015, Google Inc. All rights reserved.
# Use of this source code is governed by a BSD-style license that can
# be found in the LICENSE file.

import base64
import unittest

from vtproto import topodata_pb2

from vtdb import vtgate_client

import environment
import tablet
import utils

# shards
shard_0_master = tablet.Tablet()
shard_0_rdonly = tablet.Tablet()

shard_1_master = tablet.Tablet()
shard_1_rdonly = tablet.Tablet()


def setUpModule():
  try:
    environment.topo_server().setup()

    setup_procs = [
        shard_0_master.init_mysql(),
        shard_0_rdonly.init_mysql(),
        shard_1_master.init_mysql(),
        shard_1_rdonly.init_mysql(),
        ]
    utils.Vtctld().start()
    utils.VtGate().start()
    utils.wait_procs(setup_procs)
  except:
    tearDownModule()
    raise


def tearDownModule():
  if utils.options.skip_teardown:
    return

  teardown_procs = [
      shard_0_master.teardown_mysql(),
      shard_0_rdonly.teardown_mysql(),
      shard_1_master.teardown_mysql(),
      shard_1_rdonly.teardown_mysql(),
      ]
  utils.wait_procs(teardown_procs, raise_on_error=False)

  environment.topo_server().teardown()
  utils.kill_sub_processes()
  utils.remove_tmp_files()

  shard_0_master.remove_tree()
  shard_0_rdonly.remove_tree()
  shard_1_master.remove_tree()
  shard_1_rdonly.remove_tree()


class TestCustomSharding(unittest.TestCase):
  """Test a custom-shared keyspace."""

  def _vtdb_conn(self):
    protocol, addr = utils.vtgate.rpc_endpoint(python=True)
    return vtgate_client.connect(protocol, addr, 30.0)

  def _insert_data(self, shard, start, count, table='data'):
    sql = 'insert into ' + table + '(id, name) values (%(id)s, %(name)s)'
    conn = self._vtdb_conn()
    cursor = conn.cursor(
        tablet_type='master', keyspace='test_keyspace',
        shards=[shard],
        writable=True)
    for x in xrange(count):
      bindvars = {
          'id': start+x,
          'name': 'row %d' % (start+x),
      }
      conn.begin()
      cursor.execute(sql, bindvars)
      conn.commit()
    conn.close()

  def _check_data(self, shard, start, count, table='data'):
    sql = 'select name from ' + table + ' where id=%(id)s'
    conn = self._vtdb_conn()
    cursor = conn.cursor(
        tablet_type='master', keyspace='test_keyspace',
        shards=[shard])
    for x in xrange(count):
      bindvars = {
          'id': start+x,
      }
      cursor.execute(sql, bindvars)
      qr = cursor.fetchall()
      self.assertEqual(len(qr), 1)
      v = qr[0][0]
      self.assertEqual(v, 'row %d' % (start+x))
    conn.close()

  def test_custom_end_to_end(self):
    """Runs through the common operations of a custom sharded keyspace.

    Tests creation with one shard, schema change, reading / writing
    data, adding one more shard, reading / writing data from both
    shards, applying schema changes again, and reading / writing data
    from both shards again.
    """

    utils.run_vtctl(['CreateKeyspace', 'test_keyspace'])

    # start the first shard only for now
    shard_0_master.init_tablet('master', 'test_keyspace', '0')
    shard_0_rdonly.init_tablet('rdonly', 'test_keyspace', '0')
    for t in [shard_0_master, shard_0_rdonly]:
      t.create_db('vt_test_keyspace')
      t.start_vttablet(wait_for_state=None)
    for t in [shard_0_master, shard_0_rdonly]:
      t.wait_for_vttablet_state('SERVING')

    utils.run_vtctl(['InitShardMaster', 'test_keyspace/0',
                     shard_0_master.tablet_alias], auto_log=True)
    utils.run_vtctl(['RebuildKeyspaceGraph', 'test_keyspace'], auto_log=True)

    self._check_shards_count_in_srv_keyspace(1)
    s = utils.run_vtctl_json(['GetShard', 'test_keyspace/0'])
    self.assertEqual(len(s['served_types']), 3)

    # create a table on shard 0
    sql = '''create table data(
id bigint auto_increment,
name varchar(64),
primary key (id)
) Engine=InnoDB'''
    utils.run_vtctl(['ApplySchema', '-sql=' + sql, 'test_keyspace'],
                    auto_log=True)

    # reload schema everywhere so the QueryService knows about the tables
    for t in [shard_0_master, shard_0_rdonly]:
      utils.run_vtctl(['ReloadSchema', t.tablet_alias], auto_log=True)

    # insert data on shard 0
    self._insert_data('0', 100, 10)

    # re-read shard 0 data
    self._check_data('0', 100, 10)

    # create shard 1
    shard_1_master.init_tablet('master', 'test_keyspace', '1')
    shard_1_rdonly.init_tablet('rdonly', 'test_keyspace', '1')
    for t in [shard_1_master, shard_1_rdonly]:
      t.start_vttablet(wait_for_state=None)
    for t in [shard_1_master, shard_1_rdonly]:
      t.wait_for_vttablet_state('NOT_SERVING')
    s = utils.run_vtctl_json(['GetShard', 'test_keyspace/1'])
    self.assertEqual(len(s['served_types']), 3)

    utils.run_vtctl(['InitShardMaster', 'test_keyspace/1',
                     shard_1_master.tablet_alias], auto_log=True)
    utils.run_vtctl(['CopySchemaShard', shard_0_rdonly.tablet_alias,
                     'test_keyspace/1'], auto_log=True)
    for t in [shard_1_master, shard_1_rdonly]:
      utils.run_vtctl(['RefreshState', t.tablet_alias], auto_log=True)
      t.wait_for_vttablet_state('SERVING')

    # rebuild the keyspace serving graph now that the new shard was added
    utils.run_vtctl(['RebuildKeyspaceGraph', 'test_keyspace'], auto_log=True)

    # insert data on shard 1
    self._insert_data('1', 200, 10)

    # re-read shard 1 data
    self._check_data('1', 200, 10)

    # create a second table on all shards
    sql = '''create table data2(
id bigint auto_increment,
name varchar(64),
primary key (id)
) Engine=InnoDB'''
    utils.run_vtctl(['ApplySchema', '-sql=' + sql, 'test_keyspace'],