Update app.py via AI Editor

app.py

@@ -7,16 +7,16 @@ import os
 import tempfile
 import shutil
 import logging
-from flask import send_file, make_response
+from flask import send_file
 import threading
-import pickle
 from PyPDF2 import PdfReader, PdfWriter
 import re
+import zipfile
+import base64
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
 
-# Session storage and lock management
 SESSION_DATA = {}
 SESSION_LOCKS = {}
 
@@ -54,7 +54,6 @@ def extract_text_headers(reader, page_num):
     try:
         page = reader.pages[page_num]
         text = page.extract_text() or ""
-        # Extract the first non-blank line as a potential header
         lines = [line.strip() for line in text.split('\n') if line.strip()]
         header = lines[0] if lines else ""
         return header
@@ -104,10 +103,8 @@ def intelligent_pdf_split(input_path, session_dir, max_mb=5, min_split_mb=4):
         chapter = is_chapter_header(header)
         split_here = False
 
-        # Force split if over max size
         if size >= max_mb:
             split_here = True
-        # Prefer to split between min_split_mb and max_mb at logical points
         elif size >= min_split_mb:
             if blank or chapter or (header and header != last_header):
                 split_here = True
@@ -118,11 +115,9 @@ def intelligent_pdf_split(input_path, session_dir, max_mb=5, min_split_mb=4):
             current_writer = PdfWriter()
         last_header = header
 
-    # Add final split if not already
     if last_split_at < n_pages:
         splits.append((last_split_at, n_pages))
 
-    # Write split files
     split_files = []
     for idx, (start, end) in enumerate(splits):
         writer = PdfWriter()
@@ -135,11 +130,16 @@ def intelligent_pdf_split(input_path, session_dir, max_mb=5, min_split_mb=4):
         split_files.append({'filename': os.path.basename(out_path), 'size': size, 'path': out_path})
     return split_files
 
-# Dash app setup
+def make_zip_of_splits(split_files, session_dir):
+    zip_path = os.path.join(session_dir, "split_files.zip")
+    with zipfile.ZipFile(zip_path, 'w') as zipf:
+        for file in split_files:
+            zipf.write(file['path'], arcname=file['filename'])
+    return zip_path
+
 external_stylesheets = [dbc.themes.BOOTSTRAP]
 app = dash.Dash(__name__, external_stylesheets=external_stylesheets, suppress_callback_exceptions=True)
 server = app.server
-
 app.title = "Intelligent PDF Splitter"
 
 app.layout = dbc.Container(
@@ -169,8 +169,8 @@ app.layout = dbc.Container(
                                         multiple=False,
                                         accept='.pdf'
                                     ),
-                                    html.Div(id='file-info'),
-                                    dbc.Button("Clear Session", id='clear-session', color='secondary', className='mt-2'),
+                                    html.Div(id='file-info', className='mb-4'),
+                                    dbc.Button("Clear Session", id='clear-session', color='secondary', className='mt-2 mb-2'),
                                     dcc.Loading(
                                         id="loading", type="default",
                                         children=[html.Div(id='split-results')]
@@ -196,84 +196,170 @@ app.layout = dbc.Container(
     Input('upload-pdf', 'contents'),
     State('upload-pdf', 'filename'),
     Input('clear-session', 'n_clicks'),
+    Input({'type': 'delete-upload-btn', 'index': 0}, 'n_clicks'),
+    Input('split-btn', 'n_clicks'),
     State('session-store', 'data'),
     prevent_initial_call='initial_duplicate'
 )
-def handle_upload(contents, filename, clear_n, session_data):
+def handle_upload(contents, filename, clear_n, delete_upload_n, split_n, session_data):
     trigger = ctx.triggered_id
     session_id = get_session_id()
     flask.g.session_id = session_id
     session_dir = get_session_dir(session_id)
     lock = get_session_lock(session_id)
 
+    if session_data is None:
+        session_data = {}
+
     if trigger == 'clear-session':
         clean_session(session_id)
         resp_data = {}
         return "", "", resp_data
 
-    # If user returns, restore state
-    if not contents and session_data and 'split_files' in session_data:
-        split_files = session_data.get('split_files', [])
-        file_info = html.Div(f"Previous upload: {session_data.get('orig_filename', '')}")
-        results = [
-            html.H5("Split Files:"),
-            html.Ul([
-                html.Li([
-                    f"{fi['filename']} ({fi['size']:.2f} MB) ",
-                    dbc.Button("Download", id={'type': 'download-btn', 'index': idx}, href=f"/download/{session_id}/{fi['filename']}", color='primary', size='sm')
-                ]) for idx, fi in enumerate(split_files)
+    if trigger == {'type': 'delete-upload-btn', 'index': 0}:
+        orig_filename = session_data.get('orig_filename', '')
+        pdf_path = os.path.join(session_dir, orig_filename)
+        if os.path.exists(pdf_path):
+            os.remove(pdf_path)
+        session_data = {}
+        if os.path.exists(session_dir):
+            for file in os.listdir(session_dir):
+                os.remove(os.path.join(session_dir, file))
+        return "", "", {}
+
+    if trigger == 'upload-pdf':
+        if not contents:
+            return "", "", {}
+
+        if not allowed_file(filename):
+            return html.Div("Only .pdf files are allowed.", style={'color': 'red'}), "", {}
+
+        try:
+            header, b64data = contents.split(',', 1)
+            pdf_bytes = base64.b64decode(b64data)
+            pdf_path = os.path.join(session_dir, filename)
+            with open(pdf_path, 'wb') as f:
+                f.write(pdf_bytes)
+            logging.info(f"PDF uploaded and saved to {pdf_path} for session {session_id}")
+
+            session_data = {
+                'orig_filename': filename,
+                'split_files': None,
+                'zip_ready': False,
+            }
+            file_info = dbc.Row([
+                dbc.Col(html.Div(f"Uploaded: {filename} ({len(pdf_bytes)/1024/1024:.2f} MB)"), width=9, style={'display': 'flex', 'alignItems': 'center'}),
+                dbc.Col(
+                    dbc.Button("Delete", id={'type': 'delete-upload-btn', 'index': 0}, color='danger', n_clicks=0, className='ms-5'),
+                    width=3, style={'display': 'flex', 'justifyContent': 'end'}
+                )
+            ], className='mb-3', align='center', style={'marginTop': "15px", 'marginBottom': '25px'})
+            split_results = dbc.Row([
+                dbc.Col(
+                    dbc.Button("Split PDF", id='split-btn', color='primary', className='mb-3 mt-2', n_clicks=0, style={'width': '180px', 'fontWeight': 'bold'}),
+                    width=12, style={'display': 'flex', 'justifyContent': 'center'}
+                )
             ])
-        ]
-        return file_info, results, session_data
+            return file_info, split_results, session_data
+        except Exception as e:
+            logging.error(f"Error processing PDF: {e}")
+            return html.Div(f"Error: {e}", style={'color': 'red'}), "", {}
 
-    if not contents:
-        return "", "", {}
+    if session_data.get('orig_filename') and not session_data.get('split_files'):
+        # If user refreshes after upload but before split, restore file info and split button
+        file_info = dbc.Row([
+            dbc.Col(html.Div(f"Uploaded: {session_data['orig_filename']}"), width=9, style={'display': 'flex', 'alignItems': 'center'}),
+            dbc.Col(
+                dbc.Button("Delete", id={'type': 'delete-upload-btn', 'index': 0}, color='danger', n_clicks=0, className='ms-5'),
+                width=3, style={'display': 'flex', 'justifyContent': 'end'}
+            )
+        ], className='mb-3', align='center', style={'marginTop': "15px", 'marginBottom': '25px'})
+        split_results = dbc.Row([
+            dbc.Col(
+                dbc.Button("Split PDF", id='split-btn', color='primary', className='mb-3 mt-2', n_clicks=0, style={'width': '180px', 'fontWeight': 'bold'}),
+                width=12, style={'display': 'flex', 'justifyContent': 'center'}
+            )
+        ])
+        return file_info, split_results, session_data
 
-    if not allowed_file(filename):
-        return html.Div("Only .pdf files are allowed.", style={'color': 'red'}), "", {}
+    if trigger == 'split-btn':
+        orig_filename = session_data.get('orig_filename')
+        if not orig_filename:
+            return html.Div("No file to split.", style={'color': 'red'}), "", session_data
+        pdf_path = os.path.join(session_dir, orig_filename)
+        if not os.path.exists(pdf_path):
+            return html.Div("Uploaded file not found. Please upload again.", style={'color': 'red'}), "", {}
+        try:
+            with lock:
+                split_files = intelligent_pdf_split(pdf_path, session_dir)
+                zip_path = make_zip_of_splits(split_files, session_dir)
+            session_data['split_files'] = split_files
+            session_data['zip_ready'] = True
+            file_info = dbc.Row([
+                dbc.Col(html.Div(f"Uploaded: {orig_filename}"), width=9, style={'display': 'flex', 'alignItems': 'center'}),
+                dbc.Col(
+                    dbc.Button("Delete", id={'type': 'delete-upload-btn', 'index': 0}, color='danger', n_clicks=0, className='ms-5'),
+                    width=3, style={'display': 'flex', 'justifyContent': 'end'}
+                )
+            ], className='mb-3', align='center', style={'marginTop': "15px", 'marginBottom': '25px'})
+            split_files_list = html.Ul([
+                html.Li([
+                    f"{fi['filename']} ({fi['size']:.2f} MB)"
+                ]) for fi in split_files
+            ])
+            download_zip_btn = dbc.Button(
+                "Download All (ZIP)", color="primary", size="lg", className='mb-3 mt-4',
+                href=f"/download_zip/{session_id}/split_files.zip"
+            )
+            results = [
+                html.H5("Split Files:"),
+                split_files_list,
+                html.Div(download_zip_btn, style={'marginTop': '30px'})
+            ]
+            logging.info(f"PDF split into {len(split_files)} chunks for session {session_id}, zip ready.")
+            return file_info, results, session_data
+        except Exception as e:
+            logging.error(f"Error splitting PDF: {e}")
+            return html.Div(f"Error: {e}", style={'color': 'red'}), "", session_data
 
-    try:
-        # Save file
-        header, b64data = contents.split(',', 1)
-        import base64
-        pdf_bytes = base64.b64decode(b64data)
-        pdf_path = os.path.join(session_dir, filename)
-        with open(pdf_path, 'wb') as f:
-            f.write(pdf_bytes)
-        logging.info(f"PDF uploaded and saved to {pdf_path} for session {session_id}")
-
-        # Split PDF with lock
-        with lock:
-            split_files = intelligent_pdf_split(pdf_path, session_dir)
+    # Restore split results if user refreshes after splitting
+    if session_data.get('split_files'):
+        split_files = session_data['split_files']
+        orig_filename = session_data.get('orig_filename', '')
+        file_info = dbc.Row([
+            dbc.Col(html.Div(f"Uploaded: {orig_filename}"), width=9, style={'display': 'flex', 'alignItems': 'center'}),
+            dbc.Col(
+                dbc.Button("Delete", id={'type': 'delete-upload-btn', 'index': 0}, color='danger', n_clicks=0, className='ms-5'),
+                width=3, style={'display': 'flex', 'justifyContent': 'end'}
+            )
+        ], className='mb-3', align='center', style={'marginTop': "15px", 'marginBottom': '25px'})
+        split_files_list = html.Ul([
+            html.Li([
+                f"{fi['filename']} ({fi['size']:.2f} MB)"
+            ]) for fi in split_files
+        ])
+        download_zip_btn = dbc.Button(
+            "Download All (ZIP)", color="primary", size="lg", className='mb-3 mt-4',
+            href=f"/download_zip/{session_id}/split_files.zip"
+        )
         results = [
             html.H5("Split Files:"),
-            html.Ul([
-                html.Li([
-                    f"{fi['filename']} ({fi['size']:.2f} MB) ",
-                    dbc.Button("Download", id={'type': 'download-btn', 'index': idx}, href=f"/download/{session_id}/{fi['filename']}", color='primary', size='sm')
-                ]) for idx, fi in enumerate(split_files)
-            ])
+            split_files_list,
+            html.Div(download_zip_btn, style={'marginTop': '30px'})
         ]
-        file_info = html.Div(f"Uploaded: {filename} ({len(pdf_bytes)/1024/1024:.2f} MB)")
-        session_data = {
-            'orig_filename': filename,
-            'split_files': split_files,
-        }
-        logging.info(f"PDF split into {len(split_files)} chunks for session {session_id}")
         return file_info, results, session_data
-    except Exception as e:
-        logging.error(f"Error processing PDF: {e}")
-        return html.Div(f"Error: {e}", style={'color': 'red'}), "", {}
 
-@app.server.route('/download/<session_id>/<filename>')
-def download_split_file(session_id, filename):
+    return "", "", session_data
+
+@app.server.route('/download_zip/<session_id>/<filename>')
+def download_zip_file(session_id, filename):
     session_dir = get_session_dir(session_id)
     file_path = os.path.join(session_dir, filename)
     if os.path.exists(file_path):
-        logging.info(f"Serving file {file_path} for session {session_id}")
-        return send_file(file_path, mimetype='application/pdf', as_attachment=True, download_name=filename)
+        logging.info(f"Serving zip file {file_path} for session {session_id}")
+        return send_file(file_path, mimetype='application/zip', as_attachment=True, download_name=filename)
     else:
-        logging.error(f"File not found for download: {file_path}")
+        logging.error(f"ZIP file not found for download: {file_path}")
         return "File not found", 404
 
 @app.callback(