import google.protobuf.text_format as text_format
import time
import argparse
import logging
import yt.wrapper as yt
from uuid import uuid4
from saas.protos.shards_pb2 import TShards
from saas.tools.devops.lib23.deploy_manager_api import DeployManagerApiClient
from saas.tools.devops.lib23.saas_slot import Slot


def parse_cmd_args():

    description = """
Publishing replic state to yt in standaline indexer format
You need to have defined OAuth tokens for YT in your ENV
"""

    parser = argparse.ArgumentParser(description=description)
    parser.add_argument(
        '-t', '--ctype',
        required=True,
        type=str,
        default='stable',
        help='ctype of service'
    )

    parser.add_argument(
        '--index-path',
        required=True,
        type=str,
        help='Yt path where to store table'
    )

    parser.add_argument(
        '-p', '--result-path',
        required=True,
        type=str,
        help='Yt path where to publish table'
    )

    parser.add_argument(
        '-yt', '--yt-proxy',
        required=True,
        type=str,
        help='Yt cluster'
    )
    parser.add_argument(
        '-s', '--service',
        required=True,
        type=str,
        help='Service to detach'
    )

    parser.add_argument(
        '-d', '--dry-run',
        default=False,
        action='store_true',
        help='Print controller commands and exit'
    )

    parser.add_argument(
        '-i', '--replic-id',
        type=int,
        default=1,
        help='replica specification'
    )

    return parser.parse_args()


def get_detach_params(yt_proxy, path):
    return {
        'async': 'yes',
        'upload_to_yt': '1',
        'yt_proxy': yt_proxy,
        'yt_path': path
    }


def publish_index(intervals, timestamp, index_path, comment=""):
    shards = TShards()
    for i in range(len(intervals)):
        shard_min, shard_max = intervals[i]
        shard_name = 'detach_{0}_{1}_{2}'.format(timestamp, shard_min, shard_max)
        table_path = index_path + '[({0}u,{1}u)]'.format(shard_min, shard_max)
        shards.Shard.add(Name=shard_name,
                         ShardMin=shard_min,
                         ShardMax=shard_max,
                         Timestamp=timestamp,
                         YTIndexTable=table_path)

    shards.Comment = comment
    return text_format.MessageToString(shards)


def main():
    args = parse_cmd_args()
    ctype = args.ctype
    service_name = args.service
    replic_id = args.replic_id
    index_path = args.index_path
    publish_path = args.result_path
    yt_proxy = args.yt_proxy
    client = DeployManagerApiClient()
    shard_instances = client.get_slots_by_interval(ctype, service_name)
    shard_paths = []
    intervals = []
    current_timestamp = int(time.time())
    yt.config.set_proxy(yt_proxy)
    tmp_dir = '//tmp/saas/detach_to_yt/' + str(current_timestamp) + "-" + uuid4().hex[:8]

    if not args.dry_run:
        yt.create(path=tmp_dir, type='map_node', recursive=True, force=True)
        yt.create(path=index_path, type='table', recursive=True, force=True)
        yt.create(path=publish_path, type='file', recursive=True, force=True)

    for shard_interval in shard_instances:
        interval_detached = False
        for slot_info in shard_interval['slots']:
            if (not interval_detached and slot_info['replic_id'] == replic_id
                        and slot_info['result.controller_status'] == 'Active'):
                host, port = slot_info['slot'].split(':')
                path = tmp_dir + '/' + slot_info['interval']
                slot = Slot(host, port, shards_min=slot_info['$shards_min$'],
                            shards_max=slot_info['$shards_max$'])

                if not args.dry_run:
                    slot.detach(get_detach_params(yt_proxy, path))
                print("detaching", slot, get_detach_params(yt_proxy, path))
                shard_paths.append(path)
                intervals.append((slot.shards_min, slot.shards_max))
                interval_detached = True
        if not interval_detached:
            logging.error('cannot find active instance with such replic id for interval %d %d',
                          shard_interval['$shards_min$'],
                          shard_interval['$shards_max$'])

    if not args.dry_run:
        while len(yt.list(tmp_dir)) < len(shard_paths):
            time.sleep(10)
        yt.concatenate(shard_paths, index_path)
        yt.run_sort(args.index_path,
                    sort_by=['shard_min', 'shard_max',
                             'segment_id', 'name', 'part_index'])
    publish_comment = 'Detached from backends, ' + str(args)
    published_index = publish_index(intervals, current_timestamp, index_path, publish_comment)
    if args.dry_run:
        print(published_index)
    else:
        yt.write_file(publish_path, published_index)
        yt.remove(path=tmp_dir, recursive=True, force=True)
