"""Apply Zambia province and district overrides to geo JSON datasets."""
import json
from pathlib import Path

from django.core.management.base import BaseCommand

from home.geo.zambia_geo import ZAMBIA_REGION_CODES, zambia_districts, zambia_region_entries


class Command(BaseCommand):
    help = 'Update Zambia provinces and districts in geo JSON files'

    def handle(self, *args, **options):
        data_dir = Path(__file__).resolve().parents[2] / 'geo' / 'data'
        regions_path = data_dir / 'regions_by_country.json'
        cities_path = data_dir / 'cities_by_region.json'

        regions_by_country = json.loads(regions_path.read_text(encoding='utf-8'))
        cities_by_region = json.loads(cities_path.read_text(encoding='utf-8'))

        districts = zambia_districts()
        regions_by_country['ZM'] = zambia_region_entries()

        zm_city_keys = [key for key in cities_by_region if key.startswith('ZM|')]
        for key in zm_city_keys:
            del cities_by_region[key]

        for province_name, district_list in districts.items():
            code = ZAMBIA_REGION_CODES[province_name]
            sorted_districts = sorted(district_list, key=str.casefold)
            cities_by_region[f'ZM|{code}'] = sorted_districts
            cities_by_region[f'ZM|{province_name}'] = sorted_districts

        regions_path.write_text(
            json.dumps(regions_by_country, ensure_ascii=False, separators=(',', ':')),
            encoding='utf-8',
        )
        cities_path.write_text(
            json.dumps(cities_by_region, ensure_ascii=False, separators=(',', ':')),
            encoding='utf-8',
        )

        total = sum(len(v) for v in districts.values())
        self.stdout.write(
            self.style.SUCCESS(
                f'Updated Zambia geo data: {len(districts)} provinces, {total} districts.'
            )
        )