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):
    global UPDATED_ASSIGNMENTS

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

            # Group updated data by shift and save to UPDATED_ASSIGNMENTS
            UPDATED_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"]
            }

            UPDATED_ASSIGNMENTS["creator_names"] = list(
                set(record["Main Account"] for record in updated_data if "Main Account" in record)
            )

            save_processed_files(UPDATED_ASSIGNMENTS, PROCESSED_FOLDER)
            return "Assignments updated successfully!", pd.DataFrame(updated_data)

        # Debugging: Verify creators file before processing
        print("DEBUG: Original Creators File Columns")
        creators_df = pd.read_excel(creators_file)
        print(creators_df.columns)

        # Ensure creators file has correct columns
        column_mapping = {
            "Creator": "Creator",
            "Total earnings": "Total earnings",
            "Subscription": "Subscription",
            "Active Fans": "ActiveFans",
            "Total active fans": "ActiveFans",
        }
        creators_df.rename(columns={k: v for k, v in column_mapping.items() if k in creators_df.columns}, inplace=True)

        # Debugging: Verify creators file after renaming
        print("DEBUG: Renamed Creators File Columns")
        print(creators_df.columns)

        required_columns = ["Creator", "ActiveFans"]
        missing_columns = [col for col in required_columns if col not in creators_df.columns]
        if missing_columns:
            raise KeyError(f"Missing required columns in creators file: {missing_columns}")

        # Save the processed creators file
        creators_df.to_excel(os.path.join(PROCESSED_FOLDER, "creators_file.xlsx"), index=False)
        main_assignments = assign_main_accounts(creators_file, chatter_files)
        UPDATED_ASSIGNMENTS = main_assignments
        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







# Global variable to store updated assignments from the main account generation
UPDATED_ASSIGNMENTS = {}

def generate_full_schedule():
    global UPDATED_ASSIGNMENTS

    # Check if UPDATED_ASSIGNMENTS is populated
    if UPDATED_ASSIGNMENTS:
        # Convert the in-memory assignments to DataFrames
        chatter_files = []
        for shift in ["overnight", "day", "prime"]:
            if shift in UPDATED_ASSIGNMENTS:
                df = pd.DataFrame(UPDATED_ASSIGNMENTS[shift])
                chatter_files.append(df)

        creators_file = pd.DataFrame({"Creator": UPDATED_ASSIGNMENTS.get("creator_names", [])})
    else:
        # Fall back to reading processed files
        chatter_files = [
            pd.read_excel(os.path.join(PROCESSED_FOLDER, "Updated_overnight_file.xlsx")),
            pd.read_excel(os.path.join(PROCESSED_FOLDER, "Updated_day_file.xlsx")),
            pd.read_excel(os.path.join(PROCESSED_FOLDER, "Updated_prime_file.xlsx")),
        ]
        creators_file = pd.read_excel(os.path.join(PROCESSED_FOLDER, "creators_file.xlsx"))

    # Debugging: Ensure chatter files and creators file are populated
    print("DEBUG: Chatter Files and Creators File for Schedule Generation")
    for chatter_file in chatter_files:
        print(chatter_file.head())
    print(creators_file.head())

    try:
        # Generate schedules using in-memory or processed data
        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()