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 the 'data' attribute
        with open(os.path.join(UPLOAD_FOLDER, "creators_file.xlsx"), "wb") as f:
            f.write(creators_file.read())

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

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

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

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


def generate_main_accounts():
    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."

    try:
        main_assignments = assign_main_accounts(creators_file, chatter_files)
        save_processed_files(main_assignments, PROCESSED_FOLDER)
        return "Main accounts generated successfully!"
    except Exception as e:
        return f"Error during main account generation: {e}"


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")
            overnight_file = gr.File(label="Overnight File")
            day_file = gr.File(label="Day File")
            prime_file = gr.File(label="Prime File")
            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)

            generate_main_btn.click(generate_main_accounts, outputs=[generate_main_status])

        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()