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