#!/usr/bin/env python3

import argparse
import datetime
import requests
import time
from pathlib import Path

import balancer.production.x.perf_tests.config as config
import balancer.production.x.perf_tests.golovan as golovan
import balancer.production.x.perf_tests.run as run


STARTREK_TOKEN = '/.startrek/token'
SLEEP_INTERVAL = 5  # for checking status
GOLOVAN_LAG = 10  # golovan need some time to get and process data
GOLOVAN_INTERVAL = 55  # standart period for calculating result in golovan
CLEANUP_SLEEP = 60 * 4  # time to sleep with lowest weight during cleanup
CPU_LIMITER_DISABLED = 'cpu_limiter_disabled'  # file name

WARDEN_DC = ['MAN', 'SAS', 'VLA']
WARDEN_COMPONENT = 'l7'

# Stopping criteria
TERM_CRITICAL_FAILS = 0.5
KNOSS_CRITICAL_FAILS = 1.2
DESYNC_CRITICAL = 10


def cleanup(opts):
    print('Cleanup')
    set_weight_and_sleep(opts, '1', CLEANUP_SLEEP, False)
    script = 'chown ' + opts.cont_user_group + ' ' + opts.controls_dir
    run.script(opts, script)
    script = 'rm ' + opts.controls_dir + '/' + CPU_LIMITER_DISABLED
    run.script(opts, script)
    script = 'rm ' + opts.controls_dir + '/' + opts.weight_file
    run.script(opts, script)


def prepare_to_start(opts):
    weight_file = opts.controls_dir + '/' + opts.weight_file
    check_word = 'exists42'
    script = 'if [ -f \"' + weight_file + '\" ]; then echo \"' + check_word + '\"; fi'
    need_cleanup = True if run.host_script(opts, script, False).find(check_word) != -1 else False
    if not need_cleanup and opts.hosts.has_old():
        need_cleanup = True if run.host_script(opts, script, True).find(check_word) != -1 else False

    if need_cleanup:
        print('Cleanup should be done first')
        cleanup(opts)

    script = 'chown root:root ' + opts.controls_dir \
             + ' && touch ' + opts.controls_dir + '/' + CPU_LIMITER_DISABLED \
             + ' && echo "10" > ' + opts.controls_dir + '/' + opts.weight_file
    run.script(opts, script)


def restart_balancer(opts):
    script = 'echo "Restarting balancer..."' \
             ' && curl -4i "http://localhost:' + str(opts.hosts.admin_port) + \
             '/admin?action=graceful_shutdown&timeout=40s&close_timeout=5s"' \
             ' && while [ `ps aux | grep "balancer-' + str(opts.hosts.port) + \
             ' balancer.cfg" | wc -l` != 5 ]; do echo "wait"; sleep 1; done' \
             ' && echo "Restarted!"' \
             ' && date'
    run.script(opts, script)


def check_fails(opts):
    fails = golovan.total_fail_summ(opts.hosts.new, opts.hosts.gencfg, opts.monitoring_params,
                                    time.time() - GOLOVAN_LAG, GOLOVAN_INTERVAL - GOLOVAN_LAG)
    total = golovan.total_requests_summ(opts.hosts.new, opts.hosts.gencfg, opts.monitoring_params,
                                        time.time() - GOLOVAN_LAG, GOLOVAN_INTERVAL - GOLOVAN_LAG)
    assert total

    perc_failed = float(fails) / total * 100.0

    return perc_failed


def check_desyncs_per_cycle(opts):
    desync = golovan.worker_cpu_usage_time_desync_summ(opts.hosts.new, opts.hosts.gencfg, opts.monitoring_params,
                                                       time.time() - GOLOVAN_LAG, GOLOVAN_INTERVAL - GOLOVAN_LAG)
    conts = golovan.worker_conts_cycle_summ(opts.hosts.new, opts.hosts.gencfg, opts.monitoring_params,
                                            time.time() - GOLOVAN_LAG, GOLOVAN_INTERVAL - GOLOVAN_LAG)

    assert conts

    desync_perc = float(desync) / conts * 100.0

    return desync_perc


def check_is_ok(opts):
    perc_failed = check_fails(opts)
    crit_perc_failed = TERM_CRITICAL_FAILS if opts.is_term else KNOSS_CRITICAL_FAILS
    if perc_failed >= crit_perc_failed:
        print('Too many fails: %.2f' % perc_failed, '%')
        return False

    desync_perc = check_desyncs_per_cycle(opts)
    if desync_perc >= DESYNC_CRITICAL:
        print('Too many desyncs per cycle: %.2f' % desync_perc, '%')
        return False

    return True


def set_weight_and_sleep(opts, weight, timeout, check=True):
    print(datetime.datetime.now().strftime('%H:%M:%S '), 'Set', weight, 'and sleep', timeout)
    end_time = time.time() + timeout

    file = opts.controls_dir + '/' + opts.weight_file
    script = 'echo ' + weight + ' > ' + file + ' && cat ' + file + ' && date'
    run.script(opts, script)

    while end_time > time.time():
        check_start = time.time()
        if check and not check_is_ok(opts):
            return False
        check_time = time.time() - check_start
        if SLEEP_INTERVAL > check_time:
            time.sleep(SLEEP_INTERVAL - check_time)

    return True


def make_highload(opts, max_weight):
    if opts.has_target_rp5s():  # set 80% weight on first step
        #  get current rp5s, check dist to target, increase weight steps
        rp5s = golovan.total_requests_summ(opts.hosts.new, opts.hosts.gencfg, opts.monitoring_params,
                                           time.time() - GOLOVAN_LAG, GOLOVAN_INTERVAL - GOLOVAN_LAG)
        assert rp5s
        print('Start rp5s = ', rp5s)
        weight = 10 * (opts.target_rp5s / rp5s)
        timeout = 60 * 4
        start_time = round(time.time())
        for i in [0.8, 1.5, 1.3, 1.2, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1]:
            weight = int(weight * i)
            if weight >= max_weight:
                weight = max_weight
                print('Make target highload: weight = ', weight, ', timeout = ', timeout)
                set_weight_and_sleep(opts, str(weight), timeout)
                break
            print('Make target highload: weight = ', weight, ', timeout = ', timeout)
            if not set_weight_and_sleep(opts, str(weight), timeout):
                break
    else:
        weight = 20
        timeout = 90
        start_time = round(time.time())
        print('Make highload: weight = ', weight, ', timeout = ', timeout)
        while set_weight_and_sleep(opts, str(weight), timeout):
            weight += 10
            if weight >= max_weight:
                weight = max_weight
                print('Make highload: weight = ', weight, ', timeout = ', timeout)
                set_weight_and_sleep(opts, str(weight), timeout)
                break
            print('Make highload: weight = ', weight, ', timeout = ', timeout)

    end_time = round(time.time())
    set_weight_and_sleep(opts, '1', 15, False)  # just in case here, cleanup would set correct weight

    return start_time, end_time


def calculate_host_result(host, gencfg, monitoring_params, end):
    rp5s = golovan.total_requests_summ(host, gencfg, monitoring_params, end, GOLOVAN_INTERVAL)
    assert rp5s > 0

    return rp5s


def calculate_result(opts, end):
    new_rp5s = calculate_host_result(opts.hosts.new, opts.hosts.gencfg, opts.monitoring_params, end)
    if not opts.hosts.has_old():
        old_rp5s = 0
    else:
        old_rp5s = calculate_host_result(opts.hosts.old, opts.hosts.gencfg, opts.monitoring_params, end)

    return new_rp5s, old_rp5s


def make_snapshot(opts, start, end):
    path = opts.monitoring_path + '/?from=' + str(start) + '&to=' + str(end)

    yasm_url = 'https://yasm.yandex-team.ru'
    yasm_screen_url = 'https://s.yasm.yandex-team.ru' + path + '&width=1600&height=900&static=true'

    response = requests.get(yasm_screen_url)
    assert response.url

    snapshot_in_comment = '<{График\n((' + yasm_url + path + ' ' + response.url + '))}>'
    return snapshot_in_comment


def get_host_info(opts, old_host):
    build_script = './balancer --version 2>&1 | grep \'arcadia.yandex.ru/arc/tags/balancer/\''
    output = run.host_script(opts, build_script, old_host)
    build = output.split('/')[-2]

    kernel_script = 'uname -srm'
    kernel = 'Ядро: ' + run.host_script(opts, kernel_script, old_host)

    porto_script = 'portod --version | grep running'
    porto = 'Porto ' + run.host_script(opts, porto_script, old_host)

    return build, kernel, porto


def format_rps_result(rp5s):
    rps = rp5s / 5
    if rps > 1000:
        rps /= 1000
        rp5s /= 1000
        return '**' + ('%.1f' % rps) + 'k** rps (' + ('%.1f' % rp5s) + 'k rp5s)'
    else:
        return '**' + ('%.1f' % rps) + '** rps (' + ('%.1f' % rp5s) + ' rp5s)'


def post_comment(opts, new_rp5s, old_rp5s, snapshot_in_comment):
    f = open(str(Path.home()) + STARTREK_TOKEN, 'r')
    startrek_token = f.read().rstrip()
    date = '{0:%d.%m.%Y}'.format(datetime.datetime.now())

    new_build, new_kernel, new_porto = get_host_info(opts, False)
    new_rps_result = '\nРезультат прогрузки ' + format_rps_result(new_rp5s)
    if opts.hosts.has_old():
        old_build, old_kernel, old_porto = get_host_info(opts, True)
        old_rps_result = '. Результат на старом билде ' + old_build + ' ' + format_rps_result(old_rp5s)

    comment_text = '**' + date + '**\n' + opts.hosts.new.cont_name + '\n'
    comment_text += 'Build: ' + new_build + '\n' + new_kernel + '\n'
    comment_text += new_porto + '\n' + new_rps_result
    if opts.hosts.has_old():
        comment_text += old_rps_result
    comment_text += '\n\n((https://datalens.yandex-team.ru/57g5qn5zhniyx-proverka-perfa-l7-balanserov Дашборд с результатами прогрузок))'
    comment_text += '\n((https://warden.z.yandex-team.ru/components/l7/trainings Результаты в Warden))'
    comment_text += '\n\n' + snapshot_in_comment

    headers = {'Authorization': 'OAuth {}'.format(startrek_token), 'Content-Type': 'application/json'}
    url = 'https://st-api.yandex-team.ru/v2/issues/' + opts.ticket + '/comments'

    r = requests.post(url, headers=headers, json={'text': comment_text})
    r.raise_for_status()

    comment_url = 'https://st.yandex-team.ru/' + opts.ticket + '#' + r.json()['longId']
    return comment_url


def post_result_to_warden(opts, rp5s, comment_url):
    date = '{0:%Y-%m-%d}'.format(datetime.datetime.now())
    dc = opts.hosts.new.name[:3].upper()
    target_rp5s = opts.target_rp5s if opts.has_target_rp5s() else 1
    result = []
    for i in WARDEN_DC:
        if dc == i:
            result.append({"dc": i, "rps": rp5s, "target": target_rp5s})
        else:
            result.append({"dc": i, "rps": 1, "target": 0})

    url = 'http://uchenki.yandex-team.ru/api/v1/training/training_results'
    headers = {'Content-Type': 'application/json'}
    data = {
        "issue": comment_url,
        "full_component_name": WARDEN_COMPONENT + '/' + opts.warden_service,
        "date": date,
        "result": result
    }

    print('Post to Warden', url, headers, data)

    r = requests.post(url, headers=headers, json=data, verify=False)
    r.raise_for_status()


def run_perf_test(opts):
    try:
        prepare_to_start(opts)
        restart_balancer(opts)

        set_weight_and_sleep(opts, '10', 60 * 4, False)

        max_weight = 200 if opts.is_term else 100
        start, end = make_highload(opts, max_weight)

        set_weight_and_sleep(opts, '1', 60 * 5, False)
        set_weight_and_sleep(opts, '10', 10, False)

        new_rp5s, old_rp5s = calculate_result(opts, end)

        if opts.report and opts.has_ticket() and opts.has_warden():
            delta = 60 * 6 * 1000  # 6 min for both borders, to have some picture of before and after
            snapshot = make_snapshot(opts, start * 1000 - delta, end * 1000 + delta)

            comment_url = post_comment(opts, new_rp5s, old_rp5s, snapshot)

            post_result_to_warden(opts, new_rp5s, comment_url)

            print('\nResult: ', comment_url)
        else:
            print('\n\n\nResult: ==============================================')
            print(opts.hosts.new.cont_name, ':', str(int(new_rp5s / 5)), ' rps')
            if opts.hosts.has_old():
                print(opts.hosts.old.cont_name, ':', str(int(old_rp5s / 5)), ' rps')
            print('======================================================\n')

    finally:
        print('Final cleanup')
        cleanup(opts)


def show_usage():
    print('Usage: main.py [-c|--config path_to_config] [-r|--report y|n]')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config')
    parser.add_argument('--l7')
    parser.add_argument('-r', '--report')
    args = parser.parse_args()

    if args.config:
        opts = config.get_options(args.config)
    elif args.l7 == 'exp':
        opts = config.get_options('configs/kubr_exp.yaml')
    elif args.l7 == 'kubr':
        opts = config.get_options('configs/kubr.yaml')
    elif args.l7 == 'morda':
        opts = config.get_options('configs/morda.yaml')
    elif args.l7 == 'search':
        opts = config.get_options('configs/search.yaml')
    else:
        show_usage()
        return

    if args.report:
        if args.report not in ['y', 'n']:
            show_usage()
            return
        if args.report == 'y':
            opts.report = True

    config.print_completed_config(opts)
    print('\nMonitoring: https://yasm.yandex-team.ru{}\n'.format(opts.monitoring_path))

    run_perf_test(opts)


if __name__ == '__main__':
    main()
