Update app.py

app.py

@@ -41,6 +41,8 @@ def upload_files(creators_file, overnight_file, day_file, prime_file):
 
 
 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"),
@@ -56,14 +58,9 @@ def generate_main_accounts(updated_data=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 = {
+
+            # 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()
@@ -71,12 +68,16 @@ def generate_main_accounts(updated_data=None):
                 for shift in ["overnight", "day", "prime"]
             }
 
-            # Save updated assignments
-            save_processed_files(assignments, PROCESSED_FOLDER)
+            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)
 
         # Generate new assignments if no updated data is provided
         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
@@ -95,24 +96,47 @@ def generate_main_accounts(updated_data=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."
+# 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: