#python manage.py import_necta path/to/results.csv --exam PSLE --year 2023

import csv
from django.core.management.base import BaseCommand
from exams.models import (
    Student, School, Region, District, Ward,
    ExamType, ExamSitting, ExamSubject, ExamSubjectMark
)
from django.db import transaction

GRADE_MAP = {
    "A": 5, "B": 4, "C": 3, "D": 2, "E": 1, "F": 0
}

SUBJECT_CACHE = {}
SCHOOL_CACHE = {}
STUDENT_CACHE = {}
EXAMTYPE_CACHE = {}


def get_or_create_subject(code, eng_name="Subject"):
    """Cached subject fetch (VERY fast for millions of rows)."""
    if code in SUBJECT_CACHE:
        return SUBJECT_CACHE[code]

    subject, created = ExamSubject.objects.get_or_create(
        subject_code=code,
        defaults={"subject_name_eng": eng_name, "subject_name_sw": eng_name}
    )
    SUBJECT_CACHE[code] = subject
    return subject


def get_or_create_examtype(code):
    if code in EXAMTYPE_CACHE:
        return EXAMTYPE_CACHE[code]

    exam_type, created = ExamType.objects.get_or_create(
        code=code,
        defaults={"name": code}
    )
    EXAMTYPE_CACHE[code] = exam_type
    return exam_type


def get_or_create_school(code, name="Unknown School"):
    if code in SCHOOL_CACHE:
        return SCHOOL_CACHE[code]

    school, created = School.objects.get_or_create(
        school_code=code,
        defaults={
            "name": name,
            "region": None,
            "district": None,
            "ward": None
        }
    )
    SCHOOL_CACHE[code] = school
    return school


def get_or_create_student(cand_no, prem_no, sex, name, school):
    if cand_no in STUDENT_CACHE:
        return STUDENT_CACHE[cand_no]

    student, created = Student.objects.get_or_create(
        cand_no=cand_no,
        defaults={
            "prem_no": prem_no,
            "sex": sex,
            "name": name,
            "school": school
        }
    )
    STUDENT_CACHE[cand_no] = student
    return student


class Command(BaseCommand):
    help = "Import NECTA-style CSV results into normalized tables"

    def add_arguments(self, parser):
        parser.add_argument("csv_file", type=str, help="Full path to CSV file")
        parser.add_argument("--exam", type=str, default="PSLE", help="Exam type code: PSLE, CSEE, FTNA")
        parser.add_argument("--year", type=int, required=True, help="Exam year")

    @transaction.atomic
    def handle(self, *args, **options):
        csv_path = options["csv_file"]
        exam_code = options["exam"]
        year = options["year"]

        exam_type = get_or_create_examtype(exam_code)

        self.stdout.write(self.style.WARNING(f"Importing NECTA results from {csv_path}"))
        self.stdout.write(self.style.NOTICE("This may take time for large datasets..."))

        sitting_bulk = []
        subjectmarks_bulk = []

        with open(csv_path, newline='', encoding="utf-8") as file:
            reader = csv.DictReader(file)

            for row in reader:
                # Required fields from CSV
                school_code = row["SCHOOL_COD"].strip()
                cand_no = row["CAND_NO"].strip()
                prem_no = row["PREM_NO"].strip()
                sex = row["SEX"].strip()
                name = row["NAMES"].strip()

                avg_grade = row.get("AVERAGE_GRADE", "").strip()
                subjects_raw = row.get("RESULTS", "")

                # Create school
                school = get_or_create_school(school_code)

                # Create student
                student = get_or_create_student(cand_no, prem_no, sex, name, school)

                # Create exam sitting
                sitting = ExamSitting(
                    student=student,
                    school=school,
                    exam_type=exam_type,
                    year=year,
                    avg_grade=avg_grade,
                    avg_numeric=GRADE_MAP.get(avg_grade, 0)
                )
                sitting_bulk.append(sitting)

                # Process each subject-grade pair
                if subjects_raw:
                    subjects_list = subjects_raw.split(",")
                    for item in subjects_list:
                        try:
                            subject_name, grade = item.strip().split(" - ")
                            grade = grade.strip()
                        except ValueError:
                            continue

                        # subject code example: "Kiswahili" = 11
                        subject_code = subject_name[:3].upper()

                        subject = get_or_create_subject(subject_code, subject_name)

                        subjectmarks_bulk.append(
                            (cand_no, subject, grade)
                        )

        # SAVE EXAM SITTINGS
        created_sittings = ExamSitting.objects.bulk_create(
            sitting_bulk,
            batch_size=2000
        )
        self.stdout.write(self.style.SUCCESS(f"Created {len(created_sittings)} exam sitting records."))

        # MAP sitting by cand_no for fast lookup
        sitting_map = {s.student.cand_no: s for s in ExamSitting.objects.filter(year=year)}

        # SAVE SUBJECT MARKS
        bulk_marks = []
        for cand_no, subject, grade in subjectmarks_bulk:
            bulk_marks.append(
                ExamSubjectMark(
                    sitting=sitting_map[cand_no],
                    subject=subject,
                    grade=grade,
                    numeric_score=GRADE_MAP.get(grade, 0)
                )
            )

        ExamSubjectMark.objects.bulk_create(bulk_marks, batch_size=5000)

        self.stdout.write(self.style.SUCCESS("Import Completed Successfully!"))