Update app.py

app.py

@@ -542,6 +542,7 @@ def main():
         st.session_state.discovered_files = []
         st.session_state.current_url = None
         st.session_state.google_creds = None
+        st.session_state.selected_files = []
 
     st.title("Advanced File Downloader")
     
@@ -590,8 +591,9 @@ def main():
                             custom_ext_list=custom_extensions.split(',') if custom_extensions else [],
                             sublink_limit=max_sublinks
                         )
-                        st.session_state.discovered_files = files
-                        st.session_state.current_url = url
+                        if files:
+                            st.session_state.discovered_files = files
+                            st.session_state.current_url = url
                         return files
                     
                 files = asyncio.run(run_deep_search())
@@ -603,17 +605,22 @@ def main():
                     with col1:
                         if st.button("Select All"):
                             st.session_state.selected_files = list(range(len(files)))
+                            st.experimental_rerun()
                         if st.button("Clear Selection"):
                             st.session_state.selected_files = []
+                            st.experimental_rerun()
                     
                     # File selection
                     selected_files = st.multiselect(
                         "Select files to download",
-                        range(len(files)),
-                        default=st.session_state.get('selected_files', []),
+                        options=list(range(len(files))),
+                        default=st.session_state.selected_files,
                         format_func=lambda x: f"{files[x]['filename']} ({files[x]['size']})"
                     )
                     
+                    # Update session state
+                    st.session_state.selected_files = selected_files
+                    
                     if selected_files:
                         col1, col2, col3, col4 = st.columns(4)
                         with col1:
@@ -681,6 +688,97 @@ def main():
                 else:
                     st.warning("No files found.")
 
+        # Display current files if they exist in session state
+        elif st.session_state.discovered_files:
+            files = st.session_state.discovered_files
+            st.success(f"Found {len(files)} files!")
+            
+            # Select All/Clear Selection buttons
+            col1, col2 = st.columns([1, 4])
+            with col1:
+                if st.button("Select All"):
+                    st.session_state.selected_files = list(range(len(files)))
+                    st.experimental_rerun()
+                if st.button("Clear Selection"):
+                    st.session_state.selected_files = []
+                    st.experimental_rerun()
+            
+            # File selection
+            selected_files = st.multiselect(
+                "Select files to download",
+                options=list(range(len(files))),
+                default=st.session_state.selected_files,
+                format_func=lambda x: f"{files[x]['filename']} ({files[x]['size']})"
+            )
+            
+            # Update session state
+            st.session_state.selected_files = selected_files
+            
+            if selected_files:
+                col1, col2, col3, col4 = st.columns(4)
+                with col1:
+                    download_dir = st.text_input("Download Directory", value="./downloads")
+                with col2:
+                    create_zip = st.checkbox("Create ZIP file", value=True)
+                with col3:
+                    delete_after = st.checkbox("Delete after creating ZIP")
+                with col4:
+                    upload_to_drive = st.checkbox("Upload to Google Drive")
+                
+                if st.button("Download Selected"):
+                    if not os.path.exists(download_dir):
+                        os.makedirs(download_dir)
+                        
+                    async def download_files():
+                        downloaded_paths = []
+                        progress_bar = st.progress(0)
+                        status_text = st.empty()
+                        
+                        async with DownloadManager(use_proxy=use_proxy, proxy=proxy) as dm:
+                            for i, idx in enumerate(selected_files):
+                                progress = (i + 1) / len(selected_files)
+                                file_info = files[idx]
+                                
+                                status_text.text(f"Downloading {file_info['filename']}... ({i+1}/{len(selected_files)})")
+                                progress_bar.progress(progress)
+                                
+                                path = await dm.download_file(
+                                    file_info,
+                                    download_dir,
+                                    st.session_state.current_url
+                                )
+                                if path:
+                                    downloaded_paths.append(path)
+                            
+                            status_text.empty()
+                            progress_bar.empty()
+                            return downloaded_paths
+                    
+                    downloaded = asyncio.run(download_files())
+                    
+                    if downloaded:
+                        st.success(f"Successfully downloaded {len(downloaded)} files")
+                        
+                        if create_zip or upload_to_drive:
+                            zip_path = create_zip_file(downloaded, download_dir)
+                            st.success(f"Created ZIP file: {zip_path}")
+                            
+                            if upload_to_drive and st.session_state.get('google_creds'):
+                                with st.spinner("Uploading to Google Drive..."):
+                                    drive_id = google_drive_upload(zip_path, st.session_state.google_creds)
+                                    if not isinstance(drive_id, str) or not drive_id.startswith("Error"):
+                                        st.success(f"Uploaded to Google Drive. File ID: {drive_id}")
+                                    else:
+                                        st.error(drive_id)
+                            
+                            if delete_after:
+                                for path in downloaded:
+                                    try:
+                                        os.remove(path)
+                                    except Exception as e:
+                                        st.warning(f"Could not delete {path}: {e}")
+                                st.info("Deleted original files after ZIP creation")
+
     elif mode == "Bing Search":
         st.header("Bing Search Mode")
         query = st.text_input("Enter search query")
@@ -707,12 +805,11 @@ def main():
                                                 custom_ext_list=custom_extensions.split(',') if custom_extensions else [],
                                                 sublink_limit=max_sublinks
                                             )
-                                            # Reuse the same file handling logic as Manual URL mode
                                             if files:
                                                 st.session_state.discovered_files = files
                                                 st.session_state.current_url = url
-                                                st.success(f"Found {len(files)} files!")
-                                                # Add file selection and download UI here (same as Manual URL mode)
+                                                st.session_state.selected_files = []
+                                                st.experimental_rerun()
                                             else:
                                                 st.warning("No files found on this page.")
                             else: