from collections import defaultdict

from django.db.models import Count

from apps.authentication.models import Organization
from apps.authentication.services.service import get_all_org_child
from apps.herd.models import RancherOrganizationLink, Herd
from apps.livestock.models import LiveStock


class RancherOrganizationService:

    def orgs_linked_rancher(self, org: Organization, org_type_key: str):

        if org.type.key != 'ADM':
            org_ids = list(
                o.id for o in get_all_org_child(org)
            )
        else:
            org_ids = list(
                Organization.objects.filter(
                    type__key=org_type_key
                ).values_list('id', flat=True)
            )

        organizations = (
            Organization.objects
            .filter(id__in=org_ids, type__key=org_type_key)
            .select_related('province', 'city')
            .prefetch_related('service_area')
        )

        rancher_counts = {
            r['organization_id']: r['cnt']
            for r in (
                RancherOrganizationLink.objects
                .filter(organization_id__in=org_ids)
                .values('organization_id')
                .annotate(cnt=Count('rancher_id'))
            )
        }

        herd_counts = {
            r['organization_id']: r['cnt']
            for r in (
                RancherOrganizationLink.objects
                .filter(organization_id__in=org_ids)
                .values('organization_id')
                .annotate(cnt=Count('rancher__herd'))
            )
        }

        livestock_per_herd = {
            row['herd_id']: row['cnt']
            for row in (
                LiveStock.objects
                .values('herd_id')
                .annotate(cnt=Count('id'))
            )
        }

        livestock_counts = defaultdict(int)

        herd_org_map = (
            Herd.objects
            .filter(
                rancher__organization_links__organization_id__in=org_ids
            )
            .values(
                'id',
                'rancher__organization_links__organization_id'
            )
        )

        for row in herd_org_map:
            livestock_counts[
                row['rancher__organization_links__organization_id']
            ] += livestock_per_herd.get(row['id'], 0)

        return [
            {
                "id": org.id,
                "name": org.name,
                "org_service_area": [
                    {"id": city.id, "name": city.name}
                    for city in org.service_area.all()
                ],
                "org_purchase_policy": org.purchase_policy,
                "province": org.province.name,
                "province_id": org.province.id,
                "city": org.city.name,
                "city_id": org.city.id,

                "rancher_count": rancher_counts.get(org.id, 0),
                "herd_count": herd_counts.get(org.id, 0),
                "livestock_count": livestock_counts.get(org.id, 0),
            }
            for org in organizations
        ]