from collections import defaultdict

from django.core.management.base import BaseCommand
from django.conf import settings

from intranet.search.core.storages import FerrymanTableStorage
from intranet.search.core.storages.job_storage import job_status
from intranet.search.yt.jobs.merge_doc_tables import Job


class Command(BaseCommand):
    help = 'Команда запуска мержа запушенных в ferryman таблиц'

    def add_arguments(self, parser):
        parser.add_argument('--tables', '-t', action='store', default='',
                            help='comma separated list of tables to merge')

    @job_status
    def handle(self, **options):
        self.storage = FerrymanTableStorage()

        to_merge = defaultdict(list)
        for table in self.get_tables(**options):
            info = self.storage.parse_table_path(table)
            destination = self.storage.get_yt_table_path(search=info['search'], index=info['index'])
            to_merge[destination].append(table)

        self.run_merge(to_merge)

        for tables in to_merge.values():
            self.storage.delete(path__in=tables)

    def get_tables(self, **options):
        if options['tables']:
            return [self.storage.get_absolute_path(t) for t in options['tables'].split(',')]
        return [t['path'] for t in self.storage.get_finished()]

    def run_merge(self, tables):
        merge_job = Job({'env': settings.YENV_TYPE})
        merge_job.run(tables)
