#!/usr/bin/env python
#
# Copyright 2013, 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 logging
import time
import unittest

from vtproto import topodata_pb2

from vtdb import keyrange
from vtdb import keyrange_constants
from vtdb import vtgate_client

import environment
import tablet
import utils

# source keyspace, with 4 tables
source_master = tablet.Tablet()
source_replica = tablet.Tablet()
source_rdonly1 = tablet.Tablet()
source_rdonly2 = tablet.Tablet()

# destination keyspace, with just two tables
destination_master = tablet.Tablet()
destination_replica = tablet.Tablet()
destination_rdonly1 = tablet.Tablet()
destination_rdonly2 = tablet.Tablet()


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

    setup_procs = [
        source_master.init_mysql(),
        source_replica.init_mysql(),
        source_rdonly1.init_mysql(),
        source_rdonly2.init_mysql(),
        destination_master.init_mysql(),
        destination_replica.init_mysql(),
        destination_rdonly1.init_mysql(),
        destination_rdonly2.init_mysql(),
        ]
    utils.Vtctld().start()
    utils.wait_procs(setup_procs)
  except:
    tearDownModule()
    raise


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

  if utils.vtgate:
    utils.vtgate.kill()
  teardown_procs = [
      source_master.teardown_mysql(),
      source_replica.teardown_mysql(),
      source_rdonly1.teardown_mysql(),
      source_rdonly2.teardown_mysql(),
      destination_master.teardown_mysql(),
      destination_replica.teardown_mysql(),
      destination_rdonly1.teardown_mysql(),
      destination_rdonly2.teardown_mysql(),
  ]
  utils.wait_procs(teardown_procs, raise_on_error=False)

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

  source_master.remove_tree()
  source_replica.remove_tree()
  source_rdonly1.remove_tree()
  source_rdonly2.remove_tree()
  destination_master.remove_tree()
  destination_replica.remove_tree()
  destination_rdonly1.remove_tree()
  destination_rdonly2.remove_tree()


class TestVerticalSplit(unittest.TestCase):

  def setUp(self):
    self.insert_index = 0

    self._init_keyspaces_and_tablets()
    utils.VtGate().start(cache_ttl='0s')

    # create the schema on the source keyspace, add some values
    self._create_source_schema()
    self._insert_initial_values()

  def tearDown(self):
    # kill everything
    tablet.kill_tablets([source_master, source_replica, source_rdonly1,
                         source_rdonly2, destination_master,
                         destination_replica, destination_rdonly1,
                         destination_rdonly2])
    utils.vtgate.kill()

  def _init_keyspaces_and_tablets(self):
    utils.run_vtctl(['CreateKeyspace', 'source_keyspace'])
    utils.run_vtctl(
        ['CreateKeyspace', '--served_from',
         'master:source_keyspace,replica:source_keyspace,rdonly:'
         'source_keyspace',
         'destination_keyspace'])
    source_master.start_vttablet(
        wait_for_state=None, target_tablet_type='replica',
        init_keyspace='source_keyspace', init_shard='0')
    source_replica.start_vttablet(
        wait_for_state=None, target_tablet_type='replica',
        init_keyspace='source_keyspace', init_shard='0')
    source_rdonly1.start_vttablet(
        wait_for_state=None, target_tablet_type='rdonly',
        init_keyspace='source_keyspace', init_shard='0')
    source_rdonly2.start_vttablet(
        wait_for_state=None, target_tablet_type='rdonly',
        init_keyspace='source_keyspace', init_shard='0')

    destination_master.start_vttablet(
        wait_for_state=None, target_tablet_type='replica',
        init_keyspace='destination_keyspace', init_shard='0')
    destination_replica.start_vttablet(
        wait_for_state=None, target_tablet_type='replica',
        init_keyspace='destination_keyspace', init_shard='0')
    destination_rdonly1.start_vttablet(
        wait_for_state=None, target_tablet_type='rdonly',
        init_keyspace='destination_keyspace', init_shard='0')
    destination_rdonly2.start_vttablet(
        wait_for_state=None, target_tablet_type='rdonly',
        init_keyspace='destination_keyspace', init_shard='0')

    # wait for the tablets
    all_setup_tablets = [
        source_master, source_replica, source_rdonly1, source_rdonly2,
        destination_master, destination_replica, destination_rdonly1,
        destination_rdonly2]
    for t in all_setup_tablets:
      t.wait_for_vttablet_state('NOT_SERVING')

    # check SrvKeyspace
    self._check_srv_keyspace('ServedFrom(master): source_keyspace\n'
                             'ServedFrom(rdonly): source_keyspace\n'
                             'ServedFrom(replica): source_keyspace\n')

    # reparent to make the tablets work
    utils.run_vtctl(['InitShardMaster', '-force', 'source_keyspace/0',
                     source_master.tablet_alias], auto_log=True)
    utils.run_vtctl(['InitShardMaster', '-force', 'destination_keyspace/0',
                     destination_master.tablet_alias], auto_log=True)

    for t in all_setup_tablets:
      t.wait_for_vttablet_state('SERVING')

  def _create_source_schema(self):
    create_table_template = '''create table %s(
id bigint not null,
msg varchar(64),
primary key (id),
index by_msg (msg)
) Engine=InnoDB'''
    create_view_template = 'create view %s(id, msg) as select id, msg from %s'

    for t in ['moving1', 'moving2', 'staying1', 'staying2']:
      utils.run_vtctl(['ApplySchema',
                       '-sql=' + create_table_template % (t),
                       'source_keyspace'],
                      auto_log=True)
    utils.run_vtctl(['ApplySchema',
                     '-sql=' + create_view_template % ('view1', 'moving1'),
                     'source_keyspace'],
                    auto_log=True)
    for t in [source_master, source_replica, source_rdonly1, source_rdonly2]:
      utils.run_vtctl(['ReloadSchema', t.tablet_alias])

    # Add a table to the destination keyspace which should be ignored.
    utils.run_vtctl(['ApplySchema',
                     '-sql=' + create_table_template % 'extra1',
                     'destination_keyspace'],
                    auto_log=True)
    for t in [destination_master, destination_replica,
              destination_rdonly1, destination_rdonly2]:
      utils.run_vtctl(['ReloadSchema', t.tablet_alias])

  def _insert_initial_values(self):
    self.moving1_first = self._insert_values('moving1', 100)
    self.moving2_first = self._insert_values('moving2', 100)
    staying1_first = self._insert_values('staying1', 100)
    staying2_first = self._insert_values('staying2', 100)
    self._check_values(source_master, 'vt_source_keyspace', 'moving1',
                       self.moving1_first, 100)
    self._check_values(source_master, 'vt_source_keyspace', 'moving2',
                       self.moving2_first, 100)
    self._check_values(source_master, 'vt_source_keyspace', 'staying1',
                       staying1_first, 100)
    self._check_values(source_master, 'vt_source_keyspace', 'staying2',
                       staying2_first, 100)
    self._check_values(source_master, 'vt_source_keyspace', 'view1',
                       self.moving1_first, 100)

    # Insert data directly because vtgate would redirect us.
    destination_master.mquery(
        'vt_destination_keyspace',
        "insert into %s (id, msg) values(%d, 'value %d')" % ('extra1', 1, 1),
        write=True)
    self._check_values(destination_master, 'vt_destination_keyspace', 'extra1',
                       1, 1)

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

  # insert some values in the source master db, return the first id used
  def _insert_values(self, table, count):
    result = self.insert_index
    conn = self._vtdb_conn()
    cursor = conn.cursor(
        tablet_type='master', keyspace='source_keyspace',
        keyranges=[keyrange.KeyRange(keyrange_constants.NON_PARTIAL_KEYRANGE)],
        writable=True)
    for _ in xrange(count):
      conn.begin()
      cursor.execute("insert into %s (id, msg) values(%d, 'value %d')" % (
          table, self.insert_index, self.insert_index), {})
      conn.commit()
      self.insert_index += 1
    conn.close()
    return result

  def _check_values(self, t, dbname, table, first, count):
    logging.debug(
        'Checking %d values from %s/%s starting at %d', count, dbname,
        table, first)
    rows = t.mquery(
        dbname, 'select id, msg from %s where id>=%d order by id limit %d' %
        (table, first, count))
    self.assertEqual(count, len(rows), 'got wrong number of rows: %d != %d' %
                     (len(rows), count))
    for i in xrange(count):