from itertools import groupby, chain
import logging


logger = logging.getLogger(__name__)


def get_heads_for_persons(persons_ids, all_heads=False):
    from fb.staff.models import GroupMembership, Group, Person
    groups = (
        GroupMembership.objects
        .filter(
            person__in=persons_ids,
            is_primary=True,
        )
        .order_by('person_id', '-is_chief', '-is_deputy')
    )
    find_groups_head = set()
    user_membership = {}

    for p_id, memberships in groupby(groups, key=lambda x: x.person_id):
        for membership in memberships:
            user_membership[p_id] = membership
            find_groups_head.add(membership.group_id)
            break
            # If group is primary group for person, find head of group.
            # Otherwise use random group.

    group_paths = dict(
        Group.objects
        .filter(staff_id__in=find_groups_head)
        .values_list('staff_id', 'path')
    )
    # Paths of groups

    ancestor_paths_by_staff_id = {
        staff_id: list(Group._get_ancestors_pathes(path))
        for staff_id, path in group_paths.iteritems()
    }

    ancestor_paths = set(chain(*ancestor_paths_by_staff_id.itervalues()))

    heads = (
        GroupMembership.objects
        .filter(
            is_chief=True,
            group__group_structure__path__in=ancestor_paths,
        )
        .order_by(
            'group__group_structure__path',
            '-is_primary',
        )
        .values('group__group_structure__path', 'person_id')
    )

    get_path = lambda x: x['group__group_structure__path']

    heads_by_path = {}
    for p, _groups in groupby(heads, key=get_path):
        for g in _groups:
            heads_by_path[p] = g['person_id']
            break

    heads_by_person_id = {}
    for p_id, membership in user_membership.iteritems():
        _group_paths = ancestor_paths_by_staff_id.get(membership.group_id, None)
        if _group_paths is None:
            continue
        _group_paths = reversed(_group_paths)

        if membership.is_chief:
            # skip own group
            next(_group_paths)

        _heads = [
            heads_by_path[p] for p in _group_paths if p in heads_by_path and heads_by_path[p] != p_id
        ]

        heads_by_person_id[p_id] = None
        if p_id == Person.volozh_staff_id:
            _heads = [Person.bunina_staff_id]
        if all_heads:
            heads_by_person_id[p_id] = _heads
        elif _heads:
            heads_by_person_id[p_id] = _heads[0]

    return heads_by_person_id


def get_heads_for_groups(groups_ids, all_heads=False):
    from fb.staff.models import GroupMembership, Group

    group_paths = dict(
        Group.objects
        .filter(staff_id__in=groups_ids)
        .values_list('staff_id', 'path')
    )
    # Paths of groups

    ancestor_paths_by_staff_id = {
        staff_id: list(Group._get_ancestors_pathes(path))
        for staff_id, path in group_paths.iteritems()
    }

    ancestor_paths = set(chain(*ancestor_paths_by_staff_id.itervalues()))

    heads = (
        GroupMembership.objects
        .filter(
            is_chief=True,
            group__group_structure__path__in=ancestor_paths,
        )
        .order_by(
            'group__group_structure__path',
            '-is_primary',
        )
        .values('group__group_structure__path', 'person_id')
    )

    get_path = lambda x: x['group__group_structure__path']

    heads_by_path = {}
    for p, _groups in groupby(heads, key=get_path):
        for g in _groups:
            heads_by_path[p] = g['person_id']
            break

    heads_by_group_id = {}
    for g_id in groups_ids:
        _group_paths = ancestor_paths_by_staff_id.get(g_id, [])
        _group_paths = reversed(_group_paths)

        _heads = [
            heads_by_path[p] for p in _group_paths if p in heads_by_path
        ]

        heads_by_group_id[g_id] = None
        if all_heads:
            heads_by_group_id[g_id] = _heads
        elif _heads:
            heads_by_group_id[g_id] = _heads[0]

    return heads_by_group_id
