from itertools import groupby
from typing import Any, Dict

from staff.departments.tree_lib.abstract_entity_info import AbstractEntityInfo
from staff.departments.tree_lib.tree_builder import TreeBuilder


class TreeExpander(object):
    MAX_EXPANDED_ROWS_COUNT = 100

    def __init__(self, info_provider: AbstractEntityInfo, url):
        self.info_provider = info_provider
        self.url = url
        self._target = None

    @property
    def target(self):
        if self._target is None:
            self._target = self.info_provider.departments_query().get(url=self.url)

        return self._target

    def _target_is_deep(self):
        return (
            self.info_provider
            .total_entities_count_query()
            .filter(
                department__tree_id=self.target['tree_id'],
                department__lft__gte=self.target['lft'],
                department__rght__lte=self.target['rght'],
            )
            .count()
        ) > self.MAX_EXPANDED_ROWS_COUNT

    def _get_departments(self, target_is_deep):
        departments_qs = self.info_provider.departments_query()

        if target_is_deep:
            return departments_qs.filter(parent_id=self.target['id'])

        return departments_qs.filter(
            tree_id=self.target['tree_id'],
            lft__gte=self.target['lft'],
            rght__lte=self.target['rght'],
        )

    def _get_departments_entities(self, departments):
        all_ids = [d['id'] for d in departments]
        entities_qs = (
            self.info_provider
            .full_entities_query()
            .filter(department_id__in=all_ids)
            .order_by('department')
        )
        return {
            dep_id: list(entities)
            for dep_id, entities
            in groupby(self.info_provider.fill_list(entities_qs), lambda p: p['department_id'])
        }

    def expand(self) -> Dict[str, Any]:
        target_is_deep = self._target_is_deep()
        departments = self._get_departments(target_is_deep)

        builder = TreeBuilder(self.info_provider)
        if target_is_deep:
            entities_map = self._get_departments_entities([self.target])
            target = builder.get_as_list([self.target], entities_map)[0]
            target['descendants'] = builder.get_as_short_tree(departments)
        else:
            entities_map = self._get_departments_entities(departments)
            target = builder.get_as_tree(departments, entities_map)[0]

        target['is_deep'] = target_is_deep
        return target


class ProposalTreeExpander(TreeExpander):
    MAX_EXPANDED_ROWS_COUNT = 200


class ProposalVacanciesTreeExpander(ProposalTreeExpander):

    def _get_departments_entities(self, departments):
        entities_qs = (
            self.info_provider
            .full_entities_query()
            .order_by('department')
        )
        if not entities_qs:
            return {}

        return {
            self._target['id']: list(self.info_provider.fill_list(entities_qs))
        }
