from django.core.management.base import BaseCommand
from yt.wrapper import TablePath
from logging import getLogger

logger = getLogger(__name__)


class Mapper:
    def __init__(self, columns):
        self.columns = columns

    def __call__(self, row):
        if not row['deleted']:
            yield {key: row[key] for key in self.columns}


class Command(BaseCommand):
    help = 'Migrate from yt table for ferryman tables to yt cache tables'
    yt_table_schema = [
        {'name': 'url', 'type': 'string'},
        {'name': 'timestamp', 'type': 'int64'},
        {'name': 'raw', 'type': 'any'},
    ]
    columns = ['url', 'timestamp', 'raw']
    table_prefix = '//home/intrasearch/data'
    tables = 'at-,at-clubs,at-posts,doc-,doc-external,st-,wiki-'

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        from intranet.search.core.yt_client import client, convert_static_to_dynamic_table
        self.client = client
        self.convert_static_to_dynamic_table = convert_static_to_dynamic_table

    def add_arguments(self, parser):
        from django.conf import settings
        parser.add_argument('--env', action='store',
                            choices=['development', 'testing', 'production'],
                            default=settings.YENV_TYPE,
                            help=("Type of environment. Choices: "
                                  "development, testing, production"))
        parser.add_argument('--tables', default=self.tables,
                            help='Migration tables')

    def migrate(self, table_from, table_to):
        if not self.client.exists(table_from):
            logger.warning(f'Table {table_from} does not exists')
            return
        if not self.client.exists(table_to):
            self.client.create('table', table_to, recursive=True,
                               attributes={'schema': self.yt_table_schema})
        elif self.client.get_attribute(table_to, 'dynamic'):
            logger.info(f'Dynamic table {table_to} already exists')
            return
        table_to = TablePath(table_to, append=True)
        mapper = Mapper(self.columns)
        self.run_map(mapper, table_from, table_to, self.client)
        self.convert_static_to_dynamic_table(
            table_to,
            self.yt_table_schema,
            [self.yt_table_schema[0]['name']]
        )

    def run_map(self, mapper, table_from, table_to, client):
        client.run_map(mapper, table_from, table_to)

    def handle(self, **options):
        tables = options['tables'].split(',')
        for table in tables:
            table_from = '/'.join([self.table_prefix, options['env'], table])
            table_to = '/'.join([self.table_prefix, options['env'], 'cache-' + table])
            self.migrate(table_from, table_to)
