"""EES Singer tap — extracts KS2, KS4, Census, Admissions data.

Each stream targets a specific CSV file within an EES release ZIP.
The EES data uses 'school_urn' for school-level records and 'z' for
suppressed values. Column names vary by file — schemas declare all
columns needed by downstream dbt staging models.

Phonics has no school-level data on EES and is not included.
"""

from __future__ import annotations

import io
import zipfile

import requests
from singer_sdk import Stream, Tap
from singer_sdk import typing as th

CONTENT_API_BASE = (
    "https://content.explore-education-statistics.service.gov.uk/api"
)
TIMEOUT = 120


def get_content_release_id(publication_slug: str) -> str:
    """Return the latest release ID via the EES content API."""
    url = f"{CONTENT_API_BASE}/publications/{publication_slug}/releases/latest"
    resp = requests.get(url, timeout=TIMEOUT)
    resp.raise_for_status()
    return resp.json()["id"]


def _slug_to_time_period(slug: str) -> str | None:
    """Convert a release slug like '2022-23' to a time_period like '202223'."""
    parts = slug.split("-")
    if len(parts) == 2 and len(parts[0]) == 4 and len(parts[1]) == 2:
        return parts[0] + parts[1]
    return None


def get_all_releases(publication_slug: str) -> list[dict]:
    """Return all releases for a publication as dicts with 'id' and 'time_period'.

    The EES content API paginates with a 'paging' envelope when there are many
    releases. This function follows all pages so no historical release is missed.
    """
    result = []
    page = 1
    while True:
        url = f"{CONTENT_API_BASE}/publications/{publication_slug}/releases?page={page}&pageSize=20"
        resp = requests.get(url, timeout=TIMEOUT)
        resp.raise_for_status()
        data = resp.json()

        # API returns either a plain list or a paginated object with a "results" key
        if isinstance(data, list):
            releases = data
            total_pages = 1
        else:
            releases = data.get("results", [])
            paging = data.get("paging", {})
            total_pages = paging.get("totalPages", 1)

        for r in releases:
            time_period = _slug_to_time_period(r.get("slug", ""))
            result.append({"id": r["id"], "time_period": time_period})

        if page >= total_pages:
            break
        page += 1

    return result


def download_release_zip(release_id: str) -> zipfile.ZipFile:
    """Download all data files for a release as a ZIP."""
    url = f"{CONTENT_API_BASE}/releases/{release_id}/files"
    resp = requests.get(url, timeout=300, stream=True)
    resp.raise_for_status()
    return zipfile.ZipFile(io.BytesIO(resp.content))


class EESDatasetStream(Stream):
    """Base stream for an EES dataset extracted from a release ZIP.

    Subclasses set _target_filename to a keyword that appears in the
    target CSV path inside the ZIP (substring match, not exact).
    Subclasses may set _column_renames to map messy CSV column names to
    clean Singer field names before yielding records.
    """

    replication_key = None
    _publication_slug: str = ""
    _target_filename: str = ""  # keyword that appears in the CSV path
    _urn_column: str = "school_urn"  # column name for URN in the CSV
    _encoding: str = "utf-8"  # CSV file encoding (some DfE files use latin-1)
    _column_renames: dict = {}  # CSV column name → Singer field name

    def get_records(self, context):
        import pandas as pd

        latest_only = self.config.get("latest_only", False)
        if latest_only:
            release_id = get_content_release_id(self._publication_slug)
            releases = [{"id": release_id, "time_period": None}]
        else:
            releases = get_all_releases(self._publication_slug)

        self.logger.info(
            "Found %d release(s) for %s", len(releases), self._publication_slug
        )

        for release in releases:
            release_id = release["id"]
            release_time_period = release["time_period"]
            self.logger.info(
                "Downloading release %s for %s", release_id, self._publication_slug
            )
            zf = download_release_zip(release_id)

            # Find the target file (substring match)
            all_files = zf.namelist()
            target = None
            for name in all_files:
                if self._target_filename in name and name.endswith(".csv"):
                    target = name
                    break

            if not target:
                self.logger.warning(
                    "File matching '%s' not found in release %s. Available: %s",
                    self._target_filename,
                    release_id,
                    [n for n in all_files if n.endswith(".csv")],
                )
                continue

            self.logger.info("Reading %s from ZIP", target)
            with zf.open(target) as f:
                df = pd.read_csv(f, dtype=str, keep_default_na=False, encoding=self._encoding)

            # Strip BOM from first column name — handles both:
            # - UTF-8 BOM decoded as Unicode (\ufeff) when read with utf-8/utf-8-sig
            # - UTF-8 BOM bytes decoded as Latin-1 (ï»¿) when read with latin-1
            cols = list(df.columns)
            if cols:
                cols[0] = cols[0].lstrip("\ufeff").lstrip("ï»¿")
                df.columns = cols

            # Normalise URN column name — older files use 'URN' instead of 'urn'
            urn_col = self._urn_column
            col_lower = {c.lower(): c for c in df.columns}
            actual_urn = col_lower.get(urn_col.lower())
            if actual_urn and actual_urn != urn_col:
                df.rename(columns={actual_urn: urn_col}, inplace=True)

            # Filter to school-level data if the column exists
            if "geographic_level" in df.columns:
                df = df[df["geographic_level"] == "School"]

            # Drop rows with no URN (LA/category aggregates that slip through the level filter)
            if urn_col in df.columns:
                df = df[df[urn_col].notna() & (df[urn_col] != "")]

            self.logger.info("Emitting %d school-level rows from release %s", len(df), release_id)

            for _, row in df.iterrows():
                record = row.to_dict()
                # Apply subclass column renames (messy CSV names → clean Singer fields)
                for old, new in self._column_renames.items():
                    if old in record:
                        record[new] = record.pop(old)
                # Normalise URN column to 'school_urn' for consistency
                if self._urn_column in record and self._urn_column != "school_urn":
                    record["school_urn"] = record.pop(self._urn_column)
                # Inject time_period from release metadata if absent (older files
                # don't include this column — it's only in the filename/release slug)
                if "time_period" not in record and release_time_period:
                    record["time_period"] = release_time_period
                yield record


# ── KS2 Attainment (long format: one row per school × subject × breakdown) ──

class EESKS2AttainmentStream(EESDatasetStream):
    name = "ees_ks2_attainment"
    primary_keys = ["school_urn", "time_period", "subject", "breakdown_topic", "breakdown"]
    _publication_slug = "key-stage-2-attainment"
    _target_filename = "ks2_school_attainment_data"
    schema = th.PropertiesList(
        th.Property("time_period", th.StringType, required=True),
        th.Property("school_urn", th.StringType, required=True),
        th.Property("school_laestab", th.StringType),
        th.Property("school_name", th.StringType),
        th.Property("breakdown_topic", th.StringType, required=True),
        th.Property("breakdown", th.StringType, required=True),
        th.Property("subject", th.StringType, required=True),
        th.Property("expected_standard_pupil_percent", th.StringType),
        th.Property("higher_standard_pupil_percent", th.StringType),
        th.Property("average_scaled_score", th.StringType),
        th.Property("progress_measure_score", th.StringType),
        th.Property("progress_measure_lower_conf_interval", th.StringType),
        th.Property("progress_measure_upper_conf_interval", th.StringType),
        th.Property("absent_or_not_able_to_access_percent", th.StringType),
        th.Property("working_towards_expected_standard_pupil_percent", th.StringType),
        th.Property("absent_or_disapplied_percent", th.StringType),
        th.Property("higher_standard", th.StringType),
        th.Property("progress_measure_unadjusted", th.StringType),
        th.Property("progress_measure_description", th.StringType),
    ).to_dict()


# ── KS2 Information (wide format: one row per school, context/demographics) ──

class EESKS2InfoStream(EESDatasetStream):
    name = "ees_ks2_info"
    primary_keys = ["school_urn", "time_period"]
    _publication_slug = "key-stage-2-attainment"
    _target_filename = "ks2_school_information_data"
    schema = th.PropertiesList(
        th.Property("time_period", th.StringType, required=True),
        th.Property("school_urn", th.StringType, required=True),
        th.Property("school_laestab", th.StringType),
        th.Property("school_name", th.StringType),
        th.Property("nftype", th.StringType),
        th.Property("reldenom", th.StringType),
        th.Property("agerange", th.StringType),
        th.Property("totpups", th.StringType),
        th.Property("telig", th.StringType),
        th.Property("belig", th.StringType),
        th.Property("gelig", th.StringType),
        th.Property("ptfsm6cla1a", th.StringType),
        th.Property("ptnotfsm6cla1a", th.StringType),
        th.Property("ptealgrp2", th.StringType),
        th.Property("ptmobn", th.StringType),
        th.Property("psenelk", th.StringType),
        th.Property("psenele", th.StringType),
        th.Property("psenelek", th.StringType),
        th.Property("telig_3yr", th.StringType),
    ).to_dict()


# ── KS4 Performance (long format: one row per school × breakdown × sex) ─────
# File: 202425_performance_tables_schools_revised.csv (156 cols)
# Dimensions: breakdown_topic, breakdown, sex, disadvantage_status, etc.
# Metrics are already in separate columns (attainment8_average, progress8_average, etc.)

class EESKS4PerformanceStream(EESDatasetStream):
    name = "ees_ks4_performance"
    primary_keys = ["school_urn", "time_period", "breakdown_topic", "breakdown", "sex"]
    _publication_slug = "key-stage-4-performance"
    _target_filename = "performance_tables_schools"
    schema = th.PropertiesList(
        th.Property("time_period", th.StringType, required=True),
        th.Property("school_urn", th.StringType, required=True),
        th.Property("school_laestab", th.StringType),
        th.Property("school_name", th.StringType),
        th.Property("establishment_type_group", th.StringType),
        th.Property("breakdown_topic", th.StringType, required=True),
        th.Property("breakdown", th.StringType, required=True),
        th.Property("sex", th.StringType, required=True),
        th.Property("disadvantage_status", th.StringType),
        th.Property("first_language", th.StringType),
        th.Property("prior_attainment", th.StringType),
        th.Property("mobility", th.StringType),
        # Pupil counts
        th.Property("pupil_count", th.StringType),
        th.Property("pupil_percent", th.StringType),
        # Attainment 8
        th.Property("attainment8_sum", th.StringType),
        th.Property("attainment8_average", th.StringType),
        # English & Maths
        th.Property("engmath_entering_total", th.StringType),
        th.Property("engmath_entering_percent", th.StringType),
        th.Property("engmath_95_total", th.StringType),
        th.Property("engmath_95_percent", th.StringType),
        th.Property("engmath_94_total", th.StringType),
        th.Property("engmath_94_percent", th.StringType),
        # EBacc
        th.Property("ebacc_entering_total", th.StringType),
        th.Property("ebacc_entering_percent", th.StringType),
        th.Property("ebacc_95_total", th.StringType),
        th.Property("ebacc_95_percent", th.StringType),
        th.Property("ebacc_94_total", th.StringType),
        th.Property("ebacc_94_percent", th.StringType),
        th.Property("ebacc_aps_sum", th.StringType),
        th.Property("ebacc_aps_average", th.StringType),
        # Progress 8
        th.Property("progress8_pupil_count", th.StringType),
        th.Property("progress8_sum", th.StringType),
        th.Property("progress8_average", th.StringType),
        th.Property("progress8_lower_95_ci", th.StringType),
        th.Property("progress8_upper_95_ci", th.StringType),
        # Progress 8 elements
        th.Property("progress8eng_average", th.StringType),
        th.Property("progress8mat_average", th.StringType),
        th.Property("progress8ebacc_average", th.StringType),
        th.Property("progress8open_average", th.StringType),
        # GCSE grades
        th.Property("gcse_91_total", th.StringType),
        th.Property("gcse_91_percent", th.StringType),
        # EBacc subject entry/achievement
        th.Property("ebacceng_entering_percent", th.StringType),
        th.Property("ebaccmat_entering_percent", th.StringType),
        th.Property("ebaccsci_entering_percent", th.StringType),
        th.Property("ebacchum_entering_percent", th.StringType),
        th.Property("ebacclan_entering_percent", th.StringType),
    ).to_dict()


# ── KS4 Information (wide format: one row per school, context/demographics) ──
# File: 202425_information_about_schools_provisional.csv (38 cols)

class EESKS4InfoStream(EESDatasetStream):
    name = "ees_ks4_info"
    primary_keys = ["school_urn", "time_period"]
    _publication_slug = "key-stage-4-performance"
    _target_filename = "information_about_schools"
    schema = th.PropertiesList(
        th.Property("time_period", th.StringType, required=True),
        th.Property("school_urn", th.StringType, required=True),
        th.Property("school_laestab", th.StringType),
        th.Property("school_name", th.StringType),
        th.Property("establishment_type_group", th.StringType),
        th.Property("reldenom", th.StringType),
        th.Property("admpol_pt", th.StringType),
        th.Property("egender", th.StringType),
        th.Property("agerange", th.StringType),
        th.Property("allks_pupil_count", th.StringType),
        th.Property("allks_boys_count", th.StringType),
        th.Property("allks_girls_count", th.StringType),
        th.Property("endks4_pupil_count", th.StringType),
        th.Property("ks2_scaledscore_average", th.StringType),
        th.Property("sen_with_ehcp_pupil_percent", th.StringType),
        th.Property("sen_pupil_percent", th.StringType),
        th.Property("sen_no_ehcp_pupil_percent", th.StringType),
        th.Property("attainment8_diffn", th.StringType),
        th.Property("progress8_diffn", th.StringType),
        th.Property("progress8_banding", th.StringType),
    ).to_dict()


# ── Census (school-level pupil characteristics) ─────────────────────────────
# File: spc_school_level_underlying_data_YYYY.csv (269 cols, in supporting-files/)
# Uses 'urn' not 'school_urn'. Filename has yearly suffix that changes.

class EESCensusStream(EESDatasetStream):
    name = "ees_census"
    primary_keys = ["school_urn", "time_period"]
    _publication_slug = "school-pupils-and-their-characteristics"
    _target_filename = "spc_school_level_underlying_data"
    _urn_column = "urn"
    _encoding = "latin-1"

    # Map verbose CSV column names to clean Singer field names
    _column_renames = {
        "headcount of pupils":                                                        "total_pupils",
        "headcount total female":                                                     "female_pupils",
        "headcount total male":                                                       "male_pupils",
        "full-time pupils":                                                           "full_time_pupils",
        "part-time pupils":                                                           "part_time_pupils",
        "total boarders":                                                             "total_boarders",
        "number of pupils known to be eligible for free school meals":                "fsm_eligible_n",
        "% of pupils known to be eligible for free school meals":                    "fsm_pct",
        "% of pupils whose first language is known or believed to be other than English": "eal_pct",
        "% of pupils who are a young carer":                                          "young_carer_pct",
        "% of pupils classified as white British ethnic origin":                      "ethnicity_white_british_pct",
        "% of pupils classified as any other white background ethnic origin":         "ethnicity_white_other_pct",
        "% of pupils classified as Gypsy/Roma ethnic origin":                        "ethnicity_gypsy_roma_pct",
        "% of pupils classified as white and black Caribbean ethnic origin":          "ethnicity_mixed_wbc_pct",
        "% of pupils classified as white and black African ethnic origin":            "ethnicity_mixed_wba_pct",
        "% of pupils classified as white and Asian ethnic origin":                   "ethnicity_mixed_wa_pct",
        "% of pupils classified as any other mixed background ethnic origin":         "ethnicity_mixed_other_pct",
        "% of pupils classified as Indian ethnic origin":                             "ethnicity_asian_indian_pct",
        "% of pupils classified as Pakistani ethnic origin":                          "ethnicity_asian_pakistani_pct",
        "% of pupils classified as Bangladeshi ethnic origin":                        "ethnicity_asian_bangladeshi_pct",
        "% of pupils classified as any other Asian background ethnic origin":         "ethnicity_asian_other_pct",
        "% of pupils classified as Caribbean ethnic origin":                          "ethnicity_black_caribbean_pct",
        "% of pupils classified as African ethnic origin":                            "ethnicity_black_african_pct",
        "% of pupils classified as any other black background ethnic origin":         "ethnicity_black_other_pct",
        "% of pupils classified as Chinese ethnic origin":                            "ethnicity_chinese_pct",
        "% of pupils classified as any other ethnic group ethnic origin":             "ethnicity_other_pct",
        "% of pupils unclassified":                                                   "ethnicity_unclassified_pct",
    }

    schema = th.PropertiesList(
        th.Property("time_period", th.StringType, required=True),
        th.Property("school_urn", th.StringType, required=True),
        th.Property("school_name", th.StringType),
        th.Property("laestab", th.StringType),
        th.Property("phase_type_grouping", th.StringType),
        # Pupil counts
        th.Property("total_pupils", th.StringType),
        th.Property("female_pupils", th.StringType),
        th.Property("male_pupils", th.StringType),
        th.Property("full_time_pupils", th.StringType),
        th.Property("part_time_pupils", th.StringType),
        th.Property("total_boarders", th.StringType),
        # FSM
        th.Property("fsm_eligible_n", th.StringType),
        th.Property("fsm_pct", th.StringType),
        # EAL & young carers
        th.Property("eal_pct", th.StringType),
        th.Property("young_carer_pct", th.StringType),
        # Ethnicity percentages
        th.Property("ethnicity_white_british_pct", th.StringType),
        th.Property("ethnicity_white_other_pct", th.StringType),
        th.Property("ethnicity_gypsy_roma_pct", th.StringType),
        th.Property("ethnicity_mixed_wbc_pct", th.StringType),
        th.Property("ethnicity_mixed_wba_pct", th.StringType),
        th.Property("ethnicity_mixed_wa_pct", th.StringType),
        th.Property("ethnicity_mixed_other_pct", th.StringType),
        th.Property("ethnicity_asian_indian_pct", th.StringType),
        th.Property("ethnicity_asian_pakistani_pct", th.StringType),
        th.Property("ethnicity_asian_bangladeshi_pct", th.StringType),
        th.Property("ethnicity_asian_other_pct", th.StringType),
        th.Property("ethnicity_black_caribbean_pct", th.StringType),
        th.Property("ethnicity_black_african_pct", th.StringType),
        th.Property("ethnicity_black_other_pct", th.StringType),
        th.Property("ethnicity_chinese_pct", th.StringType),
        th.Property("ethnicity_other_pct", th.StringType),
        th.Property("ethnicity_unclassified_pct", th.StringType),
    ).to_dict()


# ── Admissions ───────────────────────────────────────────────────────────────
# File: AppsandOffers_YYYY_SchoolLevelDDMMYYYY.csv (37 cols, in supporting-files/)
# Wide format, no geographic_level column. Uses school_urn.

class EESAdmissionsStream(EESDatasetStream):
    name = "ees_admissions"
    primary_keys = ["school_urn", "time_period"]
    _publication_slug = "primary-and-secondary-school-applications-and-offers"
    _target_filename = "SchoolLevel"
    _encoding = "utf-8-sig"  # UTF-8 with BOM — sig variant strips the BOM automatically
    schema = th.PropertiesList(
        th.Property("time_period", th.StringType, required=True),
        th.Property("school_urn", th.StringType, required=True),
        th.Property("school_name", th.StringType),
        th.Property("school_laestab_as_used", th.StringType),
        th.Property("school_phase", th.StringType),
        th.Property("entry_year", th.StringType),
        # Places and offers
        th.Property("total_number_places_offered", th.StringType),
        th.Property("number_preferred_offers", th.StringType),
        th.Property("number_1st_preference_offers", th.StringType),
        th.Property("number_2nd_preference_offers", th.StringType),
        th.Property("number_3rd_preference_offers", th.StringType),
        # Applications
        th.Property("times_put_as_any_preferred_school", th.StringType),
        th.Property("times_put_as_1st_preference", th.StringType),
        th.Property("times_put_as_2nd_preference", th.StringType),
        th.Property("times_put_as_3rd_preference", th.StringType),
        # Proportions
        th.Property("proportion_1stprefs_v_1stprefoffers", th.StringType),
        th.Property("proportion_1stprefs_v_totaloffers", th.StringType),
        # Cross-LA
        th.Property("all_applications_from_another_LA", th.StringType),
        th.Property("offers_to_applicants_from_another_LA", th.StringType),
        # Context
        th.Property("establishment_type", th.StringType),
        th.Property("denomination", th.StringType),
        th.Property("FSM_eligible_percent", th.StringType),
        th.Property("admissions_policy", th.StringType),
        th.Property("urban_rural", th.StringType),
    ).to_dict()


# Note: Phonics (phonics-screening-check-attainment) has NO school-level data
# on EES. Only national and LA-level files are published.


# ── KS2 National Headlines (national level only — one row per year) ───────────
# Dataset: "Key stage 2 attainment: national headlines"
# URL: https://explore-education-statistics.service.gov.uk/data-catalogue/data-set/
#      58bb4b03-c6df-447f-bb7e-b82970c4d974/csv
# This is a stable data-catalogue CSV endpoint (not a versioned release ZIP).
# Covers 2015/16 → latest; COVID years (2019/20, 2020/21) are suppressed ('x').

_KS2_NATIONAL_CSV_URL = (
    "https://explore-education-statistics.service.gov.uk/data-catalogue/"
    "data-set/58bb4b03-c6df-447f-bb7e-b82970c4d974/csv"
)

_KS2_NATIONAL_COL_MAP = {
    "pt_rwm_exp":    "rwm_expected_pct",
    "pt_rwm_high":   "rwm_high_pct",
    "pt_read_exp":   "reading_expected_pct",
    "pt_read_high":  "reading_high_pct",
    "pt_mat_exp":    "maths_expected_pct",
    "pt_mat_high":   "maths_high_pct",
    "pt_writta_exp": "writing_expected_pct",
    "pt_writta_gd":  "writing_gd_pct",
    "pt_gps_exp":    "gps_expected_pct",
    "pt_gps_high":   "gps_high_pct",
    "pt_scita_exp":  "science_expected_pct",
    "avg_readscore": "reading_avg_score",
    "avg_gpsscore":  "gps_avg_score",
    "avg_matscore":  "maths_avg_score",
}


class EESKs2NationalStream(Stream):
    """National KS2 headline averages — one row per academic year.

    Fetches the DfE EES data-catalogue CSV directly (stable URL, not versioned
    release ZIP). Filters to geographic_level == 'National' and
    school_type == 'AllSchools' so only the England-wide headline row per year
    is emitted. COVID years (2019/20, 2020/21) are naturally absent (suppressed
    with 'x' → treated as null downstream in dbt staging).
    """

    name = "ees_ks2_national"
    primary_keys = ["time_period"]
    replication_key = None

    schema = th.PropertiesList(
        th.Property("time_period", th.StringType, required=True),
        th.Property("rwm_expected_pct",   th.StringType),
        th.Property("rwm_high_pct",       th.StringType),
        th.Property("reading_expected_pct", th.StringType),
        th.Property("reading_high_pct",   th.StringType),
        th.Property("maths_expected_pct", th.StringType),
        th.Property("maths_high_pct",     th.StringType),
        th.Property("writing_expected_pct", th.StringType),
        th.Property("writing_gd_pct",     th.StringType),
        th.Property("gps_expected_pct",   th.StringType),
        th.Property("gps_high_pct",       th.StringType),
        th.Property("science_expected_pct", th.StringType),
        th.Property("reading_avg_score",  th.StringType),
        th.Property("gps_avg_score",      th.StringType),
        th.Property("maths_avg_score",    th.StringType),
    ).to_dict()

    def get_records(self, context):
        import pandas as pd

        self.logger.info("Downloading KS2 national headlines: %s", _KS2_NATIONAL_CSV_URL)
        resp = requests.get(_KS2_NATIONAL_CSV_URL, timeout=60)
        resp.raise_for_status()

        df = pd.read_csv(
            io.BytesIO(resp.content),
            dtype=str,
            keep_default_na=False,
        )

        # Normalise column names to lowercase
        df.columns = [c.strip().lower() for c in df.columns]

        # Keep only the England national headline row per year
        if "geographic_level" in df.columns:
            df = df[df["geographic_level"].str.strip().str.lower() == "national"]
        if "school_type" in df.columns:
            df = df[df["school_type"].str.strip().str.lower() == "allschools"]

        self.logger.info("Emitting %d national KS2 rows", len(df))

        for _, row in df.iterrows():
            record = {"time_period": row.get("time_period", "").strip()}
            for csv_col, field in _KS2_NATIONAL_COL_MAP.items():
                record[field] = row.get(csv_col, "").strip()
            yield record


# ── KS4 National Headlines (national level only — one row per year) ──────────
# Dataset: "National characteristics summary data" (Key stage 4 performance).
# Official England state-funded headline measures, 2018/19 → latest.
# Suppressed values ('z', 'x') → NULL downstream. Progress 8 is legitimately
# absent in years with no KS2 baseline (e.g. 2024/25) — that is DfE policy,
# not missing data.

_KS4_NATIONAL_CSV_URL = (
    "https://explore-education-statistics.service.gov.uk/data-catalogue/"
    "data-set/1b649e16-01e8-435b-a814-56be2faf9054/csv"
)

_KS4_NATIONAL_COL_MAP = {
    "attainment8_average":    "attainment_8_score",
    "progress8_average":      "progress_8_score",
    "engmath_94_percent":     "english_maths_standard_pass_pct",
    "engmath_95_percent":     "english_maths_strong_pass_pct",
    "ebacc_entering_percent": "ebacc_entry_pct",
    "ebacc_94_percent":       "ebacc_standard_pass_pct",
    "ebacc_95_percent":       "ebacc_strong_pass_pct",
    "ebacc_aps_average":      "ebacc_avg_score",
}


class EESKs4NationalStream(Stream):
    """National KS4 headline averages — one row per academic year.

    Filters to geographic_level == 'National', establishment_type_group ==
    'All state-funded', breakdown_topic == 'Total', breakdown == 'Total'
    so only the England-wide all-pupils row per year is emitted.
    """

    name = "ees_ks4_national"
    primary_keys = ["time_period"]
    replication_key = None

    schema = th.PropertiesList(
        th.Property("time_period", th.StringType, required=True),
        *[th.Property(out, th.StringType) for out in _KS4_NATIONAL_COL_MAP.values()],
    ).to_dict()

    def get_records(self, context):
        import pandas as pd

        self.logger.info("Downloading KS4 national headlines: %s", _KS4_NATIONAL_CSV_URL)
        resp = requests.get(_KS4_NATIONAL_CSV_URL, timeout=60)
        resp.raise_for_status()

        df = pd.read_csv(io.BytesIO(resp.content), dtype=str, keep_default_na=False)
        df.columns = [c.strip().lower() for c in df.columns]

        for col, want in [
            ("geographic_level", "national"),
            ("establishment_type_group", "all state-funded"),
            ("breakdown_topic", "total"),
            ("breakdown", "total"),
        ]:
            if col in df.columns:
                df = df[df[col].str.strip().str.lower() == want]

        self.logger.info("Emitting %d national KS4 rows", len(df))
        for _, row in df.iterrows():
            record = {"time_period": row.get("time_period", "").strip()}
            for csv_col, field in _KS4_NATIONAL_COL_MAP.items():
                record[field] = row.get(csv_col, "").strip()
            yield record


# ── Legacy KS2 (pre-COVID wide format from DfE performance tables) ────────────
# The DfE "Compare School Performance" site published school-level KS2 CSVs
# in a wide format (one row per school, ~300 columns).  EES only has school-level
# data from 2022-23 onwards.  This stream loads the 2015-16 to 2018-19 CSVs
# from a configurable base URL and maps the old column names to match the
# staging output schema so dbt can union them with the EES data.

# Column mapping: old DfE CSV column → Singer field name (matches stg output)
_LEGACY_KS2_COLUMN_MAP = {
    "URN": "urn",
    "TOTPUPS": "total_pupils",
    "TELIG": "eligible_pupils",
    "PTRWM_EXP": "rwm_expected_pct",
    "PTRWM_HIGH": "rwm_high_pct",
    "PTREAD_EXP": "reading_expected_pct",
    "PTREAD_HIGH": "reading_high_pct",
    "READ_AVERAGE": "reading_avg_score",
    "READPROG": "reading_progress",
    "PTWRITTA_EXP": "writing_expected_pct",
    "PTWRITTA_HIGH": "writing_high_pct",
    "WRITPROG": "writing_progress",
    "PTMAT_EXP": "maths_expected_pct",
    "PTMAT_HIGH": "maths_high_pct",
    "MAT_AVERAGE": "maths_avg_score",
    "MATPROG": "maths_progress",
    "PTGPS_EXP": "gps_expected_pct",
    "PTGPS_HIGH": "gps_high_pct",
    "GPS_AVERAGE": "gps_avg_score",
    # Absence / unable to access
    "PTREAD_AT": "reading_absence_pct",
    "PTMAT_AT": "maths_absence_pct",
    "PTGPS_AT": "gps_absence_pct",
    # Gender breakdown for Reading, Writing & Maths
    "PTRWM_EXP_B": "rwm_expected_boys_pct",
    "PTRWM_HIGH_B": "rwm_high_boys_pct",
    "PTRWM_EXP_G": "rwm_expected_girls_pct",
    "PTRWM_HIGH_G": "rwm_high_girls_pct",
    # Disadvantaged breakdown for Reading, Writing & Maths
    "PTRWM_EXP_FSM6CLA1A": "rwm_expected_disadvantaged_pct",
    "PTRWM_EXP_NotFSM6CLA1A": "rwm_expected_non_disadvantaged_pct",
    "DIFFN_RWM_EXP": "disadvantaged_gap",
    # Context
    "PTFSM6CLA1A": "disadvantaged_pct",
    "PTEALGRP2": "eal_pct",
    "PSENELK": "sen_support_pct",
    "PSENELE": "sen_ehcp_pct",
    "PTMOBN": "stability_pct",
}

class LegacyKS2Stream(Stream):
    """Stream for pre-COVID KS2 data from DfE performance tables CSVs.

    Downloads CSVs from URLs configured in legacy_ks2_urls (a mapping of
    6-digit year code → download URL), maps old DfE column names to match
    the stg_ees_ks2 output schema, and emits one record per school per year.
    """

    name = "legacy_ks2"
    primary_keys = ["urn", "year"]
    replication_key = None

    schema = th.PropertiesList(
        th.Property("urn", th.StringType, required=True),
        th.Property("year", th.StringType, required=True),
        th.Property("total_pupils", th.StringType),
        th.Property("eligible_pupils", th.StringType),
        th.Property("rwm_expected_pct", th.StringType),
        th.Property("rwm_high_pct", th.StringType),
        th.Property("reading_expected_pct", th.StringType),
        th.Property("reading_high_pct", th.StringType),
        th.Property("reading_avg_score", th.StringType),
        th.Property("reading_progress", th.StringType),
        th.Property("writing_expected_pct", th.StringType),
        th.Property("writing_high_pct", th.StringType),
        th.Property("writing_progress", th.StringType),
        th.Property("maths_expected_pct", th.StringType),
        th.Property("maths_high_pct", th.StringType),
        th.Property("maths_avg_score", th.StringType),
        th.Property("maths_progress", th.StringType),
        th.Property("gps_expected_pct", th.StringType),
        th.Property("gps_high_pct", th.StringType),
        th.Property("gps_avg_score", th.StringType),
        th.Property("reading_absence_pct", th.StringType),
        th.Property("maths_absence_pct", th.StringType),
        th.Property("gps_absence_pct", th.StringType),
        th.Property("rwm_expected_boys_pct", th.StringType),
        th.Property("rwm_high_boys_pct", th.StringType),
        th.Property("rwm_expected_girls_pct", th.StringType),
        th.Property("rwm_high_girls_pct", th.StringType),
        th.Property("rwm_expected_disadvantaged_pct", th.StringType),
        th.Property("rwm_expected_non_disadvantaged_pct", th.StringType),
        th.Property("disadvantaged_gap", th.StringType),
        th.Property("disadvantaged_pct", th.StringType),
        th.Property("eal_pct", th.StringType),
        th.Property("sen_support_pct", th.StringType),
        th.Property("sen_ehcp_pct", th.StringType),
        th.Property("stability_pct", th.StringType),
    ).to_dict()

    def get_records(self, context):
        import pandas as pd

        url_map = self.config.get("legacy_ks2_urls", {})
        if not url_map:
            self.logger.warning("legacy_ks2_urls not configured, skipping legacy KS2")
            return

        self.logger.info("Loading legacy KS2 for %d year(s)", len(url_map))

        for year_code, url in url_map.items():
            self.logger.info("Downloading %s for %s", url, year_code)
            try:
                resp = requests.get(url, timeout=120)
                resp.raise_for_status()
            except Exception as e:
                self.logger.warning("Failed to download %s: %s", url, e)
                continue

            content = resp.content

            # Auto-detect ZIP — the DfE annual archives contain both KS2 and KS4
            # CSVs in one ZIP. If the download is a ZIP, extract england_ks2final.csv;
            # otherwise treat the content as a bare CSV (legacy individual-file URLs).
            csv_bytes = None
            try:
                zf = zipfile.ZipFile(io.BytesIO(content))
                target = next(
                    (n for n in zf.namelist() if "ks2final" in n.lower() and n.endswith(".csv")),
                    None,
                )
                if target:
                    with zf.open(target) as f:
                        csv_bytes = f.read()
                    self.logger.info("Extracted %s from ZIP for %s", target, year_code)
                else:
                    self.logger.warning("No ks2final CSV found in ZIP for %s", year_code)
                    continue
            except zipfile.BadZipFile:
                # Not a ZIP — treat as a bare CSV file
                csv_bytes = content

            df = pd.read_csv(
                io.BytesIO(csv_bytes),
                dtype=str,
                keep_default_na=False,
                encoding="latin-1",
            )

            # Strip BOM
            cols = list(df.columns)
            if cols:
                cols[0] = cols[0].lstrip("\ufeff").lstrip("ï»¿")
                df.columns = cols

            # Filter to valid URNs
            if "URN" in df.columns:
                df = df[df["URN"].str.match(r"^\d+$", na=False)]

            self.logger.info("Emitting %d schools for %s", len(df), year_code)

            for _, row in df.iterrows():
                record = {"year": year_code}
                for old_col, new_col in _LEGACY_KS2_COLUMN_MAP.items():
                    val = row.get(old_col, "")
                    # Strip % suffix — old DfE CSVs use "57%" not "57"
                    if isinstance(val, str) and val.endswith("%"):
                        val = val[:-1]
                    record[new_col] = val
                yield record


# ── Legacy KS4 (pre-EES wide format from DfE performance tables) ──────────────
# The DfE "Compare School Performance" ZIPs include england_ks4final.csv in a
# wide format (one row per school, ~416 columns, uppercase abbreviated names).
# EES only hosts 2 years of KS4 data; this stream backfills 2015-16 to 2018-19.
# Column mapping: old DfE CSV column → Singer field name (matches stg output).

_LEGACY_KS4_COLUMN_MAP = {
    "URN":              "urn",
    "TPUP":             "total_pupils",
    # Attainment 8
    "ATT8SCR":          "attainment_8_score",
    # Progress 8
    "P8MEA":            "progress_8_score",
    "P8CILOW":          "progress_8_lower_ci",
    "P8CIUPP":          "progress_8_upper_ci",
    "P8MEAENG":         "progress_8_english",
    "P8MEAMAT":         "progress_8_maths",
    "P8MEAEBAC":        "progress_8_ebacc",
    "P8MEAOPEN":        "progress_8_open",
    # English & Maths pass rates (% suffix stripped at extract time)
    "PTL2BASICS_95":    "english_maths_strong_pass_pct",
    "PTL2BASICS_94":    "english_maths_standard_pass_pct",
    # EBacc
    "PTEBACC_E_PTQ_EE": "ebacc_entry_pct",
    "PTEBACC_95":       "ebacc_strong_pass_pct",
    "PTEBACC_94":       "ebacc_standard_pass_pct",
    # Context
    "PSENSE4":          "sen_ehcp_pct",
    "PSENAPK4":         "sen_support_pct",
}


class LegacyKS4Stream(Stream):
    """Stream for pre-EES KS4 data from DfE 'Compare School Performance' ZIPs.

    Downloads ZIPs from URLs configured in legacy_ks4_urls (a mapping of
    6-digit year code → download URL), extracts england_ks4final.csv from each,
    maps old DfE column names to match stg_ees_ks4 output schema, and emits
    one record per school per year.  The % suffix present on percentage columns
    (e.g. "39.60%") is stripped here so safe_numeric in dbt can cast cleanly.
    """

    name = "legacy_ks4"
    primary_keys = ["urn", "year"]
    replication_key = None

    schema = th.PropertiesList(
        th.Property("urn",                            th.StringType, required=True),
        th.Property("year",                           th.StringType, required=True),
        th.Property("total_pupils",                   th.StringType),
        th.Property("attainment_8_score",             th.StringType),
        th.Property("progress_8_score",               th.StringType),
        th.Property("progress_8_lower_ci",            th.StringType),
        th.Property("progress_8_upper_ci",            th.StringType),
        th.Property("progress_8_english",             th.StringType),
        th.Property("progress_8_maths",               th.StringType),
        th.Property("progress_8_ebacc",               th.StringType),
        th.Property("progress_8_open",                th.StringType),
        th.Property("english_maths_strong_pass_pct",  th.StringType),
        th.Property("english_maths_standard_pass_pct", th.StringType),
        th.Property("ebacc_entry_pct",                th.StringType),
        th.Property("ebacc_strong_pass_pct",          th.StringType),
        th.Property("ebacc_standard_pass_pct",        th.StringType),
        th.Property("sen_ehcp_pct",                   th.StringType),
        th.Property("sen_support_pct",                th.StringType),
    ).to_dict()

    def get_records(self, context):
        import pandas as pd

        url_map = self.config.get("legacy_ks4_urls", {})
        if not url_map:
            self.logger.warning("legacy_ks4_urls not configured, skipping legacy KS4")
            return

        self.logger.info("Loading legacy KS4 for %d year(s)", len(url_map))

        for year_code, url in url_map.items():
            self.logger.info("Downloading %s for %s", url, year_code)
            try:
                resp = requests.get(url, timeout=120)
                resp.raise_for_status()
            except Exception as e:
                self.logger.warning("Failed to download %s: %s", url, e)
                continue

            try:
                zf = zipfile.ZipFile(io.BytesIO(resp.content))
            except zipfile.BadZipFile as e:
                self.logger.warning("Not a ZIP for %s: %s", year_code, e)
                continue

            # Find england_ks4final.csv inside the ZIP
            target = next(
                (n for n in zf.namelist() if "ks4final" in n.lower() and n.endswith(".csv")),
                None,
            )
            if not target:
                self.logger.warning("england_ks4final.csv not found in ZIP for %s", year_code)
                continue

            with zf.open(target) as f:
                df = pd.read_csv(
                    f,
                    dtype=str,
                    keep_default_na=False,
                    encoding="latin-1",
                )

            # Strip BOM from first column name
            cols = list(df.columns)
            if cols:
                cols[0] = cols[0].lstrip("\ufeff").lstrip("ï»¿")
                df.columns = cols

            # Filter to school-level rows: URN must be a plain integer
            if "URN" in df.columns:
                df = df[df["URN"].str.match(r"^\d+$", na=False)]

            self.logger.info("Emitting %d schools for %s", len(df), year_code)

            for _, row in df.iterrows():
                record = {"year": year_code}
                for old_col, new_col in _LEGACY_KS4_COLUMN_MAP.items():
                    val = row.get(old_col, "")
                    # Strip % suffix — legacy DfE CSVs use "39.60%" not "39.60"
                    if isinstance(val, str) and val.endswith("%"):
                        val = val[:-1]
                    record[new_col] = val
                yield record


class TapUKEES(Tap):
    """Singer tap for UK Explore Education Statistics."""

    name = "tap-uk-ees"

    config_jsonschema = th.PropertiesList(
        th.Property("base_url", th.StringType, description="EES API base URL"),
        th.Property(
            "latest_only",
            th.BooleanType,
            description="Only fetch the latest release per publication (default: False — fetches all historical releases)",
            default=False,
        ),
        th.Property(
            "legacy_ks2_urls",
            th.ObjectType(),
            description="Mapping of 6-digit year code to download URL for legacy KS2 CSVs (e.g. {\"201819\": \"https://...\"})",
        ),
        th.Property(
            "legacy_ks4_urls",
            th.ObjectType(),
            description="Mapping of 6-digit year code to download URL for legacy KS4 ZIPs (e.g. {\"201819\": \"https://...\"})",
        ),
    ).to_dict()

    def discover_streams(self):
        return [
            EESKS2AttainmentStream(self),
            EESKS2InfoStream(self),
            EESKS4PerformanceStream(self),
            EESKS4InfoStream(self),
            EESCensusStream(self),
            EESAdmissionsStream(self),
            LegacyKS2Stream(self),
            LegacyKS4Stream(self),
            EESKs2NationalStream(self),
            EESKs4NationalStream(self),
        ]


if __name__ == "__main__":
    TapUKEES.cli()
