from multiprocessing import cpu_count, Pool

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

from apps.herd.models import Rancher
from .link_ranchers_parallel_worker import process_city


class Command(BaseCommand):
    help = "Parallel link ranchers to cooperative by city"

    def add_arguments(self, parser):
        parser.add_argument(
            '--worker',
            type=int,
            default=cpu_count() // 2
        )

    def handle(self, *args, **options):
        workers = options['workers']

        city_ids = (
            Rancher.objects.filter(city__isnull=False)
            .values_list('city_id', flat=True)
            .distinct()
        )

        self.stdout.write(
            f"Starting parallel sync for {len(city_ids)} cities "
            f"with {workers} workers"
        )

        connection.close()

        with Pool(processes=workers) as pool:
            pool.map(process_city, city_ids)

        self.stdout.write(self.style.SUCCESS("DONE ✅"))