"""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 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

        release_id = get_content_release_id(self._publication_slug)
        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.error(
                "File matching '%s' not found in ZIP. Available: %s",
                self._target_filename,
                [n for n in all_files if n.endswith(".csv")],
            )
            return

        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

        # 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)
        urn_col = self._urn_column
        if urn_col in df.columns:
            df = df[df[urn_col].notna() & (df[urn_col] != "")]

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

        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)
            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.


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"),
    ).to_dict()

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


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