From 690c2b699ccdd60af9fb0f9e254379183dd74bac Mon Sep 17 00:00:00 2001
From: CaliBrain <calibrain@l4n.xyz>
Date: Tue, 6 May 2025 08:44:43 -0400
Subject: [PATCH] Fix download for other files than epub (#174)

Fix #173
---
 app.py     | 14 +++++++-------
 backend.py |  4 ++--
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/app.py b/app.py
index e2ee8e00..e065d151 100644
--- a/app.py
+++ b/app.py
@@ -239,19 +239,19 @@ def api_local_download() -> Union[Response, Tuple[Response, int]]:
         return jsonify({"error": "No book ID provided"}), 400
 
     try:
-        file_data, file_name = backend.get_book_data(book_id)
+        file_data, book_info = backend.get_book_data(book_id)
         if file_data is None:
             # Book data not found or not available
             return jsonify({"error": "File not found"}), 404
         # Santize the file name
-        file_name = re.sub(r'[\\/:*?"<>|]', '_', file_name.strip())[:255]
+        file_name = book_info.title
+        file_name = re.sub(r'[\\/:*?"<>|]', '_', file_name.strip())[:245]
+        file_extension = book_info.format
         # Prepare the file for sending to the client
-        epub_file = io.BytesIO(file_data)
-        # Typically EPUB mime-type: 'application/epub+zip'
+        data = io.BytesIO(file_data)
         return send_file(
-            epub_file,
-            mimetype='application/epub+zip',
-            download_name=f"{file_name}.epub",
+            data,
+            download_name=f"{file_name}.{file_extension}",
             as_attachment=True
         )
 
diff --git a/backend.py b/backend.py
index 20578669..c0e62e5c 100644
--- a/backend.py
+++ b/backend.py
@@ -84,7 +84,7 @@ def queue_status() -> Dict[str, Dict[str, Any]]:
         for status_type, books in status.items()
     }
 
-def get_book_data(book_id: str) -> Tuple[Optional[bytes], str] :
+def get_book_data(book_id: str) -> Tuple[Optional[bytes], BookInfo]:
     """Get book data for a specific book, including its title.
     
     Args:
@@ -97,7 +97,7 @@ def get_book_data(book_id: str) -> Tuple[Optional[bytes], str] :
         book_info = book_queue._book_data[book_id]
         path = book_info.download_path
         with open(path, "rb") as f:
-            return f.read(), book_info.title
+            return f.read(), book_info
     except Exception as e:
         logger.error_trace(f"Error getting book data: {e}")
         book_info.download_path = None