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

def sync_aliases(token: str, repo: str):
    """
    Downloads all master parquets and the current school_aliases.csv,
    syncs all multi-name schools, and re-uploads the CSV.
    """
    print("Starting automatic CSV alias sync...")
    api = HfApi(token=token)
    
    # 1. Get all files
    files = api.list_repo_files(repo_id=repo, repo_type="dataset")
    
    # 2. Download Old Master
    print("Downloading Old Master...")
    old_master_df = pd.DataFrame()
    try:
        if "data/baseline_master.parquet" in files:
            path = hf_hub_download(repo_id=repo, filename="data/baseline_master.parquet", repo_type="dataset", token=token, force_download=True)
            old_master_df = pd.read_parquet(path)
    except Exception as e:
        print(f"Error loading old master: {e}")

    # 3. Download all mapped masters
    mapped_files = sorted([f for f in files if f.startswith("scraped_data/mapped/mapped_master_") and f.endswith(".parquet")], reverse=True)
    
    mapped_dfs = {}
    for f in mapped_files:
        stem = f.split("/")[-1].replace(".parquet", "")
        print(f"Downloading {stem}...")
        try:
            path = hf_hub_download(repo_id=repo, filename=f, repo_type="dataset", token=token, force_download=True)
            mapped_dfs[stem] = pd.read_parquet(path)
        except Exception as e:
            print(f"Error loading {f}: {e}")

    # Helper for pretty label
    def _pretty_label(stem: str, index: int) -> tuple[str, str]:
        parts = stem.split("_")
        if len(parts) >= 6:
            year, month = parts[2], parts[3]
            ym = f"{year}-{month.capitalize()}"
            if index == 0:
                return f"Latest Scraped Master ({month.capitalize()}-{year})", ym
            return f"Past Scraped Master ({month.capitalize()}-{year})", ym
        return "Scraped Master", ""

    # Build UDISE dict
    udise_aliases = {}
    
    # Process old master
    if not old_master_df.empty:
        for _, row in old_master_df.iterrows():
            udise = str(row.get("School_Udise_Code__c", "")).strip()
            name = str(row.get("School_Name__c", "")).strip()
            if udise and name and udise != "nan" and name != "nan":
                if udise not in udise_aliases:
                    udise_aliases[udise] = {}
                udise_aliases[udise][name.upper()] = {
                    "name": name,
                    "source": "Old Master (baseline_master.parquet)",
                    "year_month": "2025"
                }
                
    # Process mapped masters
    for idx, (stem, df) in enumerate(mapped_dfs.items()):
        label, ym = _pretty_label(stem, idx)
        for _, row in df.iterrows():
            udise = str(row.get("School_Udise_Code__c", "")).strip()
            name = str(row.get("School_Name__c", "")).strip()
            if udise and name and udise != "nan" and name != "nan":
                if udise not in udise_aliases:
                    udise_aliases[udise] = {}
                if name.upper() not in udise_aliases[udise]:
                    udise_aliases[udise][name.upper()] = {
                        "name": name,
                        "source": label,
                        "year_month": ym
                    }
                    
    # Load existing CSV to preserve manual edits
    print("Loading existing CSV aliases...")
    try:
        path = hf_hub_download(repo_id=repo, filename="school_aliases.csv", repo_type="dataset", token=token, force_download=True)
        existing_csv = pd.read_csv(path)
        for _, row in existing_csv.iterrows():
            udise = str(row.get("UDISE_Code", "")).strip()
            name = str(row.get("Alias_Name", "")).strip()
            if udise and name and udise != "nan" and name != "nan":
                if udise not in udise_aliases:
                    udise_aliases[udise] = {}
                udise_aliases[udise][name.upper()] = {
                    "name": name,
                    "source": str(row.get("Source", "")),
                    "year_month": str(row.get("Year_Month", "")) if pd.notna(row.get("Year_Month")) else "",
                    "last_updated": str(row.get("Last_Updated", "")) if pd.notna(row.get("Last_Updated")) else ""
                }
    except Exception as e:
        print(f"Could not load existing CSV: {e}")

    # Extract multi-name schools
    rows = []
    for udise, names_dict in udise_aliases.items():
        if len(names_dict) > 1:
            for upper_name, info in names_dict.items():
                rows.append({
                    "UDISE_Code": udise,
                    "Alias_Name": info["name"],
                    "Source": info["source"],
                    "Year_Month": info["year_month"],
                    "Last_Updated": info.get("last_updated", "")
                })

    if not rows:
        print("No multi-name schools found.")
        return

    df = pd.DataFrame(rows, columns=["UDISE_Code", "Alias_Name", "Source", "Year_Month", "Last_Updated"])
    
    # Force string type to prevent sorting errors
    df["UDISE_Code"] = df["UDISE_Code"].astype(str)
    
    # Sort by UDISE code so that all identical schools are grouped together perfectly
    df = df.sort_values(by=["UDISE_Code", "Source"]).reset_index(drop=True)

    print(f"Uploading {len(df)} synced aliases to HuggingFace...")
    try:
        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="school_aliases.csv",
            repo_id=repo,
            repo_type="dataset",
            token=token,
            commit_message=f"Auto-sync aliases after building mapped master ({len(df)} aliases)"
        )
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        print("CSV Sync complete!")
    except Exception as e:
        print(f"Failed to upload CSV: {e}")