import os
from multiprocessing import cpu_count, Pool

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

from .link_ranchers_parallel_worker import process_city

os.environ["DJANGO_ALLOW_ASYNC_UNSAFE"] = "true"


def run_worker(chunk):
    import os
    import django

    os.environ.setdefault(
        "DJANGO_SETTINGS_MODULE",
        "Rasaddam_Backend.settings"
    )
    django.setup()


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):
        from apps.herd.models import Rancher
        workers = options['worker']

        from multiprocessing import get_context

        chunks = self.prepare_chunks()

        ctx = get_context("spawn")
        with ctx.Pool(processes=options["workers"]) as pool:
            pool.map(run_worker, chunks)

        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 ✅"))