import os

from cassandra import ConsistencyLevel
from cassandra.query import BatchStatement, BatchType
from yt.wrapper.client import Yt
from webmaster3.cass import Cass, Keyspaces
from webmaster3.task import Task


def run(cluster):
    global SESSION, YT_CLIENT

    SESSION = Cass.cluster(cluster).connect()
    SESSION.set_keyspace(Keyspaces.WEBMASTER3_HISTORY)
    SESSION.default_consistency_level = ConsistencyLevel.ONE
    SESSION.default_timeout = 1200

    ins_q = SESSION.prepare("INSERT INTO webmaster3_history.legacy_tci_history (host_id, tci) VALUES (?, ?)")

    print "Reading Yt table..."
    it = YT_CLIENT.read_table('//home/webmaster/users/leonidrom/tci_history_latest')
    table_data = []
    for r in it:
        table_data.append(r)

    print "Uploading to Cassandra..."
    count = 0
    b = []
    for r in table_data:
        if count % 1000 == 0:
            print "Processed: " + str(count)
        count += 1

        b.append(r)
        if len(b) == 1000:
            write_batch(b, ins_q)
            b = []

    if len(b) > 0:
        write_batch(b, ins_q)

    print "Done."


def write_batch(b, ins_q):
    global SESSION

    batch = BatchStatement(consistency_level=ConsistencyLevel.ONE, batch_type=BatchType.UNLOGGED)

    for r in b:
        host_id = r['host_id']
        tci = r['tci']
        batch.add(ins_q, (host_id, tci))

    SESSION.execute(batch)

def get_yt_token():
    yt_token = os.getenv('YT_TOKEN', None)
    if yt_token:
        return yt_token

    from os.path import expanduser
    home_path = expanduser("~")

    try:
        with open(home_path + '/' + '.yt/token') as f:
            l = list(f)
            if l:
                return l[0].strip()
    except:
        return None


YT_CLIENT = Yt(proxy='hahn.yt.yandex.net', token=get_yt_token())

Task.run(run, "test")
