alias_store.py

"""
alias_store.py — Manages the school_aliases.csv file stored on HuggingFace.
"""

import os
import tempfile
import pandas as pd
from datetime import datetime, timezone
from huggingface_hub import HfApi, hf_hub_download

HF_TOKEN = os.getenv("HF_TOKEN", "")
HF_SCRAPER_REPO = os.getenv("HF_SCRAPER_REPO", "")
ALIAS_FILE = "school_aliases.csv"

def load_aliases() -> dict:
    """
    Download the CSV from HF, parse it, and return a dictionary identical to 
    the old JSON format so that app.py doesn't have to change at all!
    """
    if not HF_SCRAPER_REPO:
        return {}
    try:
        path = hf_hub_download(
            repo_id=HF_SCRAPER_REPO,
            filename=ALIAS_FILE,
            repo_type="dataset",
            token=HF_TOKEN or None,
            force_download=True,
        )
        df = pd.read_csv(path)
        
        data = {}
        for _, row in df.iterrows():
            code = str(row["UDISE_Code"]).strip()
            name = str(row["Alias_Name"]).strip()
            source = str(row["Source"]).strip()
            
            # Handle empty values gracefully
            ym = str(row["Year_Month"]) if pd.notna(row["Year_Month"]) else ""
            if ym == "nan": ym = ""
            
            lu = str(row["Last_Updated"]) if pd.notna(row["Last_Updated"]) else ""
            if lu == "nan": lu = ""
            
            if code not in data:
                data[code] = {"names": [], "last_updated": lu}
            
            data[code]["names"].append({
                "name": name,
                "source": source,
                "year_month": ym
            })
        return data
    except Exception as e:
        if "404" in str(e) or "not found" in str(e).lower() or "Entry Not Found" in str(e):
            return {}
        print(f"[alias_store] Error loading CSV aliases: {e}")
        return {}

def save_aliases(new_entries: list[dict]) -> str:
    """
    Upsert new alias entries into the cloud CSV.
    """
    if not HF_SCRAPER_REPO:
        return "⚠️ HF_SCRAPER_REPO is not configured — cannot save aliases."
    if not new_entries:
        return "⚠️ No entries provided."

    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    
    try:
        path = hf_hub_download(
            repo_id=HF_SCRAPER_REPO,
            filename=ALIAS_FILE,
            repo_type="dataset",
            token=HF_TOKEN or None,
            force_download=True,
        )
        df = pd.read_csv(path)
    except Exception:
        # File might not exist yet
        df = pd.DataFrame(columns=["UDISE_Code", "Alias_Name", "Source", "Year_Month", "Last_Updated"])

    added_count = 0
    skipped_count = 0
    new_rows = []

    for entry in new_entries:
        code = str(entry.get("udise_code", "")).strip()
        name = str(entry.get("alias_name", "")).strip()
        source = str(entry.get("source_label", "")).strip()
        ym = str(entry.get("year_month", "")).strip()
        
        if not code or not name:
            continue

        # Check if alias already exists (case-insensitive) for this UDISE
        if not df.empty:
            mask = (df["UDISE_Code"].astype(str).str.strip() == code) & (df["Alias_Name"].astype(str).str.strip().str.upper() == name.upper())
            if mask.any():
                skipped_count += 1
                continue

        new_rows.append({
            "UDISE_Code": code,
            "Alias_Name": name,
            "Source": source,
            "Year_Month": ym,
            "Last_Updated": today
        })
        added_count += 1

    if added_count == 0:
        return f"ℹ️ All {skipped_count} name(s) already in CSV — nothing new to save."

    # Force string type on UDISE_Code to prevent pandas from breaking when sorting mixed types!
    df["UDISE_Code"] = df["UDISE_Code"].astype(str)
    
    df = pd.concat([df, pd.DataFrame(new_rows)], ignore_index=True)
    
    # Sort by UDISE so all aliases for the same school are grouped together in Excel
    df = df.sort_values(by=["UDISE_Code", "Source"]).reset_index(drop=True)

    # Upload to HF
    try:
        api = HfApi()
        with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False, encoding="utf-8") as f:
            df.to_csv(f.name, index=False)
            tmp_path = f.name

        udise_sample = new_entries[0].get("udise_code", "unknown")
        api.upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo=ALIAS_FILE,
            repo_id=HF_SCRAPER_REPO,
            repo_type="dataset",
            token=HF_TOKEN or None,
            commit_message=f"Add aliases for UDISE {udise_sample} ({added_count} new name(s))",
        )

        msg = f"✅ Saved! {added_count} new name(s) added."
        if skipped_count:
            msg += f" ({skipped_count} duplicate(s) skipped.)"
        return msg

    except Exception as e:
        return f"❌ Failed to save aliases: {e}"


def get_names_for_udise(udise_code: str) -> list[dict]:
    """Return all known names for a specific UDISE code as a list of dicts."""
    data = load_aliases()
    return data.get(str(udise_code).strip(), {}).get("names", [])


def delete_alias(udise_code: str, alias_name: str) -> str:
    """Delete a specific alias name for a UDISE code from the cloud CSV."""
    if not HF_SCRAPER_REPO:
        return "⚠️ HF_SCRAPER_REPO is not configured — cannot delete aliases."
        
    code = str(udise_code).strip()
    name_to_delete = str(alias_name).strip().upper()
    
    if not code or not name_to_delete:
        return "⚠️ Invalid UDISE code or name."
        
    try:
        path = hf_hub_download(
            repo_id=HF_SCRAPER_REPO,
            filename=ALIAS_FILE,
            repo_type="dataset",
            token=HF_TOKEN or None,
            force_download=True,
        )
        df = pd.read_csv(path)
    except Exception:
        return "⚠️ CSV file not found on HuggingFace."
        
    mask = (df["UDISE_Code"].astype(str).str.strip() == code) & (df["Alias_Name"].astype(str).str.strip().str.upper() == name_to_delete)
    
    if not mask.any():
        return f"⚠️ Alias '{alias_name}' not found for UDISE {code}."
        
    # Keep everything EXCEPT the matching row
    df = df[~mask]
    
    try:
        api = HfApi()
        with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False, encoding="utf-8") as f:
            df.to_csv(f.name, index=False)
            tmp_path = f.name
            
        api.upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo=ALIAS_FILE,
            repo_id=HF_SCRAPER_REPO,
            repo_type="dataset",
            token=HF_TOKEN or None,
            commit_message=f"Delete alias '{alias_name}' for UDISE {code}",
        )
        return f"🗑️ Deleted '{alias_name}' successfully!"
    except Exception as e:
        return f"❌ Failed to delete alias: {e}"