Update app.py

app.py

@@ -425,12 +425,17 @@ def process_gemini_output(output):
 
 
 def generate_documentation_page():
+    import time
     # Sidebar with "Log Out" and "Back to Project" buttons
     st.sidebar.image("SimplifAI Logo Long.jpeg", use_container_width=True)
     st.sidebar.title(f"Project: {st.session_state.current_project}")
     if st.sidebar.button("Back to Project"):
         st.session_state.page = "project_view"
+        # Clear the Gemini output when backing out
+        if "generated_documentation" in st.session_state:
+            del st.session_state["generated_documentation"]
         st.rerun()
+
     if st.sidebar.button("Log Out"):
         st.session_state.authenticated = False
         st.session_state.username = None
@@ -463,19 +468,19 @@ def generate_documentation_page():
                     # Read file contents
                     file_contents = read_files(file_paths)
 
-                    # Calculate the total lines of code
-                    total_lines = sum(content.count("\n") for content in file_contents.values())
-                    
-                    # Estimate time remaining
-                    lines_per_minute = 2500
-                    estimated_time = total_lines / lines_per_minute
-                    st.info(f"Estimated time remaining: {int(estimated_time)} minute(s)")
+                    # Calculate estimated time based on lines of code
+                    total_lines = sum(len(content.splitlines()) for content in file_contents.values())
+                    estimated_time = (total_lines // 2500) * 60  # 1 minute per 2500 lines
+                    estimated_time = max(estimated_time, 1)  # Minimum 1 second for short tasks
 
-                    # Simulate real-time progress (optional for UX)
+                    # Show progress bar
                     progress_bar = st.progress(0)
-                    for i in range(int(estimated_time * 10)):  # Simulate 10 steps per minute
-                        time.sleep(6)  # Update every 6 seconds
-                        progress_bar.progress((i + 1) / (estimated_time * 10))
+                    progress_message = st.empty()
+                    for i in range(101):
+                        time.sleep(estimated_time / 100)  # Simulate the process
+                        progress_bar.progress(i)
+                        progress_message.write(f"Estimated time remaining: {int(estimated_time * (100 - i) / 100)} seconds")
+                    progress_message.empty()
 
                     # Generate documentation using Gemini
                     documentation = generate_detailed_documentation(file_contents, functionality)
@@ -486,6 +491,10 @@ def generate_documentation_page():
                     # Display the final documentation
                     st.success("Documentation generated successfully!")
                     st.text_area("Generated Documentation", documentation, height=600)
+
+                    # Remove the progress bar once the output is ready
+                    progress_bar.empty()
+
                 except Exception as e:
                     st.error(f"An error occurred: {e}")
             else:
@@ -505,49 +514,37 @@ def generate_documentation_page():
         markdown_path = os.path.join(user_folder, f"{st.session_state.current_project}_Documentation.md")
 
         # Show appropriate buttons based on file generation status
-        pdf_generated = os.path.exists(pdf_path)
-        markdown_generated = os.path.exists(markdown_path)
-
-        # Generate and replace buttons dynamically
-        if not pdf_generated:
+        if not os.path.exists(pdf_path):
             if st.button("Generate PDF"):
                 try:
                     generate_pdf(documentation, pdf_path)
-                    st.session_state.pdf_generated = True  # Set a flag for PDF generation
-                    st.rerun()  # Trigger a rerun to show the download button
+                    st.success("PDF generated successfully!")
                 except Exception as e:
                     st.error(f"Failed to generate PDF: {e}")
         else:
-            try:
-                with open(pdf_path, "rb") as pdf_file:
-                    st.download_button(
-                        label="Download PDF",
-                        data=pdf_file.read(),
-                        file_name=f"{st.session_state.current_project}_Documentation.pdf",
-                        mime="application/pdf",
-                    )
-            except FileNotFoundError:
-                st.write("Click 'Generate PDF' to be able to download and store the documentation as a PDF")
-
-        if not markdown_generated:
+            with open(pdf_path, "rb") as pdf_file:
+                st.download_button(
+                    label="Download PDF",
+                    data=pdf_file.read(),
+                    file_name=f"{st.session_state.current_project}_Documentation.pdf",
+                    mime="application/pdf",
+                )
+
+        if not os.path.exists(markdown_path):
             if st.button("Generate Markdown File"):
                 try:
                     generate_markdown_file(documentation, markdown_path)
-                    st.session_state.markdown_generated = True  # Set a flag for Markdown generation
-                    st.rerun()  # Trigger a rerun to show the download button
+                    st.success("Markdown file generated successfully!")
                 except Exception as e:
                     st.error(f"Failed to generate Markdown file: {e}")
         else:
-            try:
-                with open(markdown_path, "rb") as markdown_file:
-                    st.download_button(
-                        label="Download Markdown File",
-                        data=markdown_file.read(),
-                        file_name=f"{st.session_state.current_project}_Documentation.md",
-                        mime="text/markdown",
-                    )
-            except FileNotFoundError:
-                st.write("Click 'Generate Markdown File' to be able to download and store the documentation as a markdown file")
+            with open(markdown_path, "rb") as markdown_file:
+                st.download_button(
+                    label="Download Markdown File",
+                    data=markdown_file.read(),
+                    file_name=f"{st.session_state.current_project}_Documentation.md",
+                    mime="text/markdown",
+                )