#!/usr/bin/python3
"""
Шорткаты для создания реплицированной таблицы ролей и реплик этой таблицы в YT.

0. Предварительно обязательно нужно создать все необходимые папки на всех репликах,
   в которых планируем создавать таблицы.

1. Создаём реплицированную таблицу
$ ./yt_export.py new-replicated --cluster markov --name //home/idm/testing/roles

2. Создаём синхронную и асинхронные реплики
$ ./yt_export.py new-replica --cluster markov --name //home/idm/testing/roles \
      --replica-cluster seneca-vla --replica-name //home/idm/testing/roles --sync
$ ./yt_export.py new-replica --cluster markov --name //home/idm/testing/roles --replica-cluster seneca-man
$ ./yt_export.py new-replica --cluster markov --name //home/idm/testing/roles --replica-cluster seneca-sas

Подробнее о репл.таблицах можно почитать здесь
https://yt.yandex-team.ru/docs/description/dynamic_tables/replicated_dynamic_tables
"""
import argparse
import subprocess
import sys


SCHEMA = """
<"unique_keys" = %true;"strict" = %true;>
[
    {
        "name" = "tvm_id";
        "required" = %true;
        "type_v3" = "uint32";
        "sort_order" = "ascending";
    };
    {
        "name" = "revision";
        "required" = %true;
        "type_v3" = "uint64";
        "sort_order" = "ascending";
    };
    {
        "name" = "unixtime";
        "required" = %true;
        "type_v3" = "timestamp";
    };
    {
        "name" = "meta";
        "required" = %false;
        "type_v3" = {"type_name" = "optional"; "item" = "yson"};
    };
    {
        "name" = "blob";
        "required" = %true;
        "type_v3" = "string";
    };
]
"""


class YTCommand:

    def __init__(self, args):
        parser = argparse.ArgumentParser()
        self.add_arguments(parser)
        self.options = parser.parse_args(args)

    def add_arguments(self, parser):
        pass

    def yt(self, command, *, remove_line_breaks=True, schema=SCHEMA, **kwargs):
        if remove_line_breaks:
            command = command.replace('\n', ' ')
        command = command.format(o=self.options, schema=schema, **kwargs)
        return self._subprocess(f'ya tool yt {command}')

    def mount_table(self, cluster=None, name=None):
        cluster = cluster or self.options.cluster
        name = name or self.options.name
        self.yt("--proxy {cluster} mount-table {name}", cluster=cluster, name=name)

    @staticmethod
    def _subprocess(command):
        result = subprocess.run(command, shell=True, capture_output=True)
        if result.returncode != 0:
            print(result.stderr.decode('utf-8'))
            raise subprocess.CalledProcessError(result.returncode, command)
        print(result.stdout.decode('utf-8'))
        return result


class NewReplicated(YTCommand):

    def add_arguments(self, parser):
        parser.add_argument('--cluster', required=True)
        parser.add_argument('--name', required=True)

    def __call__(self):
        self.yt("""
            --proxy {o.cluster} create replicated_table {o.name}
            --attr '{{
                dynamic=%true;
                schema={schema};
                replicated_table_options={{enable_replicated_table_tracker=%true}};
            }}'
        """)
        self.mount_table()


class NewReplica(YTCommand):

    def add_arguments(self, parser):
        parser.add_argument('--cluster', required=True)
        parser.add_argument('--name', required=True)
        parser.add_argument('--replica-cluster')
        parser.add_argument('--replica-name')
        parser.add_argument('--sync', action='store_true')

    def __call__(self):
        replica_cluster = self.options.replica_cluster or self.options.cluster
        replica_name = self.options.replica_name or self.options.name

        result = self.yt(
            """
            --proxy {o.cluster} create table_replica
            --attr '{{
                table_path="{o.name}";
                cluster_name="{replica_cluster}";
                replica_path="{replica_name}";
            }}'
            """,
            replica_cluster=replica_cluster,
            replica_name=replica_name,
        )
        upstream_id = result.stdout.decode('utf-8').strip()

        self.yt(
            """
            --proxy {replica_cluster} create table {replica_name}
            --attr '{{
                dynamic=%true; 
                schema={schema}; 
                upstream_replica_id="{upstream_id}";
            }}'
            """,
            replica_cluster=replica_cluster,
            replica_name=replica_name,
            upstream_id=upstream_id,
        )
        self.yt(
            '--proxy {o.cluster} alter-table-replica {upstream_id} --enable --mode {mode}',
            upstream_id=upstream_id,
            mode='sync' if self.options.sync else 'async',
        )
        self.mount_table(replica_cluster, replica_name)


commands = {
    'new-replicated': NewReplicated,
    'new-replica': NewReplica,
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('command', choices=commands)
    args = sys.argv[1:2]
    command_args = sys.argv[2:]
    options = parser.parse_args(args)
    command = commands[options.command](command_args)
    command()


if __name__ == '__main__':
    main()
