Update app.py

app.py

@@ -40,7 +40,7 @@ def upload_files(creators_file, overnight_file, day_file, prime_file):
         return f"Error uploading files: {e}"
 
 
-def generate_main_accounts():
+def generate_main_accounts(updated_data=None):
     creators_file = os.path.join(UPLOAD_FOLDER, "creators_file.xlsx")
     chatter_files = [
         os.path.join(UPLOAD_FOLDER, "overnight_file.xlsx"),
@@ -49,14 +49,35 @@ def generate_main_accounts():
     ]
 
     if not all(os.path.exists(path) for path in [creators_file] + chatter_files):
-        return "Missing required files. Please upload all necessary files."
+        return "Missing required files. Please upload all necessary files.", None
 
     try:
+        if updated_data is not None:
+            # If updated data is provided, save it back as the new assignments
+            assignments = {
+                shift: pd.DataFrame([record for record in updated_data if record["Shift"] == shift.capitalize()])
+                for shift in ["overnight", "day", "prime"]
+            }
+            save_processed_files(assignments, PROCESSED_FOLDER)
+            return "Assignments updated successfully!", pd.concat(assignments.values())
+
+        # Generate assignments for preview
         main_assignments = assign_main_accounts(creators_file, chatter_files)
         save_processed_files(main_assignments, PROCESSED_FOLDER)
-        return "Main accounts generated successfully!"
+
+        # 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}"
+        return f"Error during main account generation: {e}", None
+
 
 
 def generate_full_schedule():
@@ -105,8 +126,22 @@ def app():
         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],
+            )
 
-            generate_main_btn.click(generate_main_accounts, outputs=[generate_main_status])
+            # 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")
@@ -117,6 +152,7 @@ def app():
     return interface
 
 
+
 # Launch Gradio App
 if __name__ == "__main__":
     app().launch()