import gradio as gr
import pandas as pd
import tempfile
import os

# ---------------------------------------------------
# Main Processing Function
# ---------------------------------------------------
def group_students(file_path):

    try:
        if file_path is None:
            return pd.DataFrame({"Status": ["Upload a file first"]}), None

        # Read CSV or Excel
        if file_path.lower().endswith((".xlsx", ".xls")):
            df = pd.read_excel(file_path)
        else:
            df = pd.read_csv(file_path)

        # Clean column names
        df.columns = df.columns.str.strip().str.lower()

        # Auto detect columns
        name_col = None
        dob_col = None

        for col in df.columns:
            if "name" in col:
                name_col = col
            if "dob" in col or "birth" in col or "date" in col:
                dob_col = col

        if name_col is None or dob_col is None:
            return pd.DataFrame({
                "Error": ["Name or Date of Birth column not found"]
            }), None

        # Convert DOB
        df[dob_col] = pd.to_datetime(df[dob_col], errors="coerce", dayfirst=True)
        df = df.dropna(subset=[dob_col])

        # Extract month
        df["Birth_Month"] = df[dob_col].dt.month_name()

        grouped = (
            df.groupby("Birth_Month")[name_col]
            .apply(lambda x: ", ".join(x.astype(str)))
            .reset_index()
        )

        # Sort calendar order
        month_order = [
            "January","February","March","April","May","June",
            "July","August","September","October","November","December"
        ]

        grouped["Birth_Month"] = pd.Categorical(
            grouped["Birth_Month"],
            categories=month_order,
            ordered=True
        )

        grouped = grouped.sort_values("Birth_Month")

        # Create downloadable CSV
        output_path = os.path.join(
            tempfile.gettempdir(),
            "birth_month_grouped.csv"
        )

        grouped.to_csv(output_path, index=False)

        return grouped, output_path

    except Exception as e:
        # show error directly in table
        return pd.DataFrame({"Error": [str(e)]}), None


# ---------------------------------------------------
# Gradio UI
# ---------------------------------------------------
with gr.Blocks() as demo:

    gr.Markdown("""
    # 🎂 Student Birth Month Grouping

    Upload CSV or Excel file containing:
    - Student Name
    - Date of Birth
    """)

    file_input = gr.File(
        label="Upload File",
        type="filepath"
    )

    run_btn = gr.Button("Generate List")

    table_output = gr.Dataframe(label="Grouped Students")

    download_output = gr.File(label="Download CSV")

    run_btn.click(
        group_students,
        inputs=file_input,
        outputs=[table_output, download_output]
    )

demo.launch()