# coding: utf-8

from saas.tools.devops.billing_send.quotas_lib import AbcServiceQuota
from saas.tools.devops.billing_send.pq_sender import Sender
from saas.library.python.token_store import PersistentTokenStore

import yt.wrapper as yt

import copy
import json
import time

YT_LOCKS_PATH = '//home/saas/ssm/billing_locks'
MIN_INTERVAL_SEC = 600


def get_base_metric(service_quota):
    id = 'saas_' + str(service_quota.abc_service.service_id) + '_' + service_quota.abc_service.slug
    return {
        'id': id,
        'abc_id': service_quota.abc_service.service_id,
        'version': 'v1alpha1',
        'source_id': 'saas-billing-send',
        'source_wt': int(time.time()),
        'labels': {},
        'tags': {},
        'usage': {}
    }

def create_metrics_messages(service_quota):
    res = []
    cpus = service_quota.cpu_guarantee
    for cdc, cusage in cpus.items():
        res.append(get_base_metric(service_quota))
        res[-1]['id'] += '_cpu_' + cdc
        res[-1]['labels']['geo'] = cdc
        res[-1]['schema'] = 'saas.cpu.allocated.v1'
        res[-1]['usage']['quantity'] = cusage.cores
        res[-1]['usage']['unit'] = 'vcpu*second'
    ram = service_quota.ram_guarantee
    for rdc, rusage in ram.items():
        res.append(get_base_metric(service_quota))
        res[-1]['id'] += '_ram_' + rdc
        res[-1]['labels']['geo'] = rdc
        res[-1]['schema'] = 'saas.memory.allocated.v1'
        res[-1]['usage']['quantity'] = rusage.get_bytes()
        res[-1]['usage']['unit'] = 'byte*second'
    hdd = service_quota.hdd_volume
    for hdc, husage in hdd.items():
        res.append(get_base_metric(service_quota))
        res[-1]['id'] += '_hdd_' + hdc
        res[-1]['labels']['geo'] = hdc
        res[-1]['schema'] = 'saas.hdd_storage.allocated.v1'
        res[-1]['usage']['quantity'] = husage.get_bytes()
        res[-1]['usage']['unit'] = 'byte*second'
    ssd = service_quota.ssd_volume
    for sdc, susage in ssd.items():
        res.append(get_base_metric(service_quota))
        res[-1]['id'] += '_ssd_' + sdc
        res[-1]['labels']['geo'] = sdc
        res[-1]['schema'] = 'saas.ssd_storage.allocated.v1'
        res[-1]['usage']['quantity'] = susage.get_bytes()
        res[-1]['usage']['unit'] = 'byte*second'
    return res


def get_time_intervals(last_ts):
    now = int(time.time())
    if now - last_ts < MIN_INTERVAL_SEC:
        return []
    was_hour = last_ts - (last_ts % 3600)
    res = []
    ts_1 = last_ts
    ts_2 = min(was_hour + 3600, now)
    while ts_1 < now:
        res.append((ts_1, ts_2))
        ts_1 = ts_2
        ts_2 = min(ts_2 + 3600, now)
    return res


def process_abc_id(service_quota, sender):
    # get lock
    id = service_quota.abc_service.service_id
    yt_client = yt.YtClient(proxy='hahn', token=PersistentTokenStore.get_token_from_store_env_or_file('yt'))
    yt_locker_client = yt.YtClient(proxy='hahn', token=PersistentTokenStore.get_token_from_store_env_or_file('yt'))
    yt_state_path = YT_LOCKS_PATH + '/' + str(id) + '_' + service_quota.abc_service.slug
    yt_lock_path = YT_LOCKS_PATH + '/lock_' + str(id) + '_' + service_quota.abc_service.slug
    if not yt_locker_client.exists(yt_lock_path):
        yt_locker_client.write_table(yt_lock_path, {}, format=yt.JsonFormat(encoding="utf-8"))
    with yt_locker_client.Transaction():
        yt_locker_client.lock(yt_lock_path, mode='exclusive', wait_for=10000)
        if not yt_client.exists(yt_state_path):
            yt_client.write_table(yt_state_path, {}, format=yt.JsonFormat(encoding="utf-8"))
        data = list(yt_client.read_table(yt_state_path, format=yt.JsonFormat(encoding="utf-8")))
        # print data
        last_ts = int(time.time()) - MIN_INTERVAL_SEC
        new_ts = int(time.time())
        if len(data) > 0:
            last_ts = data[0].get('ts', last_ts)
            data[0]['ts'] = new_ts
        else:
            data.append({"ts": new_ts})
        metrics_templ = create_metrics_messages(service_quota)
        if len(metrics_templ) == 0:
            yt_client.write_table(yt_state_path, data, format=yt.JsonFormat(encoding="utf-8"))
            return
        intervals = get_time_intervals(last_ts)
        for intv in intervals:
            metrics = copy.deepcopy(metrics_templ)
            for metr in metrics:
                metr['usage']['start'] = intv[0]
                metr['usage']['finish'] = intv[1]
                metr['usage']['quantity'] = int(metr['usage']['quantity'] * (intv[1] - intv[0]))
            metrics = [m for m in metrics if m['usage']['quantity'] > 0]
            if len(metrics) != len(metrics_templ):
                print("Zero quantities found for abc: " + str(id))
            message = json.dumps(metrics)
            data[0]['ts'] = intv[1]
            data[0]['message'] = message
            data[0]['sent'] = 'false'
            with yt_client.Transaction():
                yt_client.write_table(yt_state_path, data, format=yt.JsonFormat(encoding="utf-8"))
            # send to lb here
            sender.send(message)
            print(metrics)
            data[0]['sent'] = 'true'
            with yt_client.Transaction():
                yt_client.write_table(yt_state_path, data, format=yt.JsonFormat(encoding="utf-8"))
    pass


def main():
    # yt get
    service_quotas_data = AbcServiceQuota.load_all_from_yt()
    sender = Sender()

    # parse abc ids
    for q in service_quotas_data:
        print(q.abc_service.service_id)
        if not q.abc_service.service_id:
            print ('Empty service id found, skipping it')
            continue
        process_abc_id(q, sender)
    pass

if __name__ == '__main__':
    main()
