from __future__ import absolute_import
import itertools
import logging
from operator import attrgetter
import os
import getpass
import random
import sys
import time
import uuid

import cql
import pycassa
from pycassa.pool import ConnectionPool
from pycassa.system_manager import SystemManager, SIMPLE_STRATEGY, \
    LEXICAL_UUID_TYPE, ASCII_TYPE, BYTES_TYPE
from pycassa.types import CompositeType, TimeUUIDType, LexicalUUIDType, UUIDType, UTF8Type
import pytest
import yaml

import kvlayer
import yakonfig

logger = logging.getLogger(__name__)
one_mb = ' ' * 2**20

@pytest.yield_fixture(scope='function')
def client(namespace_string, request):
    config_path = str(request.fspath.dirpath('config_cassandra.yaml'))
    with yakonfig.defaulted_config([kvlayer], filename=config_path,
                                   params={'namespace': namespace_string,
                                           'app_name': 'kvltest'}):
        client = kvlayer.client()
        yield client
        client.delete_namespace()

def grouper(iterable, n, fillvalue=None):
    "Collect data into fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx
    args = [iter(iterable)] * n
    return itertools.izip_longest(fillvalue=fillvalue, *args)

@pytest.mark.skipif('True')
def test_composite_column_names(client):
    '''
    examine the unique nature of cassandras "wide sorted rows";

    Basically, C* is good at storing large OrderedDict objects and
    less good at the simpler kind of key=value storage that we were
    doing earlier.
    '''
    config = client._config
    namespace = client._app_namespace
    chosen_server = client._chosen_server
    sm = SystemManager(chosen_server)
    sm.create_keyspace(namespace, SIMPLE_STRATEGY, {'replication_factor': '1'})

    family = 'test'
    sm.create_column_family(
        namespace, family, super=False,
        key_validation_class = ASCII_TYPE,
        default_validation_class = BYTES_TYPE,
        comparator_type=CompositeType(UUIDType(), UUIDType()),
        #column_name_class = LEXICAL_UUID_TYPE
        )

    logger.info( sm.describe_schema_versions() )

    pool = ConnectionPool(namespace, config['storage_addresses'],
                          max_retries=1000, pool_timeout=10, pool_size=2, timeout=120)

    cf = pycassa.ColumnFamily(pool, family)
    u1, u2, u3, u4 = uuid.uuid1(), uuid.uuid1(), uuid.uuid1(), uuid.uuid1()

    ## insert four
    cf.insert('inbound', {(u1, u2): b''})
    cf.insert('inbound', {(u1, u3): b''})
    cf.insert('inbound', {(u1, u4): b''})
    cf.insert('inbound', {(u1, u1): b'45'})

    ## insert three with duplicates
    cf.insert('inbound', {(u3, u3): b'42'})
    cf.insert('inbound', {(u3, u3): b'43'})
    cf.insert('inbound', {(u3, u2): b''})
    cf.insert('inbound', {(u3, u4): b''})

    ## insert two
    cf.insert('inbound', {(u2, u3): b''})
    cf.insert('inbound', {(u2, u4): b''})

    ## verify start/finish parameters work as expected
    assert 4 == len(cf.get('inbound', column_start=(u1,), column_finish=(u1,)))
    assert 2 == len(cf.get('inbound', column_start=(u2,), column_finish=(u2,)))
    assert 3 == len(cf.get('inbound', column_start=(u3,), column_finish=(u3,)))

    ## check delete of a single column
    cf.remove('inbound', columns=[(u2, u3)])
    assert 1 == len(cf.get('inbound', column_start=(u2,), column_finish=(u2,)))

    ## check that the last inserted value appears at the expected
    ## location in the OrderedDict
    assert '43' == cf.get('inbound', column_start=(u3,), column_finish=(u3,)).items()[1][1]

    rec1 = cf.get('inbound', column_start=(u3,u3), column_finish=(u3,u3)).items()
    assert len(rec1) == 1

    start  = uuid.UUID(int=0)
    finish = uuid.UUID(int=2**128-1)
    assert start < u1 < u2 < u3 < u4 < finish
    assert 8 == len(cf.get('inbound', column_start=(start,), column_finish=(finish,)).items())

    ## test range searching
    start  = uuid.UUID(int=u3.int - 1)
    finish = uuid.UUID(int=u3.int + 1)
    assert start.int < u3.int < finish.int
    rec2 = cf.get('inbound', column_start=(start,), column_finish=(finish,)).items()
    assert rec2[1][0] == rec1[0][0]
    assert rec2[1][1] == rec1[0][1] == '43'

    cf.insert('inbound', {(u3,u3): b''.join(map(lambda u: u.bytes, (u1,u2,u3,u4)))})
    data = cf.get('inbound', column_start=(u3,u3), column_finish=(u3,u3)).items()[0][1]

    assert [u1,u2,u3,u4] == map(lambda b: uuid.UUID(bytes=b), grouper(data, 16))

    sm.close()

@pytest.mark.skipif('True')
def test_composite_column_names_with_decomposited_keys(client):
    '''
    examine the unique nature of cassandras "wide sorted rows" using
    concatenated keys
    '''
    config = client._config
    namespace = client._app_namespace
    chosen_server = client._chosen_server
    sm = SystemManager(chosen_server)
    sm.create_keyspace(namespace, SIMPLE_STRATEGY, {'replication_factor': '1'})

    family = 'test'
    sm.create_column_family(
        namespace, family, super=False,
        key_validation_class = ASCII_TYPE,
        default_validation_class = BYTES_TYPE,
        comparator_type=UTF8Type()
        )

    logger.info( sm.describe_schema_versions() )

    pool = ConnectionPool(namespace, config['storage_addresses'],
                          max_retries=1000, pool_timeout=10, pool_size=2, timeout=120)

    cf = pycassa.ColumnFamily(pool, family)
    u1, u2, u3, u4 = uuid.uuid1(), uuid.uuid1(), uuid.uuid1(), uuid.uuid1()

    ## insert four
    cf.insert('inbound', {join_uuids(u1, u2): b''})
    cf.insert('inbound', {join_uuids(u1, u3): b''})
    cf.insert('inbound', {join_uuids(u1, u4): b''})
    cf.insert('inbound', {join_uuids(u1, u1): b'45'})

    ## insert three with duplicates
    cf.insert('inbound', {join_uuids(u3, u3): b'42'})
    cf.insert('inbound', {join_uuids(u3, u3): b'43'})
    cf.insert('inbound', {join_uuids(u3, u2): b''})
    cf.insert('inbound', {join_uuids(u3, u4): b''})

    ## insert two
    cf.insert('inbound', {join_uuids(u2, u3): b''})
    cf.insert('inbound', {join_uuids(u2, u4): b''})

    ## verify start/finish parameters work as expected
    assert 4 == len(cf.get('inbound', column_start=join_uuids(u1, num_uuids=2), column_finish=join_uuids(u1, num_uuids=2, padding='f')))
    assert 2 == len(cf.get('inbound', column_start=join_uuids(u2, num_uuids=2), column_finish=join_uuids(u2, num_uuids=2, padding='f')))
    assert 3 == len(cf.get('inbound', column_start=join_uuids(u3, num_uuids=2), column_finish=join_uuids(u3, num_uuids=2, padding='f')))

    ## check delete of a single column
    cf.remove('inbound', columns=[join_uuids(u2, u3)])
    assert 1 == len(cf.get('inbound', column_start=join_uuids(u2, num_uuids=2), column_finish=join_uuids(u2, num_uuids=2, padding='f')))

    ## check that the last inserted value appears at the expected
    ## location in the OrderedDict
    assert '43' == cf.get('inbound', column_start=join_uuids(u3, num_uuids=2), column_finish=join_uuids(u3, num_uuids=2, padding='f')).items()[1][1]

    rec1 = cf.get('inbound', column_start=join_uuids(u3,u3), column_finish=join_uuids(u3,u3)).items()
    assert len(rec1) == 1

    start  = uuid.UUID(int=0)
    finish = uuid.UUID(int=2**128-1)
    assert start < u1 < u2 < u3 < u4 < finish
    assert 8 == len(cf.get('inbound', column_start=join_uuids(start, num_uuids=2), column_finish=join_uuids(finish, num_uuids=2, padding='f')).items())

    ## test range searching
    start  = uuid.UUID(int=u3.int - 1)
    finish = uuid.UUID(int=u3.int + 1)
    assert start.int < u3.int < finish.int
    rec2 = cf.get('inbound', column_start=join_uuids(start, num_uuids=2), column_finish=join_uuids(finish, num_uuids=2, padding='f')).items()
    assert rec2[1][0] == rec1[0][0]
    assert rec2[1][1] == rec1[0][1] == '43'

    cf.insert('inbound', {join_uuids(u3,u3): b''.join(map(lambda u: u.bytes, (u1,u2,u3,u4)))})
    data = cf.get('inbound', column_start=join_uuids(u3,u3), column_finish=join_uuids(u3,u3)).items()[0][1]

    assert [u1,u2,u3,u4] == map(lambda b: uuid.UUID(bytes=b), grouper(data, 16))

    sm.close()

@pytest.mark.xfail
def test_composite_column_names_second_level_range_query(client):
    '''
    check that we can execute range queries on the second part of a
    CompositeType column name
    '''
    config = client._config
    namespace = client._app_namespace
    chosen_server = client._chosen_server
    sm = SystemManager(chosen_server)
    sm.create_keyspace(namespace, SIMPLE_STRATEGY, {'replication_factor': '1'})

    family = 'test'
    sm.create_column_family(
        namespace, family, super=False,
        key_validation_class = ASCII_TYPE,
        default_validation_class = BYTES_TYPE,
        comparator_type=CompositeType(UUIDType(), UUIDType()),
        )
    pool = ConnectionPool(namespace, config['storage_addresses'],
                          max_retries=1000, pool_timeout=10, pool_size=2, timeout=120)

    cf = pycassa.ColumnFamily(pool, family)
    u1, u2, u3, u4 = uuid.uuid1(), uuid.uuid1(), uuid.uuid1(), uuid.uuid1()

    cf.insert('inbound', {(u1, u2): b''})
    cf.insert('inbound', {(u1, u3): b''})
    cf.insert('inbound', {(u1, u4): b''})

    ## test range searching
    start  = uuid.UUID(int=u3.int - 1)
    finish = uuid.UUID(int=u3.int + 1)
    assert start.int < u3.int < finish.int
    rec3 = cf.get('inbound',
                  column_start =(u1, start),
                  column_finish=(u1, finish)).items()
    assert len(rec3) == 1
    assert rec3[0][0][1] == u3
    ####  This assert above passes!

    ####  This next part fails :-/
    ## now insert many rows -- enough that some should fall in each
    ## subrange below
    for i in xrange(1000):
        cf.insert('inbound', {(u1, uuid.uuid4()): b''})

    ## do four ranges, and expect more than zero in each
    step_size = 2**(128 - 2)
    for i in range(2**2, 0, -1):
        start =  uuid.UUID(int=(i-1) * step_size)
        finish = uuid.UUID(int=min(i * step_size, 2**128 - 1))
        recs = cf.get('inbound',
                      column_start =(u1, start),
                      column_finish=(u1, finish)).items()
        for key, val in recs:
            assert val == b''
            assert key[0] == u1
            assert key[1] < finish
            assert start < key[1]              ## this fails!!

        assert len(recs) > 0
        logger.info( '%r for %r %r' % (len(recs), start, finish))

    sm.close()

# TODO: delete this and use kvlayer._util.join_uuids(), but really drop it and use join_key_fragments
def join_uuids(*uuids, **kwargs):
    num_uuids = kwargs.pop('num_uuids', 0)
    padding = kwargs.pop('padding', '0')
    uuid_str = ''.join(map(attrgetter('hex'), uuids))
    uuid_str += padding * ((num_uuids * 32) - len(uuid_str))
    return uuid_str

# TODO: delete this and use kvlayer._util.split_uuids(), but really use split_key_fragments
def split_uuids(uuid_str):
    return map(lambda s: uuid.UUID(hex=''.join(s)), grouper(uuid_str, 32))

@pytest.mark.skipif('True')
def test_composite_column_names_second_level_range_query_with_decomposited_keys(client):
    '''
    check that we can execute range queries on the second part of a
    CompositeType column name after we unpack the composite key into a
    long string of concatenated hex forms of the UUIDs
    '''
    config = client._config
    namespace = client._app_namespace
    chosen_server = client._chosen_server
    sm = SystemManager(chosen_server)
    sm.create_keyspace(namespace, SIMPLE_STRATEGY, {'replication_factor': '1'})

    family = 'test'
    sm.create_column_family(
        namespace, family, super=False,
        key_validation_class = ASCII_TYPE,
        default_validation_class = BYTES_TYPE,
        comparator_type=UTF8Type(),
        )
    pool = ConnectionPool(namespace, config['storage_addresses'],
                          max_retries=1000, pool_timeout=10, pool_size=2, timeout=120)

    cf = pycassa.ColumnFamily(pool, family)
    u1, u2, u3, u4 = uuid.uuid1(), uuid.uuid1(), uuid.uuid1(), uuid.uuid1()

    cf.insert('inbound', {join_uuids(u1, u2): b''})
    cf.insert('inbound', {join_uuids(u1, u3): b''})
    cf.insert('inbound', {join_uuids(u1, u4): b''})

    ## test range searching
    start  = uuid.UUID(int=u3.int - 1)
    finish = uuid.UUID(int=u3.int + 1)
    assert start.int < u3.int < finish.int
    rec3 = cf.get('inbound',
                  column_start =join_uuids(u1, start),
                  column_finish=join_uuids(u1, finish)).items()
    assert len(rec3) == 1
    assert split_uuids(rec3[0][0])[1] == u3
    ####  This assert above passes!

    ####  This next part fails :-/
    ## now insert many rows -- enough that some should fall in each
    ## subrange below
    for i in xrange(1000):
        cf.insert('inbound', {join_uuids(u1, uuid.uuid4()): b''})

    ## do four ranges, and expect more than zero in each
    step_size = 2**(128 - 2)
    for i in range(2**2, 0, -1):
        start =  uuid.UUID(int=(i-1) * step_size)
        finish = uuid.UUID(int=min(i * step_size, 2**128 - 1))
        recs = cf.get('inbound',
                      column_start =join_uuids(u1, start),
                      column_finish=join_uuids(u1, finish)).items()
        for key, val in recs:
            key = split_uuids(key)
            assert val == b''
            assert key[0] == u1
            assert key[1] < finish
            assert start < key[1]   ## this passes!!

        assert len(recs) > 0
        logger.info( '%r for %r %r' % (len(recs), start, finish))

    sm.close()

@pytest.mark.skipif('True')
def test_cql_paging(client):
    '''
    read rows from a range of tokens in Cassandra
    '''
    config = client._config
    server = config['storage_addresses'][0]
    server = server.split(':')[0]
    conn = cql.connect(server, cql_version='3.0.1')
    cursor = conn.cursor()

    try:
        cursor.execute('DROP KEYSPACE %s;' % client._app_namespace)
    except:
        pass

    ## make a keyspace and a simple table
    ## cursor.execute("CREATE KEYSPACE %s WITH strategy_class = 'SimpleStrategy' AND strategy_options:replication_factor = 1;" % client._app_namespace)
    cursor.execute("CREATE KEYSPACE %s WITH REPLICATION = {'class' : 'SimpleStrategy', 'replication_factor' : 1};" % client._app_namespace)
    cursor.execute("USE %s;" % client._app_namespace)
    cursor.execute('CREATE TABLE data (k int PRIMARY KEY, v varchar);')

    ## put some data in the table
    count = 0
    start_time = time.time()
    for idx in xrange(10):
        cursor.execute("INSERT INTO data (k, v) VALUES (%d, '%s');" % (idx, one_mb))
        count += 1
        if count % 10 == 0:
            elapsed = time.time() - start_time
            rate = float(count) / elapsed
            logger.info( '%d in %.1f sec --> %.1f row/sec' % (count, elapsed, rate) )

    cql3_command = "SELECT k, v FROM data limit 30000;"
    logger.info( cql3_command )

    cursor.execute(cql3_command)

    start = time.time()
    for row in cursor:
        logger.info( time.time() - start )
        start = time.time()

    #print "Number of results: ", len(list(cursor))

    ## remove this test keyspace
    try:
        cursor.execute('DROP KEYSPACE %s;' % client._app_namespace)
    except:
        pass

    cursor.close()
    conn.close()


@pytest.mark.skipif('True')
def test_cql_token_range(client):
    '''
    read rows from a range of tokens in Cassandra
    '''
    config = client._config
    namespace = client._app_namespace
    chosen_server = client._chosen_server
    server = config['storage_addresses'][0]
    server = server.split(':')[0]
    conn = cql.connect(server, cql_version='3.0.1')
    cursor = conn.cursor()

    ## make a keyspace and a simple table
    ## cursor.execute("CREATE KEYSPACE %s WITH strategy_class = 'SimpleStrategy' AND strategy_options:replication_factor = 1;" % client._app_namespace)
    try:
        ## remove this test keyspace
        cursor.execute('DROP KEYSPACE %s;' % client._app_namespace)
    except:
        pass
    cursor.execute("CREATE KEYSPACE %s WITH REPLICATION = {'class' : 'SimpleStrategy', 'replication_factor' : 1};" % client._app_namespace)
    cursor.execute("USE %s;" % client._app_namespace)
    cursor.execute('CREATE TABLE data (k int PRIMARY KEY, v varchar);')

    ## put some data in the table
    count = 0
    total_inbound = 100
    start_time = time.time()
    for idx in xrange(total_inbound):
        cursor.execute("INSERT INTO data (k, v) VALUES (%d, '%s');" % (idx, ' ')) #one_mb))
        count += 1
        if count % 10 == 0:
            elapsed = time.time() - start_time
            rate = float(count) / elapsed
            logger.info( '%d in %.1f sec --> %.1f row/sec' % (count, elapsed, rate) )

    assert count == total_inbound

    ## split up the full range of tokens.
    ## Suppose there are 2**k workers:
    k = 2 # --> four workers
    token_sub_range = 2**(64 - k)
    lowest_token = -1 * 2**63  ## for murmur3hash

    total_count = 0

    for worker_i in range(2**k):
        start_token = lowest_token +      worker_i  * token_sub_range
        end_token =   lowest_token + (1 + worker_i) * token_sub_range
        end_token = min(2**63 - 1, end_token)
        if worker_i == 2**k:
            ## catch the end
            pass#end_token += 1

        limit = 1000
        cql3_command = "SELECT k, v FROM data WHERE token(k) >= %d AND token(k) < %d LIMIT %d;" % (start_token, end_token, limit)
        logger.info( cql3_command )

        cursor.execute(cql3_command)

        start = time.time()
        for row in cursor:
            logger.info( time.time() - start )
            start = time.time()

            total_count += 1

    assert total_count == total_inbound

    ## remove this test keyspace
    cursor.execute('DROP KEYSPACE %s;' % client._app_namespace)

    cursor.close()
    conn.close()
