import time

from django.core.management.base import BaseCommand
from django.db import transaction

from apps.authentication.models import Organization
from apps.herd.models import Rancher, RancherOrganizationLink


class Command(BaseCommand):
    help = "Link ranchers to cooperative of their city"

    BATCH_SIZE = 2000

    def handle(self, *args, **options):

        qs = (
            Rancher.objects
            .filter(city__isnull=False)
            .select_related('city')
        )

        total = qs.count()
        processed = 0
        created = 0
        skipped = 0
        ambiguous = 0

        start = time.time()

        buffer = []

        self.stdout.write(
            f"Starting parallel sync for {total} ranchers...."
        )

        for rancher in qs.iterator(chunk_size=self.BATCH_SIZE):

            processed += 1

            if RancherOrganizationLink.objects.filter(rancher=rancher).exists():
                skipped += 1
                continue

            orgs = list(
                Organization.objects.filter(
                    city=rancher.city,
                    type__key='CO'
                )
            )

            if not orgs:
                skipped += 1
                continue

            for org in orgs:
                buffer.append(
                    RancherOrganizationLink(
                        rancher=rancher,
                        organization=org
                    )
                )

            if len(orgs) > 1:
                ambiguous += 1

            print(len(buffer))
            if len(buffer) >= self.BATCH_SIZE:
                created += self.bulk_create(buffer)
                buffer.clear()

            # progress log
            if processed % 10_000 == 0:
                elapsed = time.time() - start
                speed = processed / elapsed if elapsed else 0

                self.stdout.write(
                    self.style.NOTICE(
                        f"[{processed:,}/{total:,}] "
                        f"created={created:,} "
                        f"skipped={skipped:,} "
                        f"ambiguous={ambiguous:,} "
                        f"{speed:.0f} r/s"
                    )
                )

        if buffer:
            created += self.bulk_create(buffer)

        self.stdout.write(
            self.style.SUCCESS(
                f"DONE ✅ created={created}, skipped={skipped}, ambiguous={ambiguous}"
            )
        )

    @transaction.atomic
    def bulk_create(self, objs):
        RancherOrganizationLink.objects.bulk_create(
            objs,
            ignore_conflicts=True
        )
        return len(objs)
