hf_store.py

"""
hf_store.py
───────────
HuggingFace dataset push / pull for KYS school scraper.

Dataset folder structure (all in one repo):
  scraped_data/raw/{state_key}.parquet    — raw scraped data per state
  scraped_data/mapped/mapped_master_{date}.parquet — final mapped master
  district_reference/district_reference.parquet    — SF mirror (Dataset 1)
  mapping_rules/manual_district_mapping.parquet           — district mapping rules (Dataset 2)
  mapping_rules/manual_block_mapping.parquet              — block-level splits (Dataset 2)
  metadata.json                                    — scrape metadata
"""

import os
import json
import tempfile
from district_mapper import NON_ACTUAL_STATES
import pandas as pd
from datetime import datetime

ALL_REQUIRED_CATEGORIES = {"3", "5", "6", "7", "8", "10", "11"}

# ── Load .env if present (for local dev) ─────────────────────────────────────
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # python-dotenv not installed; rely on system env vars


# ─────────────────────────────────────────────────────────────────────────────
# Credentials helper
# ─────────────────────────────────────────────────────────────────────────────

def get_hf_credentials() -> tuple:
    """
    Returns (token, repo) from environment variables.
    Set HF_TOKEN and HF_REPO in your .env file (local)
    or as Secrets in HuggingFace Spaces settings.
    """
    token = os.environ.get("HF_TOKEN", "").strip()
    repo  = os.environ.get("HF_REPO",  "").strip()
    return token, repo


# ── Lazy imports so the app still loads if huggingface_hub isn't installed ────
def _hf():
    try:
        from huggingface_hub import HfApi
        return HfApi()
    except ImportError:
        raise RuntimeError(
            "huggingface_hub is not installed. Run: pip install huggingface_hub pyarrow"
        )

import re

def _state_key(state_name: str) -> str:
    """Normalise state name to a safe filename key, e.g. 'ANDHRA PRADESH' → 'andhra_pradesh'"""
    return re.sub(r"[^a-z0-9]+", "_", state_name.strip().lower()).strip("_")


def _safe_for_parquet(df: pd.DataFrame) -> pd.DataFrame:
    """
    Convert all object-dtype columns to string so pyarrow never
    chokes on mixed int/str/NaN in a single column.
    """
    df = df.copy()
    for col in df.columns:
        if df[col].dtype == object:
            df[col] = df[col].astype(str).replace("nan", "")
    return df


def _upload_parquet(df: pd.DataFrame, path_in_repo: str, token: str, repo: str, commit_msg: str) -> str:
    """Helper: write df to a temp parquet and upload to HF dataset. Returns commit URL."""
    api = _hf()
    df = _safe_for_parquet(df)
    with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
        tmp_path = tmp.name
    try:
        df.to_parquet(tmp_path, index=False)
        url = api.upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo=path_in_repo,
            repo_id=repo,
            repo_type="dataset",
            token=token,
            commit_message=commit_msg,
        )
    finally:
        os.unlink(tmp_path)
    return url


def _download_parquet(path_in_repo: str, token: str, repo: str) -> pd.DataFrame:
    """Helper: download a parquet file from HF dataset. Returns DataFrame or empty DF."""
    api = _hf()
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            local_path = api.hf_hub_download(
                repo_id=repo,
                filename=path_in_repo,
                repo_type="dataset",
                token=token,
                local_dir=tmpdir,
                local_dir_use_symlinks=False,
                force_download=True,
            )
            return pd.read_parquet(local_path)
    except Exception:
        return pd.DataFrame()


# ─────────────────────────────────────────────────────────────────────────────
# Dataset 1: district_reference  (SF mirror)
# ─────────────────────────────────────────────────────────────────────────────

DISTRICT_REF_PATH = "district_reference/district_reference.parquet"

def push_district_reference(df: pd.DataFrame, token: str, repo: str) -> str:
    """
    Push the SF district reference table to HF dataset.
    Expected columns: State, District, Status ('present' | 'new districts found')
    Returns commit URL.
    """
    return _upload_parquet(df, DISTRICT_REF_PATH, token, repo, "Update district_reference")


def pull_district_reference(token: str, repo: str) -> pd.DataFrame:
    """
    Pull the SF district reference from HF dataset.
    Returns DataFrame with columns: State, District, Status
    Returns empty DataFrame if not yet seeded.
    """
    df = _download_parquet(DISTRICT_REF_PATH, token, repo)
    if df.empty:
        return pd.DataFrame(columns=["State", "District", "Status"])
    return df


def seed_district_reference_from_csv(csv_path: str, token: str, repo: str) -> tuple:
    """
    Seed Dataset 1 from a CSV (like master_all_states.csv or SF mapping CSV).
    All states/districts in the CSV → Status = 'present'
    Any states in current scraped data NOT in CSV → Status = 'new districts found'
    Returns (df, commit_url)
    """
    csv_df = pd.read_csv(csv_path)
    
    # Try to find state column
    state_col = "State" if "State" in csv_df.columns else "School_State__c"
    # Try to find district column
    if "School_District__c" in csv_df.columns:
        district_col = "School_District__c"
    elif "KYS District " in csv_df.columns:
        district_col = "KYS District "
    else:
        # Fallback to the second column if it's the old SF mapping format
        district_col = csv_df.columns[1]

    csv_df = csv_df.dropna(subset=[state_col])
    csv_df["State"] = csv_df[state_col].astype(str).str.strip().str.upper()
    csv_df["District"] = csv_df[district_col].astype(str).str.strip().str.upper()
    csv_df["Status"] = "present"

    ref_df = csv_df[["State", "District", "Status"]].drop_duplicates()

    # Pull existing scraped metadata to check for any states not in the CSV
    meta = _pull_metadata(token, repo)
    csv_states = set(ref_df["State"].unique())
    new_rows = []
    for state_name in meta.keys():
        state_upper = state_name.strip().upper()
        if state_upper not in csv_states:
            # State exists in scraped data but not in SF CSV → pending
            # We don't know districts yet; they'll be added when detected
            new_rows.append({"State": state_upper, "District": "ALL", "Status": "new districts found"})

    if new_rows:
        ref_df = pd.concat([ref_df, pd.DataFrame(new_rows)], ignore_index=True)

    url = push_district_reference(ref_df, token, repo)
    return ref_df, url


def update_district_reference_add(state: str, district: str, sf_status: str, token: str, repo: str) -> str:
    """Add or update a single district in Dataset 1. Returns commit URL."""
    df = pull_district_reference(token, repo)
    state = state.strip().upper()
    district = district.strip().upper()
    # Remove existing row if present
    df = df[~((df["State"] == state) & (df["District"] == district))]
    new_row = pd.DataFrame([{"State": state, "District": district, "Status": sf_status}])
    df = pd.concat([df, new_row], ignore_index=True)
    return push_district_reference(df, token, repo)


def update_district_reference_rename(state: str, old_district: str, new_district: str, token: str, repo: str) -> str:
    """
    Rename a district in Dataset 1. Also cascades the rename to Dataset 2 mapping rules.
    Returns commit URL of Dataset 1 update.
    """
    state = state.strip().upper()
    old_district = old_district.strip().upper()
    new_district = new_district.strip().upper()

    # Update Dataset 1
    df = pull_district_reference(token, repo)
    df.loc[(df["State"] == state) & (df["District"] == old_district), "District"] = new_district
    url = push_district_reference(df, token, repo)

    # Cascade to Dataset 2 — update Dist_NEW where it matches old name
    dist_rules, block_rules = pull_mapping_rules(token, repo)
    if not dist_rules.empty:
        mask = (dist_rules["State"] == state) & (dist_rules["Dist_NEW"] == old_district)
        dist_rules.loc[mask, "Dist_NEW"] = new_district
        push_mapping_rules(dist_rules, block_rules, token, repo)

    return url


def delete_district_reference(state: str, district: str, token: str, repo: str) -> str:
    """Delete a district from Dataset 1. Returns commit URL."""
    df = pull_district_reference(token, repo)
    state = state.strip().upper()
    district = district.strip().upper()
    df = df[~((df["State"] == state) & (df["District"] == district))]
    return push_district_reference(df, token, repo)


def delete_state_reference(state: str, token: str, repo: str) -> str:
    """Delete an entire state and all its districts from Dataset 1. Returns commit URL."""
    df = pull_district_reference(token, repo)
    state = state.strip().upper()
    df = df[df["State"] != state]
    return push_district_reference(df, token, repo)


def bulk_accept_flagged_districts(token: str, repo: str) -> None:
    """
    Take all currently flagged districts, add them to Dataset 1 with 'new districts found',
    and clear the flagged_districts list from metadata.
    """
    meta = _pull_metadata(token, repo)
    flags = meta.get("flagged_districts", [])
    if not flags:
        return

    # Add to Dataset 1
    df = pull_district_reference(token, repo)
    new_rows = []
    for f in flags:
        state = f["state"].strip().upper()
        dist = f["district"].strip().upper()
        if not ((df["State"] == state) & (df["District"] == dist)).any():
            new_rows.append({"State": state, "District": dist, "Status": "new districts found"})
    
    if new_rows:
        df = pd.concat([df, pd.DataFrame(new_rows)], ignore_index=True)
        push_district_reference(df, token, repo)
    
    # Clear flags
    meta["flagged_districts"] = []
    _push_metadata(meta, token, repo)


def pull_static_geo_lookup(token: str, repo: str) -> dict:
    """Download the static UDISE geo lookup JSON from HF."""
    import json
    from huggingface_hub import hf_hub_download
    try:
        path = hf_hub_download(
            repo_id=repo,
            repo_type="dataset",
            filename="district_reference/udise_geo_lookup.json",
            token=token,
            force_download=True
        )
        with open(path, "r") as f:
            return json.load(f)
    except Exception as e:
        print(f"[WARN] Failed to load static geo lookup: {e}")
        return {"state": {}, "district": {}, "block": {}}


# ─────────────────────────────────────────────────────────────────────────────
# Dataset 2: mapping_rules  (district + block splits)
# ─────────────────────────────────────────────────────────────────────────────

DIST_MAPPING_PATH  = "mapping_rules/manual_district_mapping.parquet"
BLOCK_MAPPING_PATH = "mapping_rules/manual_block_mapping.parquet"
BASELINE_MASTER_PATH = "mapping_rules/baseline_master.parquet"

def pull_baseline_master(token: str, repo: str) -> pd.DataFrame:
    """Download the baseline master dataset for UDISE tracking analysis."""
    return _download_parquet(BASELINE_MASTER_PATH, token, repo)

def push_mapping_rules(district_df: pd.DataFrame, block_df: pd.DataFrame, token: str, repo: str) -> str:
    """
    Push both district and block mapping rule tables to Dataset 2.
    district_df columns: State, Dist_OLD, Dist_NEW, Affected_Schools
    block_df columns:    State, Dist_NEW, Block_NEW, Dist_OLD, Schools
    Returns commit URL of district_mapping upload.
    """
    url = _upload_parquet(district_df, DIST_MAPPING_PATH, token, repo, "Update district_mapping rules")
    if not block_df.empty:
        _upload_parquet(block_df, BLOCK_MAPPING_PATH, token, repo, "Update block_mapping rules")
    return url


def pull_mapping_rules(token: str, repo: str) -> tuple:
    """
    Pull both mapping rule tables from Dataset 2.
    Returns (district_df, block_df) — both may be empty DataFrames if not yet seeded.
    """
    district_df = _download_parquet(DIST_MAPPING_PATH, token, repo)
    block_df    = _download_parquet(BLOCK_MAPPING_PATH, token, repo)
    return district_df, block_df


def seed_mapping_rules_from_excel(excel_path: str, token: str, repo: str) -> str:
    """
    Seed Dataset 2 from the existing district_and_block_mapping.xlsx.
    Returns commit URL.
    """
    dist_df  = pd.read_excel(excel_path, sheet_name="District Level Mapping")
    block_df = pd.read_excel(excel_path, sheet_name="Block Level Splits")
    return push_mapping_rules(dist_df, block_df, token, repo)


def upsert_mapping_rule(state: str, dist_old: str, dist_new: str,
                        affected_schools: int, token: str, repo: str) -> str:
    """Add or update a single district mapping rule in Dataset 2."""
    dist_df, block_df = pull_mapping_rules(token, repo)
    state    = state.strip().upper()
    dist_old = dist_old.strip().upper()
    dist_new = dist_new.strip().upper()
    if dist_df.empty:
        dist_df = pd.DataFrame(columns=["State", "Dist_OLD", "Dist_NEW", "Affected_Schools"])
    # Remove existing entry
    dist_df = dist_df[~((dist_df["State"] == state) & (dist_df["Dist_OLD"] == dist_old) & (dist_df["Dist_NEW"] == dist_new))]
    new_row = pd.DataFrame([{"State": state, "Dist_OLD": dist_old, "Dist_NEW": dist_new, "Affected_Schools": affected_schools}])
    dist_df = pd.concat([dist_df, new_row], ignore_index=True)
    return push_mapping_rules(dist_df, block_df, token, repo)


def delete_mapping_rule(state: str, dist_old: str, dist_new: str, token: str, repo: str) -> str:
    """Delete a district mapping rule from Dataset 2."""
    dist_df, block_df = pull_mapping_rules(token, repo)
    state    = state.strip().upper()
    dist_old = dist_old.strip().upper()
    dist_new = dist_new.strip().upper()
    dist_df = dist_df[~((dist_df["State"] == state) & (dist_df["Dist_OLD"] == dist_old) & (dist_df["Dist_NEW"] == dist_new))]
    return push_mapping_rules(dist_df, block_df, token, repo)


# ─────────────────────────────────────────────────────────────────────────────
# Flagged new districts (stored in metadata.json under "flagged_districts" key)
# ─────────────────────────────────────────────────────────────────────────────

def get_flagged_districts(token: str, repo: str) -> list:
    """Return all districts from Dataset 1 that have Status == 'new districts found', filtered by states with raw data."""
    ref_df = pull_district_reference(token, repo)
    if ref_df.empty or "Status" not in ref_df.columns:
        return []
        
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=token)
        files = api.list_repo_files(repo_id=repo, repo_type='dataset')
        active_states = set()
        for f in files:
            if f.startswith('scraped_data/raw/') and f.endswith('.parquet'):
                # Extract state name from filename (e.g. 'navy_education_society.parquet' -> 'NAVY EDUCATION SOCIETY')
                s_name = f.replace('scraped_data/raw/', '').replace('.parquet', '').replace('_', ' ').upper()
                active_states.add(s_name)
    except:
        active_states = set(ref_df["State"].str.upper()) # fallback

    pending = ref_df[(ref_df["Status"] == "new districts found") & (ref_df["State"].str.upper().isin(active_states))]
    
    dist_df, _ = pull_mapping_rules(token, repo)
    mapped = set()
    if not dist_df.empty and "Dist_OLD" in dist_df.columns:
        mapped = set(zip(
            dist_df["State"].astype(str).str.strip().str.upper(),
            dist_df["Dist_OLD"].astype(str).str.strip().str.upper()
        ))
        
    flags = []
    for _, row in pending.iterrows():
        state = str(row.get("State", "")).strip().upper()
        district = str(row.get("District", "")).strip().upper()
        if (state, district) not in mapped:
            flags.append({
                "state": state,
                "district": district,
                "reason": "New district detected"
            })
    return flags

def rename_district_in_reference(state: str, district: str, new_name: str, token: str, repo: str):
    ref_df = pull_district_reference(token, repo)
    if ref_df.empty: return
    
    # Check if new_name already exists in this state (as in_sf or pending)
    exists = not ref_df[(ref_df["State"] == state) & (ref_df["District"] == new_name)].empty
    
    # Target row to update/delete
    mask = (ref_df["State"] == state) & (ref_df["District"] == district) & (ref_df["Status"] == "new districts found")
    
    if not ref_df[mask].empty:
        if exists:
            # new_name already exists, so this is a true mapping to an old SF name. Just delete the pending row.
            ref_df = ref_df[~mask]
        else:
            # new_name doesn't exist, so this is a rename (typo fix). Keep the row as pending, just update the name.
            ref_df.loc[mask, "District"] = new_name
        
        push_district_reference(ref_df, token, repo)


# ─────────────────────────────────────────────────────────────────────────────
# Dataset 3: scraped_data  (raw + mapped master)
# ─────────────────────────────────────────────────────────────────────────────

def push_state_file(
    df: pd.DataFrame,
    state_name: str,
    categories_scraped: list,
    token: str,
    repo: str,
) -> str:
    """
    Push a raw scraped state DataFrame to scraped_data/raw/{state_key}.parquet.
    Also runs background analysis to detect new districts and flags them.
    Returns the commit URL string.
    """
    skey = _state_key(state_name)

    # ── Add scrape date column ────────────────────────────────────────────────
    df = df.copy()
    df["Scraped_Date"] = datetime.utcnow().strftime("%Y-%m-%d")

    # ── Geo-decode if non-actual state ─────────────────────────────────────────
    from district_mapper import NON_ACTUAL_STATES, apply_geo_decode
    if state_name.strip().upper() in NON_ACTUAL_STATES:
        print(f"[INFO] {state_name} is a non-actual state. Applying static UDISE geo decoding...")
        geo_lookup = pull_static_geo_lookup(token, repo)
        df, _ = apply_geo_decode(df, geo_lookup)

    # ── Upload raw parquet ────────────────────────────────────────────────────
    url = _upload_parquet(
        df,
        f"scraped_data/raw/{skey}.parquet",
        token, repo,
        f"Raw scrape: {state_name} — categories {categories_scraped}"
    )

    # ── Update metadata.json ──────────────────────────────────────────────────
    meta = _pull_metadata(token, repo)
    existing_cats = set(meta.get(state_name, {}).get("categories", []))
    merged_cats   = sorted(existing_cats | {str(c) for c in categories_scraped})
    meta[state_name] = {
        "state_key":   skey,
        "categories":  merged_cats,
        "last_scraped": datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC"),
    }
    _push_metadata(meta, token, repo)

    # ── Background: detect new districts and flag them ────────────────────────
    try:
        _detect_and_flag_new_districts(state_name, df, token, repo)
    except Exception as e:
        print(f"[WARN] New district detection failed for {state_name}: {e}")

    return url


def _detect_and_flag_new_districts(state_name: str, df: pd.DataFrame, token: str, repo: str) -> None:
    """
    Compare districts in newly scraped df against Dataset 1 and Dataset 2.
    Any district not in Dataset 1 AND not mapped in Dataset 2 is truly new.
    Add them straight to Dataset 1 with Status = 'new districts found'.
    """
    state_upper = state_name.strip().upper()
    if state_upper in NON_ACTUAL_STATES:
        return
    ref_df = pull_district_reference(token, repo)
    dist_df, _ = pull_mapping_rules(token, repo)

    # --- 1. Run Automated UDISE Tracking Math ---
    try:
        baseline_df = pull_baseline_master(token, repo)
        if not baseline_df.empty:
            # Filter baseline to just this state to speed up inner join
            baseline_df = baseline_df[baseline_df["School_State__c"].str.strip().str.upper() == state_upper]
            if not baseline_df.empty:
                from national_analysis import run_analysis
                result = run_analysis(baseline_df, df)
                
                # Push mathematically generated mapping rules
                auto_dist = result.district_mapping
                auto_block = result.block_mapping
                
                # Merge with existing manual rules
                if not auto_dist.empty or not auto_block.empty:
                    # Drop List_of_UDISEs if present to match Dataset 2 schema
                    if "List_of_UDISEs" in auto_dist.columns:
                        auto_dist = auto_dist.drop(columns=["List_of_UDISEs"])
                    if "List_of_UDISEs" in auto_block.columns:
                        auto_block = auto_block.drop(columns=["List_of_UDISEs"])
                    
                    # Merge with existing
                    current_dist, current_block = pull_mapping_rules(token, repo)
                    
                    if not current_dist.empty and not auto_dist.empty:
                        merged_dist = pd.concat([current_dist, auto_dist], ignore_index=True)
                        merged_dist = merged_dist.drop_duplicates(subset=["State", "Dist_NEW"], keep="first")
                    else:
                        merged_dist = auto_dist if not auto_dist.empty else current_dist
                        
                    if not current_block.empty and not auto_block.empty:
                        merged_block = pd.concat([current_block, auto_block], ignore_index=True)
                        merged_block = merged_block.drop_duplicates(subset=["State", "Dist_OLD", "Block_OLD", "Block_NEW"], keep="first")
                    else:
                        merged_block = auto_block if not auto_block.empty else current_block

                    push_mapping_rules(merged_dist, merged_block, token, repo)
    except Exception as e:
        print(f"[WARN] Automated UDISE Math tracking failed for {state_name}: {e}")

    # --- 2. Flag truly new districts ---
    # Re-pull rules in case they were updated by math
    dist_df, _ = pull_mapping_rules(token, repo)

    known_districts = set()
    if not ref_df.empty and "State" in ref_df.columns:
        known_districts.update(ref_df[ref_df["State"] == state_upper]["District"].unique())
        
    mapped_districts = set()
    if not dist_df.empty and "Dist_OLD" in dist_df.columns:
        mapped_districts.update(dist_df[dist_df["State"] == state_upper]["Dist_OLD"].unique())

    if "School_District__c" in df.columns:
        scraped_districts = set(df["School_District__c"].dropna().astype(str).str.strip().str.upper().unique())
    else:
        scraped_districts = set()

    # It's a new district if it's not in Salesforce reference AND not in the mapping rules
    new_districts = scraped_districts - known_districts - mapped_districts
    if not new_districts:
        return

    # Add directly to Dataset 1 as pending
    new_rows = [{"State": state_upper, "District": d, "Status": "new districts found"} for d in new_districts]
    updated_ref = pd.concat([ref_df, pd.DataFrame(new_rows)], ignore_index=True)
    push_district_reference(updated_ref, token, repo)


def delete_state_file(state_name: str, token: str, repo: str) -> None:
    """Delete a state's raw dataset file and remove it from metadata.json, and also clean up pending newly found districts for that state."""
    api = _hf()
    skey = _state_key(state_name)
    try:
        api.delete_file(
            path_in_repo=f"scraped_data/raw/{skey}.parquet",
            repo_id=repo,
            repo_type="dataset",
            token=token,
            commit_message=f"Delete {state_name} from dataset",
        )
    except Exception:
        pass  # Might not exist

    meta = _pull_metadata(token, repo)
    if state_name in meta:
        del meta[state_name]
        _push_metadata(meta, token, repo)
        
    # Also clean up any "new districts found" for this deleted state in the District Reference
    try:
        ref_df = pull_district_reference(token, repo)
        if not ref_df.empty:
            # We want to KEEP rows that are NOT (State == state_name and Status == 'new districts found')
            original_len = len(ref_df)
            state_upper = state_name.upper()
            mask = ~((ref_df["State"].str.upper() == state_upper) & (ref_df["Status"] == "new districts found"))
            cleaned_ref = ref_df[mask]
            if len(cleaned_ref) < original_len:
                push_district_reference(cleaned_ref, token, repo)
    except Exception as e:
        print(f"Failed to cleanup district reference on state delete: {e}")


def pull_state_file(state_name: str, token: str, repo: str) -> pd.DataFrame:
    """Download a single raw state Parquet from scraped_data/raw/."""
    skey = _state_key(state_name)
    return _download_parquet(f"scraped_data/raw/{skey}.parquet", token, repo)


def pull_all_complete_states(token: str, repo: str) -> tuple:
    """
    Pull all state files from HF dataset (scraped_data/raw/).

    Returns
    -------
    (complete_dfs, partial_info, coverage_df)
      complete_dfs  – list of DataFrames for states with all 7 categories
      partial_info  – dict { state_name: categories_list } for incomplete states
      coverage_df   – DataFrame showing every state's coverage (for UI table)
    """
    api  = _hf()
    meta = _pull_metadata(token, repo)

    complete_dfs  = []
    partial_info  = {}
    coverage_rows = []

    for state_name, info in meta.items():
        if not isinstance(info, dict):
            continue
        cats = set(info.get("categories", []))
        is_complete = ALL_REQUIRED_CATEGORIES.issubset(cats)

        coverage_rows.append({
            "State":              state_name,
            "Categories Scraped": ", ".join(sorted(cats, key=lambda x: int(x) if x.isdigit() else 0)),
            "Complete?":          "✅ Yes" if is_complete else "❌ No",
            "Last Scraped":       info.get("last_scraped", "—"),
        })

        skey = info.get("state_key", _state_key(state_name))
        try:
            with tempfile.TemporaryDirectory() as tmpdir:
                local_path = api.hf_hub_download(
                    repo_id=repo,
                    filename=f"scraped_data/raw/{skey}.parquet",
                    repo_type="dataset",
                    token=token,
                    local_dir=tmpdir,
                    local_dir_use_symlinks=False,
                    force_download=True,
                )
                df = pd.read_parquet(local_path)
                if is_complete:
                    complete_dfs.append(df)
                else:
                    partial_info[state_name] = sorted(cats, key=lambda x: int(x) if x.isdigit() else 0)
        except Exception as e:
            print(f"[WARN] Could not pull {state_name}: {e}")

    coverage_df = pd.DataFrame(coverage_rows).sort_values("State").reset_index(drop=True)
    return complete_dfs, partial_info, coverage_df


def push_mapped_master(df: pd.DataFrame, token: str, repo: str) -> str:
    """
    Push the final combined mapped master sheet to scraped_data/mapped/mapped_master_{date}.parquet
    Returns the commit URL string.
    """
    from datetime import timedelta
    ist_now = datetime.utcnow() + timedelta(hours=5, minutes=30)
    date_str = ist_now.strftime("%Y_%b_%d_%I_%M_%p").lower()
    path = f"scraped_data/mapped/mapped_master_{date_str}.parquet"
    return _upload_parquet(df, path, token, repo, f"Mapped master — {len(df):,} schools")


def pull_mapped_master(token: str, repo: str) -> pd.DataFrame:
    """
    Download the latest mapped master parquet from HuggingFace dataset.
    Tries the current date, then falls back to scanning for any mapped_master file.
    Returns a DataFrame, or empty DataFrame if not found.
    """
    date_str = datetime.utcnow().strftime("%Y_%d_%b").lower()
    df = _download_parquet(f"scraped_data/mapped/mapped_master_{date_str}.parquet", token, repo)
    if not df.empty:
        return df
    # Fallback: old path (backward compatibility)
    return _download_parquet("data/mapped_master.parquet", token, repo)


# ─────────────────────────────────────────────────────────────────────────────
# Metadata helpers (internal)
# ─────────────────────────────────────────────────────────────────────────────

def check_state_scraped(state_name: str, categories: list) -> bool:
    """Check if a state AND all requested categories are already present in the HF dataset metadata."""
    token, repo = get_hf_credentials()
    if not token or not repo or not state_name:
        return False
    meta = _pull_metadata(token, repo)
    if state_name not in meta:
        return False

    existing_cats = set(meta[state_name].get("categories", []))
    req_cats = {str(c).split(" - ")[0].strip() for c in (categories or [])}

    if not req_cats:
        return True

    # Return True ONLY if requested categories exactly match the existing categories
    return req_cats == existing_cats


def _pull_metadata(token: str, repo: str) -> dict:
    api = _hf()
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            local_path = api.hf_hub_download(
                repo_id=repo,
                filename="metadata.json",
                repo_type="dataset",
                token=token,
                local_dir=tmpdir,
                local_dir_use_symlinks=False,
                force_download=True,
            )
            with open(local_path, "r", encoding="utf-8") as f:
                return json.load(f)
    except Exception:
        return {}


def _push_metadata(meta: dict, token: str, repo: str) -> None:
    api = _hf()
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".json", delete=False, encoding="utf-8"
    ) as tmp:
        json.dump(meta, tmp, indent=2)
        tmp_path = tmp.name
    try:
        api.upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo="metadata.json",
            repo_id=repo,
            repo_type="dataset",
            token=token,
            commit_message="Update metadata.json",
        )
    finally:
        os.unlink(tmp_path)


# ─────────────────────────────────────────────────────────────────────────────
# test_connection
# ─────────────────────────────────────────────────────────────────────────────
def test_connection(token: str, repo: str) -> tuple:
    """
    Test that the token and repo are valid.
    Returns (success: bool, message: str)
    """
    try:
        api = _hf()
        info = api.repo_info(repo_id=repo, repo_type="dataset", token=token)
        return True, f"✅ Connected to **{info.id}**"
    except Exception as e:
        return False, f"❌ Connection failed: {e}"

def list_mapped_master_files(token: str, repo: str) -> list[str]:
    api = _hf()
    try:
        files = api.list_repo_files(repo_id=repo, repo_type="dataset", token=token)
        return [f for f in files if f.startswith("scraped_data/mapped/mapped_master_") and f.endswith(".parquet")]
    except Exception as e:
        print(f"[ERROR] Failed to list mapped master files: {e}")
        return []

def push_national_analysis_report(df: pd.DataFrame, token: str, repo: str) -> str:
    """Push the combined national analysis report of all changed districts."""
    path = "mapping_rules/changed_districts_report.parquet"
    return _upload_parquet(df, path, token, repo, "Update national analysis district changes report")