#!/usr/bin/python3

import itertools
import json
import uuid
import argparse
import requests
import os


def set_node_pos(node, x, y, block_size):
    node['graph_position'] = {
        'width': block_size[0],
        'height': block_size[1],
        'x': x,
        'y': y
    }


def make_task(name, params, confirm):
    task = {
        'id': str(uuid.uuid4()),
        'data': {
            'name': name,
            'params': params
        },
        'flags': {
            'manually_confirmed': confirm,
            'skipped': False,
            'banned': False
        }
    }
    return task


def prepare_task(service, confirm=False):
    return make_task(
        'set_snapshot_target_state',
        {
            'service_id': service,
            'wait_for_success': True,
            'attach_committed_tickets': False,
            'target_state': 'PREPARED',
        },
        confirm=confirm,
    )


def activate_task(service, confirm=False):
    return make_task(
        'set_snapshot_target_state',
        {
            'service_id': service,
            'wait_for_success': True,
            'attach_committed_tickets': False,
            'target_state': 'ACTIVE',
        },
        confirm=confirm,
    )


def empty_task(confirm):
    return make_task('empty', {}, confirm)


def connect_tasks(lhs, rhs):
    return {
        'parent': {'id': lhs['id']},
        'child': {'id': rhs['id']},
        'id': str(uuid.uuid4())
    }


def render(services):
    y = 0
    X_SIZE = 220
    Y_SIZE = 100
    X_GAP = X_SIZE + 40
    Y_GAP = Y_SIZE + 20

    tasks = []
    dependencies = []

    def calc_task_y(index, tasks, max_tasks):
        return Y_GAP * (index + 1) * max_tasks / (tasks + 1)

    for phases in services:
        max_tasks = 1

        for phase in phases:
            if type(phase) == list:
                max_tasks = max(max_tasks, len(phase))

        start_task = empty_task(confirm=False)
        tasks.append(start_task)
        set_node_pos(start_task, 0, y + calc_task_y(0, 1, max_tasks), (X_SIZE, Y_SIZE))
        dependencies.append(connect_tasks({'id': 'start'}, start_task))

        prev_phase = [start_task]

        for phase_index, phase in enumerate(phases):
            if type(phase) != list:
                phase = [phase]

            for task_index, task in enumerate(phase):
                tasks.append(task)
                set_node_pos(
                    task,
                    (phase_index + 1) * X_GAP,
                    y + calc_task_y(task_index, len(phase), max_tasks),
                    (X_SIZE, Y_SIZE),
                )

                for prev_task in prev_phase:
                    dependencies.append(connect_tasks(prev_task, task))

            prev_phase = phase

        y += max_tasks * Y_GAP

    return tasks, dependencies, y / 2


def deploy_service(
    name,
    name_yp=None,
    testing_names=(),
    confirm_testing=False,
    dcs=('man', 'sas', 'vla'),
    dcs_yp=(None, 'sas', 'vla'),
):
    prepare_phase = []

    for testing_name in testing_names:
        prepare_phase.append(prepare_task(testing_name))

    for dc in dcs:
        prepare_phase.append(prepare_task('{}_{}'.format(name, dc)))

    phases = [prepare_phase]

    for index, testing_name in enumerate(testing_names):
        confirm = index == 0 and confirm_testing
        phases.append(activate_task(testing_name, confirm=confirm))

    for index, (dc, dc_yp) in enumerate(itertools.zip_longest(dcs, dcs_yp)):
        confirm = index == 0
        if name_yp and dc_yp:
            phases.append(activate_task('{}-{}'.format(name_yp, dc_yp), confirm=confirm))
        if dc:
            phases.append(activate_task('{}_{}'.format(name, dc), confirm=confirm and not dc_yp))

    return phases


services = [
    [
        [
            prepare_task('balancer_l7leveler_kubr_man_testing'),
            prepare_task('balancer_l7leveler_kubr_sas_testing'),
            prepare_task('balancer_l7leveler_kubr_vla_testing'),
            prepare_task('balancer_l7leveler_yaru_sas_testing'),
            prepare_task('balancer_l7leveler_kubr_man_only'),
            prepare_task('balancer-l7leveler-kubr-yp-sas-only'),
            prepare_task('balancer_l7leveler_kubr_sas_only'),
            prepare_task('balancer-l7leveler-kubr-yp-vla-only'),
            prepare_task('balancer_l7leveler_kubr_vla_only'),
        ],
        activate_task('balancer_l7leveler_kubr_man_testing'),
        activate_task('balancer_l7leveler_kubr_sas_testing'),
        activate_task('balancer_l7leveler_kubr_vla_testing'),
        activate_task('balancer_l7leveler_yaru_sas_testing'),
        activate_task('balancer_l7leveler_kubr_man_only'),
        activate_task('balancer-l7leveler-kubr-yp-sas-only'),
        activate_task('balancer_l7leveler_kubr_sas_only'),
        activate_task('balancer-l7leveler-kubr-yp-vla-only'),
        activate_task('balancer_l7leveler_kubr_vla_only'),
    ],
    deploy_service(
        'balancer_l7leveler_kubr',
        'balancer-l7leveler-kubr-yp',
        testing_names=['balancer_l7leveler_kubr_vla_experiment'],
        confirm_testing=True,
    ),
    deploy_service('balancer_knoss_search_yp'),
    deploy_service('balancer_knoss_morda_yp'),
    deploy_service(
        'rtc_balancer_knoss_fast_yp',
        testing_names=[
            'rtc_balancer_knoss_fast_testing_yp_man',
            'rtc_balancer_knoss_fast_testing_yp_sas',
            'rtc_balancer_knoss_fast_testing_vla'
        ],
    ),
    deploy_service('production_balancer_any'),
    deploy_service('balancer_l7leveler_yaru', 'balancer-l7leveler-yaru-yp'),
    deploy_service('balancer_proxy_yandex_ua', 'balancer-proxy-yandex-ua-yp'),
]


tasks, dependencies, start_y = render(services)

recipe = {
    'id': 'default',
    'content': {
        'name': 'production',
        'desc': 'Руками не редактировать, генерируется https://a.yandex-team.ru/arc/trunk/arcadia/balancer/production/recipe_generator/generate_recipe.py',
        'dependencies': dependencies,
        'tasks': tasks,
        'start': {
            'graph_position': {
                'height': 30,
                'width': 30,
                'x': -70,
                'y': start_y,
            }
        },
    },
}

parser = argparse.ArgumentParser('Balancer deploy recipe generator')
parser.add_argument('--commit', action='store_true', help='Commit new recipe to nanny')
args = parser.parse_args()

if args.commit:
    print(
        requests.post(
            'http://nanny.yandex-team.ru/v2/services_dashboards/catalog/production_balancer_l7heavy/recipes/',
            data=json.dumps(recipe),
            headers={
                'Content-Type': 'application/json',
                'Authorization': 'OAuth {}'.format(os.environ['OAUTH_NANNY']),
            },
        ).json()
    )
else:
    print(json.dumps(recipe, indent=2, sort_keys=True))
