#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Генератор solomon графиков и дашбордов

команды:
dump - сделать дамп сущностей из solomon'a в директорию projects
upload - загрузить локальные данные в solomon

формат сущностей: projects_prefix/project_id/[graphs|dashboards]/[entity_id]
projets_prefix - папка с дампом (один уровень)
project_id - id проекта
[graphs|dashboards] - графики или дашборды
entity_id - id сущности (графика, дашборда, возможно еще чего-нибудь в дальнейшем)

примеры:
projects/direct/dashboards
projects/direct-test

Примеры запука:
solomonerator dump projects/direct/dashboards
solomonerator upload projects/projects/direct/graph/4xx_responses

Чтобы применить изменения в solomon'е, нужно указать --do,
иначе скрипт работает в режиме dry-run:
solomonerator upload projects/projects/direct/graph/4xx_responses --do
         
"""

import sys
import os
import argparse
import requests
import json
from collections import OrderedDict
import tempfile
import shutil
import subprocess


SOLOMON_API_URL = 'http://solomon.yandex.net/api/v2'
SVN_PATH = ''
OAUTH_HEADER = {}
COMMON_PREFIX = "projects"
PROJECTS = ["direct", "direct-test"]
ALLOWED_TYPES = ['graphs', 'dashboards', 'clusters', 'services', 'shards', 'alerts', 'notificationChannels']
CMDS = ['upload', 'dump', 'backup']
TO_DO_MESSAGE = "#" * 30 + "\n!!! Changes are not applied !!!" + "\n!!! run with  --do to do it !!!\n" + "#" * 30


def make_request(method, url, data=""):
    headers = OAUTH_HEADER
    if method.upper() != "GET":
        headers['Content-Type'] = 'application/json'

    req = requests.request(
        method, "%s/%s/%s" % (SOLOMON_API_URL, COMMON_PREFIX, url), headers=headers, data=data
    )

    if req.status_code != 200:
        print "ERROR while requesting (%s) %s" % (method.upper(), url)
        print req.content
        return None

    return req.json(object_pairs_hook=OrderedDict)


def get_entity_type_ids(e_project, e_type):
    all_data = []
    if e_type in ['alerts', 'notificationChannels']:
        data = make_request("get", "%s/%s" % (e_project, e_type))
        while 'items' in data:
            all_data.extend(data['items'])
            if 'nextPageToken' not in data:
                break

            data = make_request("get", "%s/%s?pageToken=%s" % (e_project, e_type, data['nextPageToken']))
    else:
        data = make_request("get", "%s/%s?pageSize=all" % (e_project, e_type))
        if data:
            all_data = data['result']

    if not all_data:
        return None

    return [el['id'] for el in all_data]


def get_entity(e_project, e_type, e_id):
    return make_request("get", os.path.join(e_project, e_type, e_id))


def get_entity_local(e_prefix, e_project, e_type, e_id):
    with open(os.path.join(e_prefix, e_project, e_type, e_id), 'r') as fh:
        return json.load(fh, object_pairs_hook=OrderedDict)


def dump_entity(e_prefix, e_project, e_type, e_id, e_data):
    with open(os.path.join(e_prefix, e_project, e_type, e_id), 'w') as fh:
        json.dump(e_data, fh, indent=2)


def update_entity(e_prefix, e_project, e_type, e_id, e_data):
    make_request("put", os.path.join(e_project, e_type, e_id), data=e_data)


def create_entity(e_prefix, e_project, e_type, e_data):
    make_request("post", os.path.join(e_project, e_type), data=e_data)


def parse_entity_path(path):
    parts = path.split('/')

    if len(parts) < 1 or len(parts) > 4:
        print 'ERROR: bad entity path: %s' % path
        return []

    parts += [''] * (4 - len(parts))

    if parts[2] and parts[2] not in ALLOWED_TYPES:
        print 'ERROR: bad entity type, %s is not in [%s]' % (parts[2], "|".join(ALLOWED_TYPES))
        return []

    return parts


def run(args):
    entities = args.entities
    tmpdir = ''

    try:
        if args.cmd == "backup":
            tmpdir = tempfile.mkdtemp()
            os.chdir(tmpdir)
            subprocess.check_output(['svn', 'co', SVN_PATH])

        for entity in entities:
            print "### Processing %s" % entity
            path_parts = parse_entity_path(entity)
            if not path_parts:
                continue
    
            changed = False
            e_prefix, e_project, e_type, e_id = path_parts
    
            projects = [e_project] if e_project else PROJECTS
            solomon_list = set()

            for cur_project in projects:
                types = [e_type] if e_type else ALLOWED_TYPES
    
                for cur_type in types:
                    if args.cmd in ["dump", "backup"]:
                        ids = [e_id] if e_id else get_entity_type_ids(cur_project, cur_type)
    
                        entity_path = os.path.join(e_prefix, cur_project, cur_type)
                        if not os.path.exists(entity_path):
                            os.makedirs(entity_path)
    
                        if not ids:
                            print "WARNING: empty set of ids in %s" % os.path.join(cur_project, cur_type)
                            continue
        
                        for cur_id in ids:
                            print "Dumping %s: %s" % (cur_type[:-1], os.path.join(entity_path, cur_id))
                            data = get_entity(cur_project, cur_type, cur_id)
                            if data:
                                solomon_list.add(os.path.join(cur_project, cur_type, cur_id))
                                dump_entity(e_prefix, cur_project, cur_type, cur_id, data)
    
                    elif args.cmd == "upload":
                        entity_path = os.path.join(e_prefix, cur_project, cur_type)
                        if not os.path.isdir(entity_path):
                            print "ERROR: entity path '%s' does not exist" % entity_path
                            continue
    
                        ids = [e_id] if e_id else os.listdir(entity_path)
                        existing_ids = get_entity_type_ids(cur_project, cur_type)
                        if existing_ids is None:
                            continue
                        existing_ids = set(existing_ids)
    
                        for cur_id in ids:
                            if cur_id not in existing_ids:
                                print "Creating new %s: %s" % (cur_type[:-1], os.path.join(entity_path, cur_id))
                                changed = True
                                if args.do:
                                    create_entity(
                                        e_prefix, cur_project, cur_type,
                                        json.dumps(get_entity_local(e_prefix, cur_project, cur_type, cur_id))
                                    )
                            else:
                                data_local = get_entity_local(e_prefix, cur_project, cur_type, cur_id)
                                if data_local == get_entity(cur_project, cur_type, cur_id):
                                    continue
        
                                print "Updating %s: %s" % (cur_type[:-1], os.path.join(entity_path, cur_id))
                                changed = True
                                if args.do:
                                    update_entity(
                                        e_prefix, cur_project, cur_type, cur_id,
                                        json.dumps(data_local)
                                    )
    
            if args.cmd == "upload" and not changed:
                print "No changes"
    
        if args.cmd == "backup":
            os.chdir(os.path.join(tmpdir, os.path.basename(SVN_PATH)))

            files_to_remove = []
            for root, dirs, files in os.walk('./'):
                files = [f for f in files if not f[0] == '.']
                dirs[:] = [d for d in dirs if not d[0] == '.']
                for file in files:
                    cur_file = os.path.join(root, file)[2:]
                    if cur_file not in solomon_list:
                        files_to_remove.append(cur_file)

            if files_to_remove:
                subprocess.check_output(['svn', 'rm'] + files_to_remove)
            subprocess.check_output(['svn', 'add', '--force', '--depth', 'infinity', '.'])
            changes = subprocess.check_output(['svn', 'st', '.'])
            print '\nChanges solomon vs svn:\n%s' % (changes if changes else "Nothing new\n")
            if args.do and changes:
                subprocess.check_output([
                    'svn', 'commit', '-m',
                    u'Автоматический бэкап данных из соломона (сделан скриптом %s с машины %s)' % (
                        os.path.basename(__file__), subprocess.check_output(['hostname', '-f']).strip()
                    )
                ])
                print "Changes were successfully commited"
            elif changes:
                print TO_DO_MESSAGE
        elif args.cmd == "upload" and not args.do:
            print TO_DO_MESSAGE
            
    finally:
        if args.cmd == "backup":
            shutil.rmtree(tmpdir)
                    

def parse_args():
    global SVN_PATH

    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, description=__doc__)
    parser.add_argument('cmd', choices=CMDS, help='команда для выполнения')
    parser.add_argument('entities', metavar='entity', nargs='*', type=str, help='список сущностей')
    parser.add_argument('-t', '--token', default="~/.solomon-oauth", help='путь до oauth токена для solomon')
    parser.add_argument('-u', '--user', type=str, help='пользователь, от которого ходить в svn')
    parser.add_argument('--arcadia-key', type=str, help='путь до пользовательского ключа для аркадии')
    parser.add_argument(
        '--do', default=False, action='store_true',
        help='привезти изменения в solomon, иначе просто выводится информация'
    )
    args = parser.parse_args()

    SVN_PATH = 'svn+ssh://%sarcadia.yandex.ru/robots/trunk/directadmin-data/solomonerator/projects' % (
        args.user + '@' if args.user else ''
    )

    if args.cmd != 'backup':
        if not args.entities:
            parser.error('укажите хотя бы одну сущность')
    else:
        if args.entities:
            parser.error('команда backup не принимает отдельные сущности')
        args.entities = [os.path.basename(SVN_PATH)]

    if args.cmd == 'backup':
        if args.user and not args.arcadia_key:
            parser.error('укажите путь до ключа для пользователя %s' % args.user)
        if not args.user and args.arcadia_key:
            parser.error('укажите имя пользователя, от которого хотите сходить в svn')

    if args.arcadia_key:
        os.environ['SVN_SSH'] = 'ssh -i %s' % args.arcadia_key

    return args


def main():
    global OAUTH_HEADER
    args = parse_args()

    try:
        with open(os.path.expanduser(args.token), 'r') as fh:
            OAUTH_HEADER = {'Authorization': 'OAuth ' + fh.read().strip()}
    except:
        sys.exit("can't get solomon oauth token")

    run(args)


if __name__ == '__main__':
    main()

