# geo/management/commands/import_geo_data.py
from django.core.management.base import BaseCommand
from openpyxl import load_workbook
from geo.models import State, District, Location, LocationPart
from django.db import transaction

BATCH_SIZE = 5000  # adjust for your RAM

class Command(BaseCommand):
    help = "Import States, Districts, Locations, and LocationParts from Excel (streaming + batch safe)"

    def add_arguments(self, parser):
        parser.add_argument("filepath", type=str)
        parser.add_argument("--sheet", type=str, default=None, help="Sheet name: States, Districts, Location, LocationParts")
        parser.add_argument("--start", type=int, default=2, help="Row number to start from (default: 2)")

    def handle(self, *args, **options):
        filepath = options["filepath"]
        sheet = options["sheet"]
        start = options["start"]

        wb = load_workbook(filepath, read_only=True, data_only=True)
        self.stdout.write(self.style.NOTICE(f"Opening {filepath} (sheet={sheet or 'ALL'}) ..."))

        if sheet in (None, "States") and "States" in wb.sheetnames:
            self.import_states(wb["States"], start)

        if sheet in (None, "Districts") and "Districts" in wb.sheetnames:
            self.import_districts(wb["Districts"], start)

        if sheet in (None, "Location") and "Location" in wb.sheetnames:
            self.import_locations(wb["Location"], start)

        if sheet in (None, "LocationParts") and "LocationParts" in wb.sheetnames:
            self.import_location_parts(wb["LocationParts"], start)

        self.stdout.write(self.style.SUCCESS("✅ Import completed successfully."))

    # ----------------------------------------------------------
    # States
    # ----------------------------------------------------------
    def import_states(self, ws, start):
        self.stdout.write("Importing States ...")
        created = 0
        for i, row in enumerate(ws.iter_rows(min_row=start, values_only=True), start=start):
            if not row or not row[0]:
                continue
            state_name = str(row[0]).strip()
            _, c = State.objects.get_or_create(name=state_name)
            created += 1 if c else 0
            if created % 1000 == 0:
                self.stdout.write(f"  {created} states processed...")
        self.stdout.write(self.style.SUCCESS(f"Imported {created} States."))

    # ----------------------------------------------------------
    # Districts
    # ----------------------------------------------------------
    def import_districts(self, ws, start):
        self.stdout.write("Importing Districts ...")
        created = 0
        for i, row in enumerate(ws.iter_rows(min_row=start, values_only=True), start=start):
            if not row or not row[0]:
                continue
            # "District & State"
            parts = str(row[0]).split("&")
            if len(parts) != 2:
                continue
            district_name = parts[0].strip()
            state_name = parts[1].strip()
            state, _ = State.objects.get_or_create(name=state_name)
            _, c = District.objects.get_or_create(name=district_name, state=state)
            created += 1 if c else 0
            if created % 1000 == 0:
                self.stdout.write(f"  {created} districts processed...")
        self.stdout.write(self.style.SUCCESS(f"Imported {created} Districts."))

    # ----------------------------------------------------------
    # Locations  (1.5 L)
    # ----------------------------------------------------------
    def import_locations(self, ws, start):
        self.stdout.write("Importing Locations (streaming)...")
        batch = []
        total = 0
        for i, row in enumerate(ws.iter_rows(min_row=start, values_only=True), start=start):
            if not row or not row[0]:
                continue
            name, pincode, state_name, district_name = [str(x).strip() for x in row[:4]]
            if not (pincode and pincode.isdigit()):
                continue
            state, _ = State.objects.get_or_create(name=state_name)
            district, _ = District.objects.get_or_create(name=district_name, state=state)
            batch.append(Location(name=name, pincode=pincode, state=state, district=district))

            if len(batch) >= BATCH_SIZE:
                self._bulk_insert(Location, batch)
                total += len(batch)
                batch.clear()
                self.stdout.write(f"  Imported {total} locations so far...")

        if batch:
            self._bulk_insert(Location, batch)
            total += len(batch)

        self.stdout.write(self.style.SUCCESS(f"✅ Imported {total} Locations."))

    # ----------------------------------------------------------
    # LocationParts  (10 L)
    # ----------------------------------------------------------
    def import_location_parts(self, ws, start):
        self.stdout.write("Importing LocationParts (streaming)...")
        batch = []
        total = 0
        for i, row in enumerate(ws.iter_rows(min_row=start, values_only=True), start=start):
            if not row or not row[0]:
                continue
            nearby, pincode, state_name, district_name = [str(x).strip() for x in row[:5]]
            if not (pincode and pincode.isdigit()):
                continue
            state, _ = State.objects.get_or_create(name=state_name)
            district, _ = District.objects.get_or_create(name=district_name, state=state)
            # base = Location.objects.filter(name=base_loc, pincode=pincode).first()
            # if not base:
            #     base = Location.objects.create(name=base_loc, pincode=pincode, state=state, district=district)
            batch.append(LocationPart(
                # base_location=base,
                nearby_location=nearby,
                pincode=pincode,
                state=state,
                district=district
            ))

            if len(batch) >= BATCH_SIZE:
                self._bulk_insert(LocationPart, batch)
                total += len(batch)
                batch.clear()
                self.stdout.write(f"  Imported {total} location parts so far...")

        if batch:
            self._bulk_insert(LocationPart, batch)
            total += len(batch)

        self.stdout.write(self.style.SUCCESS(f"✅ Imported {total} LocationParts."))

    # ----------------------------------------------------------
    # Helper for safe bulk inserts
    # ----------------------------------------------------------
    @staticmethod
    def _bulk_insert(model, objs):
        try:
            with transaction.atomic():
                model.objects.bulk_create(objs, ignore_conflicts=True)
        except Exception as e:
            print("⚠️ Bulk insert error:", e)


# 🧠 How to use

# Examples:

# # Import all sheets
# python manage.py import_geo_data data.xlsx

# # Import only States
# python manage.py import_geo_data data.xlsx --sheet States

# # Resume from row 50,001 for large Location sheet
# python manage.py import_geo_data data.xlsx --sheet Location --start 50001