app.py

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

# Global Constants
UPLOAD_FOLDER = "uploads"
PROCESSED_FOLDER = "processed"
ZIP_FILE = "processed_files.zip"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(PROCESSED_FOLDER, exist_ok=True)

# Create an empty ZIP file if it doesn't exist
if not os.path.exists(ZIP_FILE):
    with zipfile.ZipFile(ZIP_FILE, "w") as zipf:
        pass

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:
        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():
    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, None

    try:
        updated_chatter_files, processed_creator_file, combined_assignments = assign_main_accounts(
            creators_file, chatter_files
        )

        # Add the `Shift` column to the combined assignments
        shifts = ["Overnight", "Day", "Prime"]
        for idx, chatter_df in enumerate(updated_chatter_files):
            chatter_df["Shift"] = shifts[idx]
            updated_chatter_files[idx] = chatter_df

        combined_assignments = pd.concat(updated_chatter_files)

        # Save updated files
        for idx, chatter_df in enumerate(updated_chatter_files):
            chatter_df.to_excel(
                os.path.join(PROCESSED_FOLDER, f"Updated_{shifts[idx].lower()}_file.xlsx"),
                index=False
            )
        processed_creator_file.to_excel(
            os.path.join(PROCESSED_FOLDER, "creators_file.xlsx"), index=False
        )

        # Create a zip file containing all processed files
        with zipfile.ZipFile(ZIP_FILE, "w") as zipf:
            for root, _, files in os.walk(PROCESSED_FOLDER):
                for file in files:
                    zipf.write(os.path.join(root, file), arcname=file)

        return "Main accounts generated successfully!", combined_assignments, processed_creator_file
    except Exception as e:
        return f"Error during main account generation: {e}", None, None

def update_assignments(assignments_df, creators_df):
    try:
        print("DEBUG: Assignments DataFrame Columns:", assignments_df.columns)

        shifts = ["Overnight", "Day", "Prime"]
        for shift in shifts:
            shift_data = assignments_df[assignments_df["Shift"] == shift]
            shift_file_path = os.path.join(PROCESSED_FOLDER, f"Updated_{shift.lower()}_file.xlsx")
            shift_data.to_excel(shift_file_path, index=False)

        creators_file_path = os.path.join(PROCESSED_FOLDER, "creators_file.xlsx")
        creators_df.to_excel(creators_file_path, index=False)

        # Update the ZIP file with the latest files
        with zipfile.ZipFile(ZIP_FILE, "w") as zipf:
            for root, _, files in os.walk(PROCESSED_FOLDER):
                for file in files:
                    zipf.write(os.path.join(root, file), arcname=file)

        return "Assignments and Creator File updated successfully!", ZIP_FILE
    except Exception as e:
        return f"Error updating assignments: {e}", None
        
def update_and_download_schedule(overnight_df, day_df, prime_df):
    try:
        # Save updated schedules to the processed folder
        overnight_path = os.path.join(PROCESSED_FOLDER, "Updated_Overnight_Schedule.xlsx")
        day_path = os.path.join(PROCESSED_FOLDER, "Updated_Day_Schedule.xlsx")
        prime_path = os.path.join(PROCESSED_FOLDER, "Updated_Prime_Schedule.xlsx")
        
        overnight_df.to_excel(overnight_path, index=False)
        day_df.to_excel(day_path, index=False)
        prime_df.to_excel(prime_path, index=False)

        # Zip the updated files for download
        with zipfile.ZipFile(ZIP_FILE, "w") as zipf:
            zipf.write(overnight_path, arcname="Updated_Overnight_Schedule.xlsx")
            zipf.write(day_path, arcname="Updated_Day_Schedule.xlsx")
            zipf.write(prime_path, arcname="Updated_Prime_Schedule.xlsx")

        return "Schedules updated and saved successfully!", ZIP_FILE
    except Exception as e:
        return f"Error updating schedules: {e}", None

def generate_full_schedule(creators_file_path, overnight_file_path, day_file_path, prime_file_path):
    try:
        # Log file paths for debugging
        print(f"DEBUG: Creators File Path: {creators_file_path}")
        print(f"DEBUG: Overnight File Path: {overnight_file_path}")
        print(f"DEBUG: Day File Path: {day_file_path}")
        print(f"DEBUG: Prime File Path: {prime_file_path}")

        # Read files
        creators_data = pd.read_excel(creators_file_path)
        overnight_data = pd.read_excel(overnight_file_path)
        day_data = pd.read_excel(day_file_path)
        prime_data = pd.read_excel(prime_file_path)

        print("DEBUG: Creators File Columns:", creators_data.columns)

        # Validate structure
        if not {"Creator", "ActiveFans"}.issubset(creators_data.columns):
            raise KeyError("The account data must contain 'Creator' and 'ActiveFans' columns.")

        # Pass data to generate_schedule
        chatter_files = [overnight_data, day_data, prime_data]
        full_schedule = generate_schedule(chatter_files, creators_data)

        # Extract individual schedules
        overnight_schedule = full_schedule["Overnight"]
        day_schedule = full_schedule["Day"]
        prime_schedule = full_schedule["Prime"]

        print("DEBUG: Overnight Schedule:", overnight_schedule.head())
        print("DEBUG: Day Schedule:", day_schedule.head())
        print("DEBUG: Prime Schedule:", prime_schedule.head())

        return overnight_schedule, day_schedule, prime_schedule
    except Exception as e:
        print(f"DEBUG: Error in generate_full_schedule: {e}")
        return f"Error generating schedule: {e}", None, None


# 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_status = gr.Textbox(label="Status", interactive=False)
            assignments_preview = gr.Dataframe(label="Main Account Assignments Preview", interactive=True)
            creators_preview = gr.Dataframe(label="Processed Creator File Preview", interactive=True)

            generate_main_btn = gr.Button("Generate Main Accounts")
            update_btn = gr.Button("Update and Download Files")
            download_zip_btn = gr.File(label="Download Processed Files", value=ZIP_FILE)

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

            update_btn.click(
                update_assignments,
                inputs=[assignments_preview, creators_preview],
                outputs=[generate_main_status, download_zip_btn],
            )

        with gr.Tab("Generate Full Schedule"):
            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")

            generate_schedule_btn = gr.Button("Generate Full Schedule")
            overnight_schedule_output = gr.Dataframe(label="Overnight Schedule", interactive=True)
            day_schedule_output = gr.Dataframe(label="Day Schedule", interactive=True)
            prime_schedule_output = gr.Dataframe(label="Prime Schedule", interactive=True)

            update_schedule_btn = gr.Button("Update and Download Schedule")
            download_schedule_btn = gr.File(label="Download Updated Schedules", value=ZIP_FILE)
            update_status = gr.Textbox(label="Update Status", interactive=False)

            # Generate Schedule
            generate_schedule_btn.click(
                generate_full_schedule,
                inputs=[creators_file, overnight_file, day_file, prime_file],
                outputs=[overnight_schedule_output, day_schedule_output, prime_schedule_output]
            )

            # Update and Download Schedule
            update_schedule_btn.click(
                update_and_download_schedule,
                inputs=[overnight_schedule_output, day_schedule_output, prime_schedule_output],
                outputs=[update_status, download_schedule_btn]
            )

    return interface

if __name__ == "__main__":
    app().launch()