import json
import os
import argparse
import re
import pandas as pd

def _state_prefix(s: str) -> str:
    """Matches the exact filename formatting used by the main scraper."""
    return re.sub(r"[^a-z0-9]+", "_", s.lower()).strip("_")

def convert_json_to_excel(state: str):
    output_dir = "output"
    
    # Format the filename exactly like the scraper does
    prefix = _state_prefix(state)
    json_filename = f"{prefix}_schools_by_category.json"
    jfile = os.path.join(output_dir, json_filename)
    
    if not os.path.exists(jfile):
        print(f"Error: Could not find data for '{state}'.")
        print(f"Looked for: {jfile}")
        return

    all_rows = []
    
    print(f"Reading {jfile}...")
    with open(jfile, "r", encoding="utf-8") as f:
        try:
            data = json.load(f)
        except json.JSONDecodeError:
            print(f"Error: {jfile} contains invalid JSON.")
            return
        
    for record in data:
        # Only process successfully scraped combinations
        if record.get("status") != "success":
            continue
            
        resp_data = record.get("response", {}).get("data", {})
        if not resp_data:
            continue
            
        content = resp_data.get("content", [])
        for sch in content:
            
            # --- Map Management Type ---
            mgmt_id = sch.get("schBroadMgmtId", 0)
            if mgmt_id == 1:
                mgmt_type = "GOVERNMENT"
            elif mgmt_id == 2:
                mgmt_type = "GOVERNMENT AIDED"
            elif mgmt_id == 3:
                mgmt_type = "PRIVATE"
            else:
                mgmt_type = "OTHER"
            
            # --- Map Status ---
            # Using schoolStatusName strictly as requested
            final_status = str(sch.get("schoolStatusName") or "").strip()

            # Base mapping provided by user
            rename_map = {
                "schoolId": "School_Id__c",
                "udiseschCode": "School_Udise_Code__c",
                "schoolName": "School_Name__c",
                "stateName": "School_State__c",
                "districtName": "School_District__c",
                "blockName": "School_Block__c",
                "villageName": "School_Village__c",
                "schLocDesc": "School_Location__c",
                "schMgmtDescSt": "state_mgmt__c",
                "schMgmtDesc": "School_Management__c",
                "schTypeDesc": "School_Type__c",
                "schoolStatusName": "School_Status__c",
                "schCatDesc": "schCategoryType__c",
                "clusterName": "School_Cluster__c",
                "schBroadMgmtId": "School_Management_Id__c",
            }

            # Build row exactly matching your headers
            row = {}
            for json_key, excel_col in rename_map.items():
                val = sch.get(json_key)
                # Apply string formatting for text fields
                if isinstance(val, str):
                    val = val.strip()
                    if excel_col == "state_mgmt__c":
                        val = val.upper()
                elif val is None:
                    val = ""
                    
                row[excel_col] = val

            # Force strict status if it wasn't in json correctly
            if not row["School_Status__c"]:
                row["School_Status__c"] = final_status

            # Append the derived Management Type
            row["School_Management_Type__c"] = mgmt_type
            
            # Fallback for state_mgmt__c if missing from JSON but requested uppercase
            if not row["state_mgmt__c"]:
                row["state_mgmt__c"] = str(sch.get("schMgmtDesc") or "").strip().upper()

            all_rows.append(row)
            
    if not all_rows:
        print(f"No school data found inside {jfile}.")
        return

    # Convert to DataFrame and Export
    df = pd.DataFrame(all_rows)
    
    # Sort by UDISE code so that the rows line up with the old master sheet
    if "School_Udise_Code__c" in df.columns:
        df = df.sort_values(by="School_Udise_Code__c").reset_index(drop=True)
    
    excel_dir = "output_excel"
    os.makedirs(excel_dir, exist_ok=True)
    output_path = os.path.join(excel_dir, f"{prefix}_Schools.xlsx")
    
    df.to_excel(output_path, index=False)
    print(f"SUCCESS: Exported {len(df):,} schools to {output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert a State's JSON output to Excel")
    parser.add_argument("--state", type=str, required=True, help="Name of the state (e.g. 'GOA' or 'ANDHRA PRADESH')")
    args = parser.parse_args()
    
    convert_json_to_excel(args.state)