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}

    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)

            school, s_new = School.objects.update_or_create(
                school_code=str(row["SCHOOL_CODE"]).strip(),
                defaults={
                    "school_regno": (
                        str(row["SCHOOL_REGNO"]).strip()
                        if "SCHOOL_REGNO" in df.columns and pd.notna(row["SCHOOL_REGNO"])
                        else None
                    ),
                    "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
                    ),
                }
            )
            created["schools"] += int(s_new)

    return created