hf_reader.py

"""
hf_reader.py — HuggingFace data layer for the School Name Resolver.

Reads:
  1. master_all_states.csv  (Old Master) from HF_RESOLVER_REPO
  2. scraped_data/mapped/*.parquet  (new masters) from HF_SCRAPER_REPO

All data is cached in memory after first load. Call refresh_all_caches() to force reload.
"""

import os
import re
import pandas as pd
from huggingface_hub import HfApi, hf_hub_download
from datetime import datetime
from rapidfuzz import process, fuzz
try:
    from admin_patterns import normalize_with_patterns_dynamic
except ImportError:
    normalize_with_patterns_dynamic = lambda s, st: s

HF_TOKEN = os.getenv("HF_TOKEN", "")
HF_SCRAPER_REPO = os.getenv("HF_SCRAPER_REPO", "")

# ─── Module-level caches ───────────────────────────────────────────────────────
_old_master_df: pd.DataFrame | None = None
_mapped_masters: dict[str, pd.DataFrame] = {}
_cache_loaded: bool = False


# ─── Internal helpers ──────────────────────────────────────────────────────────

def _load_old_master() -> pd.DataFrame:
    """Download and return baseline_master.parquet from HF_SCRAPER_REPO."""
    if not HF_SCRAPER_REPO:
        print("[hf_reader] HF_SCRAPER_REPO not set — skipping Old Master load.")
        return pd.DataFrame()
    try:
        path = hf_hub_download(
            repo_id=HF_SCRAPER_REPO,
            filename="mapping_rules/baseline_master.parquet",
            repo_type="dataset",
            token=HF_TOKEN or None,
            force_download=True,
        )
        df = pd.read_parquet(path)
        if "School_Udise_Code__c" in df.columns:
            df["School_Udise_Code__c"] = df["School_Udise_Code__c"].astype(str).str.strip()
        print(f"[hf_reader] Old Master loaded: {len(df):,} rows.")
        return df
    except Exception as e:
        print(f"[hf_reader] Could not load baseline_master.parquet: {e}")
        return pd.DataFrame()


def _load_mapped_masters() -> dict[str, pd.DataFrame]:
    """Download all parquet files from scraped_data/mapped/ in HF_SCRAPER_REPO."""
    if not HF_SCRAPER_REPO:
        print("[hf_reader] HF_SCRAPER_REPO not set — skipping mapped masters load.")
        return {}

    api = HfApi()
    results: dict[str, pd.DataFrame] = {}

    try:
        all_files = list(api.list_repo_files(
            repo_id=HF_SCRAPER_REPO,
            repo_type="dataset",
            token=HF_TOKEN or None,
        ))
        mapped_files = sorted(
            [f for f in all_files if f.startswith("scraped_data/mapped/") and f.endswith(".parquet")],
            reverse=True,
        )
        print(f"[hf_reader] Found {len(mapped_files)} mapped master(s) in scraper repo.")

        for file_path in mapped_files:
            try:
                local_path = hf_hub_download(
                    repo_id=HF_SCRAPER_REPO,
                    filename=file_path,
                    repo_type="dataset",
                    token=HF_TOKEN or None,
                    force_download=True,
                )
                df = pd.read_parquet(local_path)
                if "School_Udise_Code__c" in df.columns:
                    df["School_Udise_Code__c"] = df["School_Udise_Code__c"].astype(str).str.strip()
                stem = file_path.split("/")[-1].replace(".parquet", "")
                results[stem] = df
                print(f"[hf_reader]   Loaded: {stem}  ({len(df):,} rows)")
            except Exception as e:
                print(f"[hf_reader]   Failed to load {file_path}: {e}")

    except Exception as e:
        print(f"[hf_reader] Could not list scraper repo files: {e}")

    return results


def _ensure_loaded():
    """Load all caches if they haven't been loaded yet."""
    global _old_master_df, _mapped_masters, _cache_loaded
    if not _cache_loaded:
        print("[hf_reader] Initializing caches... This may take a few minutes if downloading from HuggingFace.")
        _old_master_df = _load_old_master()
        print("[hf_reader] Old Master loaded.")
        _mapped_masters = _load_mapped_masters()
        print(f"[hf_reader] Mapped Masters loaded: {len(_mapped_masters)} files.")
        _cache_loaded = True
        print("[hf_reader] All caches fully loaded and ready.")


def get_state_hierarchy() -> dict:
    """
    Returns a nested dictionary of the geographic hierarchy:
    { state: { district: { block: [villages] } } }
    Built by combining all loaded data sources.
    """
    _ensure_loaded()
    hier = {}
    
    dfs_to_process = []
    if _old_master_df is not None and not _old_master_df.empty:
        dfs_to_process.append(_old_master_df)
    dfs_to_process.extend(_mapped_masters.values())
    
    for df in dfs_to_process:
        if df.empty: continue
        
        s_col = "School_State__c"
        d_col = "School_District__c"
        b_col = "School_Block__c"
        v_col = "School_Village__c"
        
        if s_col not in df.columns: continue
        
        df_tmp = df.copy()
        for col in [s_col, d_col, b_col, v_col]:
            if col not in df_tmp.columns:
                df_tmp[col] = ""
            else:
                df_tmp[col] = df_tmp[col].fillna("").astype(str).str.strip().str.upper()
                
        for _, r in df_tmp.iterrows():
            st = r[s_col]
            di = r[d_col]
            bl = r[b_col]
            vi = r[v_col]
            
            if not st: continue
            if st not in hier: hier[st] = {}
            if not di: continue
            if di not in hier[st]: hier[st][di] = {}
            if not bl: continue
            if bl not in hier[st][di]: hier[st][di][bl] = set()
            if vi:
                hier[st][di][bl].add(vi)
                
    for s in hier:
        for d in hier[s]:
            for b in hier[s][d]:
                hier[s][d][b] = sorted(list(hier[s][d][b]))
                
    return hier


def search_by_name_fuzzy(query_name: str, state: str, district: str, block: str, village: str = None, max_results=10) -> list[str]:
    """
    Fuzzy searches the combined data sources for a school name.
    Returns a list of unique matched UDISE codes.
    """
    _ensure_loaded()
    if not query_name:
        return []
        
    dfs_to_process = []
    if _old_master_df is not None and not _old_master_df.empty:
        dfs_to_process.append(_old_master_df)
    dfs_to_process.extend(_mapped_masters.values())
    
    if not dfs_to_process:
        return []
        
    combined_rows = []
    
    for df in dfs_to_process:
        if df.empty: continue
        if "School_Name__c" not in df.columns or "School_Udise_Code__c" not in df.columns: continue
        
        mask = pd.Series(True, index=df.index)
        
        if state:
            if "School_State__c" in df.columns:
                mask = mask & (df["School_State__c"].astype(str).str.strip().str.upper() == state.upper())
            else:
                continue
        
        if district:
            if "School_District__c" in df.columns:
                mask = mask & (df["School_District__c"].astype(str).str.strip().str.upper() == district.upper())
            else:
                continue
                
        if block:
            if "School_Block__c" in df.columns:
                mask = mask & (df["School_Block__c"].astype(str).str.strip().str.upper() == block.upper())
            else:
                continue
                
        if village:
            if "School_Village__c" in df.columns:
                mask = mask & (df["School_Village__c"].astype(str).str.strip().str.upper() == village.upper())
            else:
                continue
                
        filtered = df[mask]
        if not filtered.empty:
            filtered_sub = filtered[["School_Udise_Code__c", "School_Name__c"]].copy()
            combined_rows.append(filtered_sub)
            
    if not combined_rows:
        return []
        
    combined_df = pd.concat(combined_rows, ignore_index=True).drop_duplicates()
    combined_df["School_Name__c"] = combined_df["School_Name__c"].astype(str).str.strip()
    combined_df = combined_df[combined_df["School_Name__c"] != ""]
    
    if combined_df.empty:
        return []
        
    choices = combined_df["School_Name__c"].tolist()
    state_for_patterns = (state or "ARUNACHAL PRADESH").upper()
    
    candidates_raw = process.extract(
        query_name,
        choices,
        scorer=fuzz.token_set_ratio,
        processor=lambda s: normalize_with_patterns_dynamic(s, state_for_patterns),
        limit=max_results,
    )
    
    matched_udises = []
    seen_udises = set()
    for choice, score, idx in candidates_raw:
        udise = combined_df.iloc[idx]["School_Udise_Code__c"]
        udise = str(udise).strip()
        if udise and udise not in seen_udises:
            seen_udises.add(udise)
            matched_udises.append(udise)
            
    return matched_udises


# ─── Public API ───────────────────────────────────────────────────────────────

def refresh_all_caches() -> str:
    """Force-reload all data from HuggingFace. Returns a status message."""
    global _old_master_df, _mapped_masters, _cache_loaded
    _cache_loaded = False
    _old_master_df = None
    _mapped_masters = {}
    _ensure_loaded()
    n_masters = len(_mapped_masters)
    sf_rows = len(_old_master_df) if _old_master_df is not None else 0
    return (
        f"✅ Caches refreshed — Old Master: {sf_rows:,} rows, "
        f"Mapped Masters: {n_masters} file(s) loaded."
    )


def _pretty_label(stem: str, idx: int) -> tuple[str, str]:
    """Turn a filename stem like 'mapped_master_2026_jul_01_10_19_pm' into a readable label and year_month."""
    match_str = re.search(r"(\d{4})_([a-zA-Z]{3})_(\d{2})", stem)
    if match_str:
        y, m_str, d = match_str.groups()
        date_str = f"{d}-{m_str.capitalize()}-{y}"
        ym = f"{y}-{m_str.capitalize()}"
    else:
        match_num = re.search(r"(\d{4})(\d{2})(\d{2})", stem)
        if match_num:
            y, m, d = match_num.groups()
            date_str = f"{d}/{m}/{y}"
            month_str = datetime.strptime(m, "%m").strftime("%b").capitalize()
            ym = f"{y}-{month_str}"
        else:
            date_str = stem
            ym = ""

    if idx == 0:
        return f"Latest Scraped Master ({date_str})", ym
    else:
        return f"Older Scraped Master ({date_str})", ym


def search_udise(udise_code: str, marksheet_name: str = "") -> list[dict]:
    """
    Search for a UDISE code across all data sources.
    Returns a list of dicts.
    """
    _ensure_loaded()
    results: list[dict] = []

    if marksheet_name and marksheet_name.strip():
        results.append({
            "source": "Marksheet (entered by you)",
            "source_type": "marksheet",
            "name": marksheet_name.strip(),
            "state": "",
            "district": "",
            "block": "",
            "year_month": datetime.now().strftime("%Y-%b").capitalize(),
        })

    if _old_master_df is not None and not _old_master_df.empty:
        mask = _old_master_df["School_Udise_Code__c"] == str(udise_code).strip()
        matches = _old_master_df[mask]
        for _, row in matches.iterrows():
            results.append({
                "source": "Old Master (baseline_master.parquet)",
                "source_type": "legacy_db",
                "name": str(row.get("School_Name__c", "")).strip(),
                "state": str(row.get("School_State__c", "")).strip(),
                "district": str(row.get("School_District__c", "")).strip(),
                "block": str(row.get("School_Block__c", "")).strip(),
                "village": str(row.get("School_Village__c", "")).strip(),
                "year_month": "2025",
            })

    for idx, (stem, df) in enumerate(_mapped_masters.items()):
        if "School_Udise_Code__c" not in df.columns:
            continue
        mask = df["School_Udise_Code__c"] == str(udise_code).strip()
        matches = df[mask]
        label, ym = _pretty_label(stem, idx)
        source_type = "latest_master" if idx == 0 else "old_master"
        for _, row in matches.iterrows():
            results.append({
                "source": label,
                "source_type": source_type,
                "name": str(row.get("School_Name__c", "")).strip(),
                "state": str(row.get("School_State__c", "")).strip(),
                "district": str(row.get("School_District__c", "")).strip(),
                "block": str(row.get("School_Block__c", "")).strip(),
                "village": str(row.get("School_Village__c", "")).strip(),
                "year_month": ym,
            })

    return results


def get_udise_choices(state: str, district: str, block: str, village: str = None) -> list[str]:
    """Returns a list of 'UDISE - School Name' for the given location filter."""
    if not state:
        return []
        
    _ensure_loaded()
    
    dfs_to_process = []
    if _old_master_df is not None and not _old_master_df.empty:
        dfs_to_process.append(_old_master_df)
    dfs_to_process.extend(_mapped_masters.values())
    
    if not dfs_to_process:
        return []
        
    combined_rows = []
    for df in dfs_to_process:
        if df.empty: continue
        if "School_Name__c" not in df.columns or "School_Udise_Code__c" not in df.columns: continue
        
        mask = pd.Series(True, index=df.index)
        if state:
            if "School_State__c" in df.columns:
                mask &= (df["School_State__c"].astype(str).str.strip().str.upper() == state.upper())
            else: continue
        if district:
            if "School_District__c" in df.columns:
                mask &= (df["School_District__c"].astype(str).str.strip().str.upper() == district.upper())
            else: continue
        if block:
            if "School_Block__c" in df.columns:
                mask &= (df["School_Block__c"].astype(str).str.strip().str.upper() == block.upper())
            else: continue
        if village:
            if "School_Village__c" in df.columns:
                mask &= (df["School_Village__c"].astype(str).str.strip().str.upper() == village.upper())
            else: continue
                
        filtered = df[mask]
        if not filtered.empty:
            filtered_sub = filtered[["School_Udise_Code__c", "School_Name__c"]].copy()
            combined_rows.append(filtered_sub)
            
    if not combined_rows:
        return []
        
    combined_df = pd.concat(combined_rows, ignore_index=True).drop_duplicates()
    combined_df["School_Udise_Code__c"] = combined_df["School_Udise_Code__c"].astype(str).str.strip()
    combined_df["School_Name__c"] = combined_df["School_Name__c"].astype(str).str.strip()
    combined_df = combined_df[(combined_df["School_Udise_Code__c"] != "") & (combined_df["School_Name__c"] != "")]
    
    if combined_df.empty:
        return []
        
    choices = (combined_df["School_Udise_Code__c"] + " - " + combined_df["School_Name__c"]).unique().tolist()
    return sorted(choices)


def get_config_status() -> dict:
    """Return a dict describing the current configuration state."""
    _ensure_loaded()
    return {
        "scraper_repo": HF_SCRAPER_REPO or "⚠️ Not set (HF_SCRAPER_REPO)",
        "sf_baseline_rows": len(_old_master_df) if _old_master_df is not None else 0,
        "mapped_masters_count": len(_mapped_masters),
        "token_set": bool(HF_TOKEN),
    }