import gradio as gr
import pandas as pd
import json
import os
import csv
import io
from datetime import datetime

# --- Persistent Storage ---
DATA_FILE = "gradebook_data.json"

def load_data():
    """Load gradebook data from JSON file."""
    if os.path.exists(DATA_FILE):
        try:
            with open(DATA_FILE, "r") as f:
                return json.load(f)
        except (json.JSONDecodeError, IOError):
            pass
    return {"pupils": [], "assignments": [], "grades": {}}

def save_data(data):
    """Save gradebook data to JSON file."""
    with open(DATA_FILE, "w") as f:
        json.dump(data, f, indent=2)

# --- Helper Functions ---

def get_grade_table():
    """Build the main grades dataframe."""
    data = load_data()
    pupils = data["pupils"]
    assignments = data["assignments"]
    grades = data["grades"]

    if not pupils:
        return pd.DataFrame({"Info": ["No pupils added yet. Go to 'Manage Pupils' to add some."]})

    if not assignments:
        return pd.DataFrame({"Pupil": pupils, "Average": ["N/A"] * len(pupils)})

    rows = []
    for pupil in pupils:
        row = {"Pupil": pupil}
        total = 0
        count = 0
        for assignment in assignments:
            key = f"{pupil}|||{assignment}"
            grade = grades.get(key, "")
            row[assignment] = grade
            if grade != "":
                try:
                    total += float(grade)
                    count += 1
                except (ValueError, TypeError):
                    pass
        row["Average"] = f"{total / count:.1f}" if count > 0 else "N/A"
        rows.append(row)
    
    df = pd.DataFrame(rows)
    # Reorder columns: Pupil first, then assignments, then Average
    cols = ["Pupil"] + assignments + ["Average"]
    df = df[cols]
    return df

def get_styled_table():
    """Return a styled grade table with color-coded averages."""
    df = get_grade_table()
    if "Average" not in df.columns or "Pupil" not in df.columns:
        return df
    
    def color_grades(val):
        try:
            v = float(val)
            if v >= 90:
                return "background-color: #c6efce; color: #006100"
            elif v >= 80:
                return "background-color: #d4edbc; color: #2e6b09"
            elif v >= 70:
                return "background-color: #ffeb9c; color: #9c6500"
            elif v >= 60:
                return "background-color: #ffc7ce; color: #9c0006"
            else:
                return "background-color: #ff9999; color: #800000"
        except (ValueError, TypeError):
            return ""

    # Apply styling to grade columns (not Pupil)
    grade_cols = [c for c in df.columns if c != "Pupil"]
    styler = df.style.map(color_grades, subset=grade_cols)
    styler = styler.format(precision=1, na_rep="")
    return styler

def get_summary_stats():
    """Generate summary statistics."""
    data = load_data()
    pupils = data["pupils"]
    assignments = data["assignments"]
    grades = data["grades"]

    if not pupils or not assignments:
        return "No data available yet. Add pupils and assignments first."

    lines = []
    lines.append(f"📊 **Gradebook Summary**")
    lines.append(f"- **Total Pupils:** {len(pupils)}")
    lines.append(f"- **Total Assignments:** {len(assignments)}")
    lines.append("")

    # Per-assignment stats
    lines.append("### 📝 Assignment Statistics")
    lines.append("| Assignment | Mean | Min | Max | Graded |")
    lines.append("|---|---|---|---|---|")
    
    for assignment in assignments:
        scores = []
        for pupil in pupils:
            key = f"{pupil}|||{assignment}"
            grade = grades.get(key, "")
            if grade != "":
                try:
                    scores.append(float(grade))
                except (ValueError, TypeError):
                    pass
        if scores:
            mean = sum(scores) / len(scores)
            lines.append(f"| {assignment} | {mean:.1f} | {min(scores):.1f} | {max(scores):.1f} | {len(scores)}/{len(pupils)} |")
        else:
            lines.append(f"| {assignment} | — | — | — | 0/{len(pupils)} |")
    
    lines.append("")

    # Per-pupil averages
    lines.append("### 🎓 Pupil Averages (ranked)")
    pupil_avgs = []
    for pupil in pupils:
        scores = []
        for assignment in assignments:
            key = f"{pupil}|||{assignment}"
            grade = grades.get(key, "")
            if grade != "":
                try:
                    scores.append(float(grade))
                except (ValueError, TypeError):
                    pass
        avg = sum(scores) / len(scores) if scores else None
        pupil_avgs.append((pupil, avg, len(scores)))
    
    pupil_avgs.sort(key=lambda x: x[1] if x[1] is not None else -1, reverse=True)
    
    lines.append("| Rank | Pupil | Average | Assignments Graded |")
    lines.append("|---|---|---|---|")
    for i, (pupil, avg, graded) in enumerate(pupil_avgs, 1):
        avg_str = f"{avg:.1f}" if avg is not None else "N/A"
        medal = "🥇" if i == 1 else "🥈" if i == 2 else "🥉" if i == 3 else f"{i}."
        lines.append(f"| {medal} | {pupil} | {avg_str} | {graded}/{len(assignments)} |")

    # Class average
    all_scores = []
    for pupil in pupils:
        for assignment in assignments:
            key = f"{pupil}|||{assignment}"
            grade = grades.get(key, "")
            if grade != "":
                try:
                    all_scores.append(float(grade))
                except (ValueError, TypeError):
                    pass
    
    if all_scores:
        lines.append("")
        lines.append(f"### 📈 Class Overview")
        lines.append(f"- **Class Average:** {sum(all_scores)/len(all_scores):.1f}")
        lines.append(f"- **Highest Grade:** {max(all_scores):.1f}")
        lines.append(f"- **Lowest Grade:** {min(all_scores):.1f}")
        lines.append(f"- **Total Grades Entered:** {len(all_scores)} / {len(pupils) * len(assignments)}")

    return "\n".join(lines)

# --- Event Handlers ---

def add_pupil(name):
    """Add a new pupil."""
    name = name.strip()
    if not name:
        return "⚠️ Please enter a name.", get_pupils_list(), get_styled_table(), gr.update(), get_summary_stats()
    
    data = load_data()
    if name in data["pupils"]:
        return f"⚠️ '{name}' already exists.", get_pupils_list(), get_styled_table(), gr.update(), get_summary_stats()
    
    data["pupils"].append(name)
    data["pupils"].sort()
    save_data(data)
    
    pupil_choices = data["pupils"]
    return f"✅ Added '{name}'.", get_pupils_list(), get_styled_table(), gr.update(choices=pupil_choices), get_summary_stats()

def remove_pupil(name):
    """Remove a pupil and their grades."""
    name = name.strip()
    data = load_data()
    
    if name not in data["pupils"]:
        return f"⚠️ '{name}' not found.", get_pupils_list(), get_styled_table(), gr.update(), get_summary_stats()
    
    data["pupils"].remove(name)
    # Remove associated grades
    keys_to_remove = [k for k in data["grades"] if k.startswith(f"{name}|||")]
    for k in keys_to_remove:
        del data["grades"][k]
    save_data(data)
    
    pupil_choices = data["pupils"]
    return f"✅ Removed '{name}'.", get_pupils_list(), get_styled_table(), gr.update(choices=pupil_choices), get_summary_stats()

def get_pupils_list():
    """Get formatted list of pupils."""
    data = load_data()
    if not data["pupils"]:
        return "No pupils added yet."
    return "\n".join([f"• {p}" for p in data["pupils"]])

def add_assignment(name):
    """Add a new assignment."""
    name = name.strip()
    if not name:
        return "⚠️ Please enter an assignment name.", get_assignments_list(), get_styled_table(), gr.update(), gr.update(), get_summary_stats()
    
    data = load_data()
    if name in data["assignments"]:
        return f"⚠️ '{name}' already exists.", get_assignments_list(), get_styled_table(), gr.update(), gr.update(), get_summary_stats()
    
    data["assignments"].append(name)
    save_data(data)
    
    assignment_choices = data["assignments"]
    return f"✅ Added assignment '{name}'.", get_assignments_list(), get_styled_table(), gr.update(choices=assignment_choices), gr.update(choices=assignment_choices), get_summary_stats()

def remove_assignment(name):
    """Remove an assignment and its grades."""
    name = name.strip()
    data = load_data()
    
    if name not in data["assignments"]:
        return f"⚠️ '{name}' not found.", get_assignments_list(), get_styled_table(), gr.update(), gr.update(), get_summary_stats()
    
    data["assignments"].remove(name)
    # Remove associated grades
    keys_to_remove = [k for k in data["grades"] if k.endswith(f"|||{name}")]
    for k in keys_to_remove:
        del data["grades"][k]
    save_data(data)
    
    assignment_choices = data["assignments"]
    return f"✅ Removed assignment '{name}'.", get_assignments_list(), get_styled_table(), gr.update(choices=assignment_choices), gr.update(choices=assignment_choices), get_summary_stats()

def get_assignments_list():
    """Get formatted list of assignments."""
    data = load_data()
    if not data["assignments"]:
        return "No assignments added yet."
    return "\n".join([f"• {a}" for a in data["assignments"]])

def enter_grade(pupil, assignment, grade):
    """Enter or update a grade for a pupil/assignment."""
    data = load_data()
    
    if not pupil or not assignment:
        return "⚠️ Please select both a pupil and an assignment.", get_styled_table(), get_summary_stats()
    
    grade = grade.strip()
    if grade == "":
        # Allow clearing a grade
        key = f"{pupil}|||{assignment}"
        if key in data["grades"]:
            del data["grades"][key]
            save_data(data)
            return f"✅ Cleared grade for {pupil} on '{assignment}'.", get_styled_table(), get_summary_stats()
        return "ℹ️ No grade to clear.", get_styled_table(), get_summary_stats()
    
    try:
        grade_val = float(grade)
        if grade_val < 0 or grade_val > 100:
            return "⚠️ Grade must be between 0 and 100.", get_styled_table(), get_summary_stats()
    except ValueError:
        return "⚠️ Please enter a valid number (0-100).", get_styled_table(), get_summary_stats()
    
    key = f"{pupil}|||{assignment}"
    data["grades"][key] = grade_val
    save_data(data)
    
    return f"✅ {pupil} → {assignment}: {grade_val}", get_styled_table(), get_summary_stats()

def batch_enter_grades(assignment, grades_text):
    """Enter grades for multiple pupils at once for a given assignment."""
    data = load_data()
    
    if not assignment:
        return "⚠️ Please select an assignment.", get_styled_table(), get_summary_stats()
    
    if not grades_text.strip():
        return "⚠️ Please enter grades.", get_styled_table(), get_summary_stats()
    
    lines = grades_text.strip().split("\n")
    success = 0
    errors = []
    
    for line in lines:
        line = line.strip()
        if not line:
            continue
        # Support formats: "Name: Grade", "Name, Grade", "Name Grade"
        parts = None
        for sep in [":", ",", "\t"]:
            if sep in line:
                parts = line.split(sep, 1)
                break
        if parts is None:
            parts = line.rsplit(" ", 1)
        
        if len(parts) != 2:
            errors.append(f"Could not parse: '{line}'")
            continue
        
        pupil_name = parts[0].strip()
        grade_str = parts[1].strip()
        
        if pupil_name not in data["pupils"]:
            errors.append(f"Pupil not found: '{pupil_name}'")
            continue
        
        try:
            grade_val = float(grade_str)
            if grade_val < 0 or grade_val > 100:
                errors.append(f"Grade out of range for '{pupil_name}': {grade_val}")
                continue
        except ValueError:
            errors.append(f"Invalid grade for '{pupil_name}': '{grade_str}'")
            continue
        
        key = f"{pupil_name}|||{assignment}"
        data["grades"][key] = grade_val
        success += 1
    
    save_data(data)
    
    msg = f"✅ Entered {success} grade(s)."
    if errors:
        msg += f"\n⚠️ {len(errors)} error(s):\n" + "\n".join(f"  • {e}" for e in errors)
    
    return msg, get_styled_table(), get_summary_stats()

def export_csv():
    """Export gradebook as CSV."""
    df = get_grade_table()
    if "Info" in df.columns:
        return None
    
    filepath = "gradebook_export.csv"
    df.to_csv(filepath, index=False)
    return filepath

def import_csv(file):
    """Import grades from a CSV file."""
    if file is None:
        return "⚠️ No file uploaded.", get_styled_table(), get_pupils_list(), get_assignments_list(), get_summary_stats()
    
    try:
        df = pd.read_csv(file.name if hasattr(file, 'name') else file)
    except Exception as e:
        return f"⚠️ Error reading CSV: {e}", get_styled_table(), get_pupils_list(), get_assignments_list(), get_summary_stats()
    
    if "Pupil" not in df.columns:
        return "⚠️ CSV must have a 'Pupil' column.", get_styled_table(), get_pupils_list(), get_assignments_list(), get_summary_stats()
    
    data = load_data()
    
    # Get assignment columns (everything except Pupil and Average)
    assignment_cols = [c for c in df.columns if c not in ("Pupil", "Average")]
    
    # Add new pupils and assignments
    for pupil in df["Pupil"]:
        pupil = str(pupil).strip()
        if pupil and pupil not in data["pupils"]:
            data["pupils"].append(pupil)
    data["pupils"].sort()
    
    for assignment in assignment_cols:
        assignment = str(assignment).strip()
        if assignment and assignment not in data["assignments"]:
            data["assignments"].append(assignment)
    
    # Import grades
    count = 0
    for _, row in df.iterrows():
        pupil = str(row["Pupil"]).strip()
        for assignment in assignment_cols:
            val = row[assignment]
            if pd.notna(val) and str(val).strip() != "":
                try:
                    grade_val = float(val)
                    key = f"{pupil}|||{assignment}"
                    data["grades"][key] = grade_val
                    count += 1
                except (ValueError, TypeError):
                    pass
    
    save_data(data)
    return f"✅ Imported {count} grades from CSV.", get_styled_table(), get_pupils_list(), get_assignments_list(), get_summary_stats()

def refresh_pupil_dropdown():
    data = load_data()
    return gr.update(choices=data["pupils"])

def refresh_assignment_dropdown():
    data = load_data()
    return gr.update(choices=data["assignments"])

def clear_all_data():
    """Reset the entire gradebook."""
    save_data({"pupils": [], "assignments": [], "grades": {}})
    return (
        "✅ All data cleared.",
        get_styled_table(),
        get_pupils_list(),
        get_assignments_list(),
        get_summary_stats(),
        gr.update(choices=[]),
        gr.update(choices=[]),
        gr.update(choices=[]),
    )

# --- Build the UI ---

css = """
.gradio-container { max-width: 1200px !important; }
h1 { text-align: center; margin-bottom: 0.5em; }
.subtitle { text-align: center; color: #666; margin-bottom: 1.5em; }
"""

with gr.Blocks(css=css, title="📚 Gradebook", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📚 Gradebook")
    gr.Markdown("<p class='subtitle'>Manage your pupils' grades in one place</p>")
    
    with gr.Tabs():
        # ===================== TAB 1: Grades Overview =====================
        with gr.Tab("📊 Grades", id="grades"):
            grades_table = gr.Dataframe(
                value=get_styled_table(),
                label="Grades Table",
                interactive=False,
                wrap=True,
                max_height=600,
                show_search="filter",
            )
            with gr.Row():
                refresh_btn = gr.Button("🔄 Refresh", variant="secondary", scale=1)
                export_btn = gr.Button("📥 Export CSV", variant="secondary", scale=1)
            export_file = gr.File(label="Download Export", visible=False)
            
            refresh_btn.click(fn=get_styled_table, outputs=[grades_table])
            
            def do_export():
                path = export_csv()
                if path:
                    return gr.update(value=path, visible=True)
                return gr.update(visible=False)
            
            export_btn.click(fn=do_export, outputs=[export_file])
        
        # ===================== TAB 2: Enter Grades =====================
        with gr.Tab("✏️ Enter Grades", id="enter"):
            gr.Markdown("### Enter a Single Grade")
            with gr.Row():
                grade_pupil = gr.Dropdown(
                    choices=load_data()["pupils"],
                    label="Pupil",
                    interactive=True,
                    scale=2,
                )
                grade_assignment = gr.Dropdown(
                    choices=load_data()["assignments"],
                    label="Assignment",
                    interactive=True,
                    scale=2,
                )
                grade_value = gr.Textbox(
                    label="Grade (0-100)",
                    placeholder="e.g. 85",
                    scale=1,
                )
            grade_btn = gr.Button("💾 Save Grade", variant="primary")
            grade_msg = gr.Markdown("")
            
            gr.Markdown("---")
            gr.Markdown("### Batch Enter Grades")
            gr.Markdown("Enter one grade per line in the format: `Pupil Name: Grade`  \n"
                        "Also supports comma or tab separation.")
            with gr.Row():
                batch_assignment = gr.Dropdown(
                    choices=load_data()["assignments"],
                    label="Assignment",
                    interactive=True,
                    scale=1,
                )
                batch_text = gr.Textbox(
                    label="Grades",
                    placeholder="Alice: 92\nBob: 85\nCharlie: 78",
                    lines=8,
                    scale=2,
                )
            batch_btn = gr.Button("💾 Save All Grades", variant="primary")
            batch_msg = gr.Markdown("")
            
            # Hidden table for updates
            grades_table_hidden = gr.Dataframe(visible=False)
            summary_hidden = gr.Markdown(visible=False)
        
        # ===================== TAB 3: Manage Pupils =====================
        with gr.Tab("👩‍🎓 Manage Pupils", id="pupils"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Add Pupil")
                    pupil_name_input = gr.Textbox(label="Pupil Name", placeholder="e.g. Alice Johnson")
                    add_pupil_btn = gr.Button("➕ Add Pupil", variant="primary")
                    
                    gr.Markdown("### Remove Pupil")
                    remove_pupil_input = gr.Textbox(label="Pupil Name to Remove", placeholder="Exact name")
                    remove_pupil_btn = gr.Button("🗑️ Remove Pupil", variant="stop")
                    
                    pupil_msg = gr.Markdown("")
                
                with gr.Column(scale=1):
                    gr.Markdown("### Current Pupils")
                    pupils_list_display = gr.Markdown(get_pupils_list())
        
        # ===================== TAB 4: Manage Assignments =====================
        with gr.Tab("📝 Manage Assignments", id="assignments"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Add Assignment")
                    assignment_name_input = gr.Textbox(label="Assignment Name", placeholder="e.g. Midterm Exam")
                    add_assignment_btn = gr.Button("➕ Add Assignment", variant="primary")
                    
                    gr.Markdown("### Remove Assignment")
                    remove_assignment_input = gr.Textbox(label="Assignment Name to Remove", placeholder="Exact name")
                    remove_assignment_btn = gr.Button("🗑️ Remove Assignment", variant="stop")
                    
                    assignment_msg = gr.Markdown("")
                
                with gr.Column(scale=1):
                    gr.Markdown("### Current Assignments")
                    assignments_list_display = gr.Markdown(get_assignments_list())
        
        # ===================== TAB 5: Statistics =====================
        with gr.Tab("📈 Statistics", id="stats"):
            stats_display = gr.Markdown(get_summary_stats())
            stats_refresh_btn = gr.Button("🔄 Refresh Statistics", variant="secondary")
            stats_refresh_btn.click(fn=get_summary_stats, outputs=[stats_display])
        
        # ===================== TAB 6: Import / Settings =====================
        with gr.Tab("⚙️ Settings", id="settings"):
            gr.Markdown("### Import Data from CSV")
            gr.Markdown("Upload a CSV file with a `Pupil` column and assignment columns. Example:\n\n"
                        "| Pupil | Homework 1 | Quiz 1 | Midterm |\n"
                        "|---|---|---|---|\n"
                        "| Alice | 95 | 88 | 92 |\n"
                        "| Bob | 82 | 76 | 85 |")
            import_file = gr.File(label="Upload CSV", file_types=[".csv"])
            import_btn = gr.Button("📤 Import", variant="primary")
            import_msg = gr.Markdown("")
            
            gr.Markdown("---")
            gr.Markdown("### ⚠️ Danger Zone")
            clear_btn = gr.Button("🗑️ Clear All Data", variant="stop")
            clear_msg = gr.Markdown("")
    
    # --- Wire up events ---
    
    # Add pupil
    add_pupil_btn.click(
        fn=add_pupil,
        inputs=[pupil_name_input],
        outputs=[pupil_msg, pupils_list_display, grades_table, grade_pupil, stats_display],
    )
    
    # Remove pupil
    remove_pupil_btn.click(
        fn=remove_pupil,
        inputs=[remove_pupil_input],
        outputs=[pupil_msg, pupils_list_display, grades_table, grade_pupil, stats_display],
    )
    
    # Add assignment
    add_assignment_btn.click(
        fn=add_assignment,
        inputs=[assignment_name_input],
        outputs=[assignment_msg, assignments_list_display, grades_table, grade_assignment, batch_assignment, stats_display],
    )
    
    # Remove assignment
    remove_assignment_btn.click(
        fn=remove_assignment,
        inputs=[remove_assignment_input],
        outputs=[assignment_msg, assignments_list_display, grades_table, grade_assignment, batch_assignment, stats_display],
    )
    
    # Enter single grade
    grade_btn.click(
        fn=enter_grade,
        inputs=[grade_pupil, grade_assignment, grade_value],
        outputs=[grade_msg, grades_table, stats_display],
    )
    
    # Batch enter grades
    batch_btn.click(
        fn=batch_enter_grades,
        inputs=[batch_assignment, batch_text],
        outputs=[batch_msg, grades_table, stats_display],
    )
    
    # Import CSV
    import_btn.click(
        fn=import_csv,
        inputs=[import_file],
        outputs=[import_msg, grades_table, pupils_list_display, assignments_list_display, stats_display],
    )
    
    # Clear all data
    clear_btn.click(
        fn=clear_all_data,
        outputs=[clear_msg, grades_table, pupils_list_display, assignments_list_display, stats_display, grade_pupil, grade_assignment, batch_assignment],
    )

demo.launch()