from collections import defaultdict

from django.core.management.base import BaseCommand
from django.db import transaction
from django.db.models.aggregates import Count

from apps.herd.models import Herd
from apps.livestock.models import LiveStock


class Command(BaseCommand):
    help = "Rebuild heavy/light counts for all herds"
    BATCH_SIZE = 2000

    def handle(self, *args, **options):
        self.stdout.write("🔄 Rebuilding herd livestock counts...")

        self.stdout.write("📊 Aggregating livestock data...")
        livestock_stats = (
            LiveStock.objects
            .filter(archive=False, herd__isnull=False)
            .values('herd_id', 'weight_type')
            .annotate(cnt=Count('id'))
        )

        herd_map = defaultdict(lambda: {"H": 0, "L": 0})

        for row in livestock_stats:
            herd_map[row['herd_id']][row['weight_type']] = row['cnt']

        self.stdout.write(f"✅ Aggregated {len(herd_map)} herds")

        self.stdout.write("✍️ Updating herds in batches...")

        herd_ids = list(herd_map.keys())

        for i in range(0, len(herd_ids), self.BATCH_SIZE):
            batch_ids = herd_ids[i:i + self.BATCH_SIZE]

            herds = Herd.objects.filter(id__in=batch_ids)

            for herd in herds:
                counts = herd_map.get(herd.id, {})
                herd.heavy_livestock_number = counts.get("H", 0)
                herd.light_livestock_number = counts.get("L", 0)

            with transaction.atomic():
                Herd.objects.bulk_update(
                    herds,
                    ['heavy_livestock_number', 'light_livestock_number']
                )

            self.stdout.write(
                f"✔ Updated {i + len(batch_ids)} / {len(herd_ids)} herds"
            )

            self.stdout.write("🎉 Herd livestock counts rebuilt successfully.")