import json
from pathlib import Path

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError
from django.db import IntegrityError, transaction

from locations.models import County, Division, Location, SubCounty, SubLocation


DEFAULT_JSON_DIRECTORY = Path(settings.BASE_DIR) / "location_json"


def clean_location_name(value):
    return str(value if value is not None else "").strip()


class Command(BaseCommand):
    help = (
        "Import Kenya administrative location JSON files from the location_json directory. "
        "The import is idempotent and can be run repeatedly."
    )

    def add_arguments(self, parser):
        parser.add_argument(
            "--directory",
            default=str(DEFAULT_JSON_DIRECTORY),
            help="Directory containing counties.json, sub_counties (2).json, divisions.json, locations.json, and sub_locations.json.",
        )
        parser.add_argument(
            "--dry-run",
            action="store_true",
            help="Validate and simulate the import without committing database changes.",
        )

    def handle(self, *args, **options):
        directory = Path(options["directory"]).expanduser().resolve()
        dry_run = bool(options["dry_run"])

        required_files = {
            "counties": directory / "counties.json",
            "sub_counties": directory / "sub_counties (2).json",
            "divisions": directory / "divisions.json",
            "locations": directory / "locations.json",
            "sub_locations": directory / "sub_locations.json",
        }

        missing = [str(path) for path in required_files.values() if not path.exists()]
        if missing:
            raise CommandError(f"Missing location JSON files: {', '.join(missing)}")

        with transaction.atomic():
            summary = {
                "counties": self._import_counties(required_files["counties"]),
                "sub_counties": self._import_sub_counties(required_files["sub_counties"]),
                "divisions": self._import_divisions(required_files["divisions"]),
                "locations": self._import_locations(required_files["locations"]),
                "sub_locations": self._import_sub_locations(required_files["sub_locations"]),
            }

            if dry_run:
                transaction.set_rollback(True)

        for label, counts in summary.items():
            self.stdout.write(
                f"{label}: created={counts['created']}, updated={counts['updated']}"
            )

        if dry_run:
            self.stdout.write(self.style.WARNING("Dry run complete. No database changes were committed."))
        else:
            self.stdout.write(self.style.SUCCESS("Kenya administrative locations imported successfully."))

    def _load_entries(self, path):
        with path.open("r", encoding="utf-8") as handle:
            payload = json.load(handle)
        if not isinstance(payload, list):
            raise CommandError(f"{path} must contain a JSON array.")
        return payload

    def _upsert(self, model, *, pk, defaults, conflict_lookup=None):
        try:
            instance, created = model.objects.update_or_create(
                pk=int(pk),
                defaults=defaults,
            )
            return instance, created
        except IntegrityError:
            if not conflict_lookup:
                raise
            instance = model.objects.filter(**conflict_lookup).first()
            if not instance:
                raise
            for field_name, value in defaults.items():
                setattr(instance, field_name, value)
            instance.save(update_fields=list(defaults.keys()))
            created = False
        return instance, created

    def _import_counties(self, path):
        counts = {"created": 0, "updated": 0}
        for entry in self._load_entries(path):
            defaults = {"name": clean_location_name(entry.get("fields", {}).get("name"))}
            if not defaults["name"]:
                raise CommandError(f"County record {entry.get('pk')} is missing a name.")
            _, created = self._upsert(
                County,
                pk=entry["pk"],
                defaults=defaults,
                conflict_lookup={"name": defaults["name"]},
            )
            counts["created" if created else "updated"] += 1
        return counts

    def _import_sub_counties(self, path):
        counts = {"created": 0, "updated": 0}
        for entry in self._load_entries(path):
            fields = entry.get("fields", {})
            county = County.objects.filter(pk=fields.get("county")).first()
            if not county:
                raise CommandError(
                    f"Sub-county record {entry.get('pk')} references missing county {fields.get('county')}."
                )
            defaults = {
                "name": clean_location_name(fields.get("name")),
                "county": county,
            }
            if not defaults["name"]:
                raise CommandError(f"Sub-county record {entry.get('pk')} is missing a name.")
            _, created = self._upsert(
                SubCounty,
                pk=entry["pk"],
                defaults=defaults,
                conflict_lookup={"county": county, "name": defaults["name"]},
            )
            counts["created" if created else "updated"] += 1
        return counts

    def _import_divisions(self, path):
        counts = {"created": 0, "updated": 0}
        for entry in self._load_entries(path):
            fields = entry.get("fields", {})
            sub_county = SubCounty.objects.filter(pk=fields.get("constituency")).first()
            if not sub_county:
                raise CommandError(
                    f"Division record {entry.get('pk')} references missing sub-county {fields.get('constituency')}."
                )
            defaults = {
                "name": clean_location_name(fields.get("name")),
                "sub_county": sub_county,
            }
            if not defaults["name"]:
                raise CommandError(f"Division record {entry.get('pk')} is missing a name.")
            _, created = self._upsert(
                Division,
                pk=entry["pk"],
                defaults=defaults,
                conflict_lookup={"sub_county": sub_county, "name": defaults["name"]},
            )
            counts["created" if created else "updated"] += 1
        return counts

    def _import_locations(self, path):
        counts = {"created": 0, "updated": 0}
        for entry in self._load_entries(path):
            fields = entry.get("fields", {})
            division = Division.objects.filter(pk=fields.get("division")).first()
            if not division:
                raise CommandError(
                    f"Location record {entry.get('pk')} references missing division {fields.get('division')}."
                )
            defaults = {
                "name": clean_location_name(fields.get("name")),
                "division": division,
            }
            if not defaults["name"]:
                raise CommandError(f"Location record {entry.get('pk')} is missing a name.")
            _, created = self._upsert(
                Location,
                pk=entry["pk"],
                defaults=defaults,
                conflict_lookup={"division": division, "name": defaults["name"]},
            )
            counts["created" if created else "updated"] += 1
        return counts

    def _import_sub_locations(self, path):
        counts = {"created": 0, "updated": 0}
        for entry in self._load_entries(path):
            fields = entry.get("fields", {})
            location = Location.objects.filter(pk=fields.get("location")).first()
            if not location:
                raise CommandError(
                    f"Sub-location record {entry.get('pk')} references missing location {fields.get('location')}."
                )
            defaults = {
                "name": clean_location_name(fields.get("name")),
                "location": location,
            }
            if not defaults["name"]:
                raise CommandError(f"Sub-location record {entry.get('pk')} is missing a name.")
            _, created = self._upsert(
                SubLocation,
                pk=entry["pk"],
                defaults=defaults,
                conflict_lookup={"location": location, "name": defaults["name"]},
            )
            counts["created" if created else "updated"] += 1
        return counts