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:
#     """
#     different services of ranchers linked to organization
#     """
#
#     def orgs_linked_rancher(self, org: Organization = None, org_type_key: str = None):
#         """
#         list of organizations with their information of rancher, herd, ....
#         """
#         if org.type.key != 'ADM':
#             organizations = get_all_org_child(org)
#             organizations = Organization.objects.filter(id__in=[item.id for item in organizations])
#         else:
#             organizations = Organization.objects.filter(type__key=org_type_key)
#
#         linked_qs = RancherOrganizationLink.objects.select_related(
#             'rancher',
#             'organization'
#         ).filter(organization__in=organizations, organization__type__key=org_type_key)
#
#         organizations = organizations.annotate(
#             rancher_count=Count(
#                 'rancher_links__rancher_id',
#                 distinct=True
#             ),
#             herd_count=Count(
#                 'rancher_links__rancher__herd',
#                 distinct=True
#             ),
#             livestock_count=Count(
#                 'rancher_links__rancher__herd__live_stock_herd',
#                 distinct=True
#             ),
#         )
#
#         return [
#             {
#                 "id": org.id,
#                 "name": org.name,
#                 "org_service_area": [{'name': city.name, 'id': city.id} 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": org.rancher_count,
#                 "herd_count": org.herd_count,
#                 "livestock_count": org.livestock_count,
#             }
#             for org in organizations
#         ]


class RancherOrganizationService:

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

        # --------------------------------------------------
        # 1. Resolve organization ids
        # --------------------------------------------------
        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')
        )

        # --------------------------------------------------
        # 2. Rancher count per organization
        # --------------------------------------------------
        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'))
            )
        }

        # --------------------------------------------------
        # 3. Herd count per organization
        # --------------------------------------------------
        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'))
            )
        }

        # --------------------------------------------------
        # 4. Livestock count (optimized path)
        #   LiveStockHerd -> Herd -> Org
        # --------------------------------------------------

        # 4-1 livestock count per herd (NO JOIN, NO DISTINCT)
        livestock_per_herd = {
            row['herd_id']: row['cnt']
            for row in (
                LiveStock.objects
                .values('herd_id')
                .annotate(cnt=Count('id'))
            )
        }

        # 4-2 sum livestock per organization
        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)

        # --------------------------------------------------
        # 5. Final response
        # --------------------------------------------------
        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
        ]