app.py

import os
import pandas as pd
import gradio as gr
from helper import assign_main_accounts, generate_schedule, save_processed_files

UPLOAD_FOLDER = "uploads"
PROCESSED_FOLDER = "processed"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(PROCESSED_FOLDER, exist_ok=True)

USERNAME = "admin"
PASSWORD = "password123"


# Helper Functions
def login(username, password):
    if username == USERNAME and password == PASSWORD:
        return "Login Successful!", True
    else:
        return "Invalid credentials. Try again.", False


def upload_files(creators_file, overnight_file, day_file, prime_file):
    try:
        # Save each file using its binary data
        with open(os.path.join(UPLOAD_FOLDER, "creators_file.xlsx"), "wb") as f:
            f.write(creators_file)

        with open(os.path.join(UPLOAD_FOLDER, "overnight_file.xlsx"), "wb") as f:
            f.write(overnight_file)

        with open(os.path.join(UPLOAD_FOLDER, "day_file.xlsx"), "wb") as f:
            f.write(day_file)

        with open(os.path.join(UPLOAD_FOLDER, "prime_file.xlsx"), "wb") as f:
            f.write(prime_file)

        return "Files uploaded successfully!"
    except Exception as e:
        return f"Error uploading files: {e}"


def generate_main_accounts(updated_data=None):
    creators_file = os.path.join(UPLOAD_FOLDER, "creators_file.xlsx")
    chatter_files = [
        os.path.join(UPLOAD_FOLDER, "overnight_file.xlsx"),
        os.path.join(UPLOAD_FOLDER, "day_file.xlsx"),
        os.path.join(UPLOAD_FOLDER, "prime_file.xlsx"),
    ]

    if not all(os.path.exists(path) for path in [creators_file] + chatter_files):
        return "Missing required files. Please upload all necessary files.", None

    try:
        if updated_data is not None:
            # Ensure updated_data is in the correct format
            if isinstance(updated_data, pd.DataFrame):
                updated_data = updated_data.to_dict(orient="records")
            
            # Validate updated data structure
            for record in updated_data:
                if not isinstance(record, dict):
                    raise ValueError(f"Invalid record in updated data: {record}")

            # Group updated data by shift
            assignments = {
                shift.lower(): [
                    {key: record[key] for key in record if key != "Shift"}  # Remove the Shift column
                    for record in updated_data if record["Shift"] == shift.capitalize()
                ]
                for shift in ["overnight", "day", "prime"]
            }

            # Save updated assignments
            save_processed_files(assignments, PROCESSED_FOLDER)
            return "Assignments updated successfully!", pd.DataFrame(updated_data)

        # Generate new assignments if no updated data is provided
        main_assignments = assign_main_accounts(creators_file, chatter_files)
        save_processed_files(main_assignments, PROCESSED_FOLDER)

        # Combine all shifts into a single DataFrame for preview
        preview_data = []
        for shift, data in main_assignments.items():
            if shift != "creator_names":  # Skip creator names
                for record in data:
                    record["Shift"] = shift.capitalize()
                    preview_data.append(record)

        preview_df = pd.DataFrame(preview_data)
        return "Main accounts generated successfully!", preview_df
    except Exception as e:
        return f"Error during main account generation: {e}", None




def generate_full_schedule():
    chatter_files = [
        os.path.join(PROCESSED_FOLDER, "Updated_overnight_file.xlsx"),
        os.path.join(PROCESSED_FOLDER, "Updated_day_file.xlsx"),
        os.path.join(PROCESSED_FOLDER, "Updated_prime_file.xlsx"),
    ]
    creators_file = os.path.join(PROCESSED_FOLDER, "creators_file.xlsx")

    if not all(os.path.exists(path) for path in chatter_files + [creators_file]):
        return "Files are missing. Please generate main accounts first."

    try:
        full_schedules = generate_schedule(chatter_files, creators_file)
        return full_schedules
    except Exception as e:
        return f"Error generating full schedule: {e}"


# Gradio Interface
def app():
    with gr.Blocks() as interface:
        with gr.Tab("Login"):
            username = gr.Textbox(label="Username")
            password = gr.Textbox(label="Password", type="password")
            login_btn = gr.Button("Login")
            login_status = gr.Textbox(label="Login Status", interactive=False)

            login_btn.click(login, inputs=[username, password], outputs=[login_status])

        with gr.Tab("Upload Files"):
            creators_file = gr.File(label="Creators File", type="binary")
            overnight_file = gr.File(label="Overnight File", type="binary")
            day_file = gr.File(label="Day File", type="binary")
            prime_file = gr.File(label="Prime File", type="binary")
            upload_btn = gr.Button("Upload Files")
            upload_status = gr.Textbox(label="Upload Status", interactive=False)

            upload_btn.click(
                upload_files,
                inputs=[creators_file, overnight_file, day_file, prime_file],
                outputs=[upload_status],
            )

        with gr.Tab("Generate Main Accounts"):
            generate_main_btn = gr.Button("Generate Main Accounts")
            generate_main_status = gr.Textbox(label="Status", interactive=False)
            preview_table = gr.DataFrame(label="Main Account Assignments Preview", interactive=True)
            update_btn = gr.Button("Update Assignments")

            # Generate assignments and preview
            generate_main_btn.click(
                generate_main_accounts,
                inputs=[],
                outputs=[generate_main_status, preview_table],
            )

            # Update assignments based on supervisor input
            update_btn.click(
                generate_main_accounts,
                inputs=[preview_table],
                outputs=[generate_main_status, preview_table],
            )

        with gr.Tab("Generate Full Schedule"):
            generate_full_btn = gr.Button("Generate Full Schedule")
            full_schedule_output = gr.Textbox(label="Full Schedule", interactive=False)

            generate_full_btn.click(generate_full_schedule, outputs=[full_schedule_output])

    return interface




# Launch Gradio App
if __name__ == "__main__":
    app().launch()