from .dates_add_days import parse_iso_date
from .education_school_year_group_data import COHORT_RULES, YEAR_GROUPS
from .education_school_year_group_types import SchoolYearGroup, UkNation


def school_year_group(birth_date: str, on_date: str, nation: UkNation) -> SchoolYearGroup:
    """The year group a child is in on a date, placed by their birth date.

    The cohort comes from the birth date itself (born on or after the cohort
    start joins a year later), never from an age on a cut-off date, which
    misplaces 29 February births where the cohort starts on 1 March.
    """
    birth = parse_iso_date(birth_date)
    on = parse_iso_date(on_date)
    if on_date < birth_date:
        raise ValueError("onDate %s is before birthDate %s" % (on_date, birth_date))
    rule = next((r for r in COHORT_RULES if r.nation == nation), None)
    if rule is None:
        raise ValueError('unknown nation "%s": use england, wales, scotland or northern-ireland' % (nation,))
    entry_year = birth.year + rule.entry_age + (1 if birth_date[5:] >= rule.cohort_start else 0)
    start_year = on.year if on_date[5:] >= rule.school_year_start else on.year - 1
    index = start_year - entry_year
    name = next((g for g in YEAR_GROUPS if g.nation == nation and g.index == index), None)
    return SchoolYearGroup(
        school_year_start="%04d-%s" % (start_year, rule.school_year_start),
        entry_year=entry_year,
        index=index,
        year_group=None if name is None else name.label,
        stage=None if name is None else name.stage,
    )