#!/usr/bin/env python
# -*- coding: utf-8 -*-i

import yt.wrapper as yt
import sys
import argparse
import logging


NUM_TRIES = 3


def parse_options():
    parser = argparse.ArgumentParser(
        description=u"Рекурсивно размаунчивает каталог в yt, параметры для yt можно передавать через переменные окружения"
    )

    parser.add_argument('dirs', type=str, nargs='+', help=u'Каталоги для размаунчивания')
    parser.add_argument('-c', '--cluster', type=str, help=u'Название кластера')
    parser.add_argument('-t', '--token-path', dest='token_path', type=str, help=u'Путь до файла с токеном')
    opts, extra = parser.parse_known_args()

    if len(extra) > 0:
        die("There are unknown parameters")

    if opts.token_path:
        yt.config["token_path"] = opts.token_path
    elif not yt.config["token_path"]:
        yt.config["token_path"] = '/etc/direct-tokens/yt_robot-direct-yt-ro'

    if opts.cluster:
        yt.config.set_proxy(opts.cluster)

    return opts


def main():
    logging.basicConfig(format=u'%(levelname)-8s [%(asctime)s] %(message)s', level=logging.INFO)
    opts = parse_options()
    success_num = 0
    fails_num = 0

    for yt_dir in opts.dirs:
        logging.info("Processing '%s':" % yt_dir)
        for table in yt.search(yt_dir, node_type=["table"],
                               object_filter=lambda obj: obj.attributes.get("tablet_state") == "mounted",
                               attributes=["tablet_state"]):
            for i in xrange(NUM_TRIES):
                try:
                    logging.info("Trying to unmount table '%s'" % table)

                    yt.unmount_table(table)

                    logging.info("Success")
                    success_num += 1
                    break
                except:
                    if i + 1 == NUM_TRIES:
                        logging.exception("Can't unmount table '%s', attempt %d" % (table, i + 1))
                        fails_num += 1
                    else:
                        logging.error("Can't unmount table '%s', attempt %d" % (table, i + 1))

    logging.info("Total successes: %d" % success_num)
    logging.info("Total fails: %d" % fails_num)
                

if __name__ == '__main__':
    main()
