import pandas as pd
from django.db import transaction
from exams.models import Region, District, Ward, School

 
def import_geo_dataframe(df):
    """
    Required columns:
    REGION, COUNCIL, WARD, SCHOOL NAME, SCHOOL_CODE, OWNERSHIP

    Optional columns:
    SCHOOL_REGNO, LATITUDE, LONGITUDE
    """

    required_cols = [
        "REGION", "COUNCIL", "WARD",
        "SCHOOL NAME", "SCHOOL_CODE", "OWNERSHIP"
    ]

    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"Missing required column: {col}")

    created = {"regions": 0, "districts": 0, "wards": 0, "schools": 0}
    updated = {"schools": 0}

    with transaction.atomic():
        for _, row in df.iterrows():

            region, r_new = Region.objects.get_or_create(
                name=str(row["REGION"]).strip().upper()
            )
            created["regions"] += int(r_new)

            district, d_new = District.objects.get_or_create(
                region=region,
                name=str(row["COUNCIL"]).strip().upper()
            )
            created["districts"] += int(d_new)

            ward, w_new = Ward.objects.get_or_create(
                district=district,
                name=str(row["WARD"]).strip().upper()
            )
            created["wards"] += int(w_new)

            # Handle school registration number carefully
            school_regno = None
            if "SCHOOL_REGNO" in df.columns and pd.notna(row["SCHOOL_REGNO"]):
                school_regno_val = str(row["SCHOOL_REGNO"]).strip()
                if school_regno_val:  # Only set if not empty
                    school_regno = school_regno_val

            # Build defaults dictionary
            defaults = {
                "name": str(row["SCHOOL NAME"]).strip(),
                "region": region,
                "district": district,
                "ward": ward,
                "ownership": str(row["OWNERSHIP"]).strip(),
                "latitude": (
                    float(row["LATITUDE"])
                    if "LATITUDE" in df.columns and pd.notna(row["LATITUDE"])
                    else None
                ),
                "longitude": (
                    float(row["LONGITUDE"])
                    if "LONGITUDE" in df.columns and pd.notna(row["LONGITUDE"])
                    else None
                ),
            }
            
            # Only add school_regno to defaults if we have a value
            if school_regno is not None:
                defaults["school_regno"] = school_regno

            # Use update_or_create with school_code as lookup
            school, s_created = School.objects.update_or_create(
                school_code=str(row["SCHOOL_CODE"]).strip(),
                defaults=defaults
            )
            
            if s_created:
                created["schools"] += 1
            else:
                updated["schools"] += 1

    return created, updated