Update main.py

main.py

@@ -8,7 +8,7 @@ import requests
 import json
 import pandas as pd
 from datetime import datetime
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS, cross_origin
 from firebase_admin import credentials, db, storage, auth
 import firebase_admin
@@ -16,6 +16,8 @@ from PIL import ImageFont, ImageDraw, Image
 import logging
 import traceback
 from video_gen import create_video
+import zipfile
+from fpdf import FPDF
 
 
 # Initialize Flask app and CORS
@@ -610,8 +612,10 @@ def view_videos():
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
-@app.route('/api/download/story_asset/<story_id>', methods=['GET'])
-def download_story_asset(story_id):
+
+
+@app.route('/api/download/story_archive/<story_id>', methods=['GET'])
+def download_story_archive(story_id):
     try:
         # --- Authentication ---
         auth_header = request.headers.get('Authorization', '')
@@ -622,7 +626,7 @@ def download_story_asset(story_id):
         if not uid:
             return jsonify({'error': 'Invalid or expired token'}), 401
 
-        # --- Check if the Story Exists and Belongs to the User ---
+        # --- Fetch the Story ---
         stories_ref = db.reference('stories')
         story_record = stories_ref.child(story_id).get()
 
@@ -631,30 +635,69 @@ def download_story_asset(story_id):
         if story_record.get('uid') != uid:
             return jsonify({'error': 'Unauthorized'}), 403
 
-        # --- Determine the Asset Type ---
-        video_url = story_record.get("video_url", None)
-        audio_url = story_record.get("audio_url", None)  # Add if audio files exist
-        image_url = story_record.get("image_url", None)  # Add if images exist
-
-        # Prioritize video > audio > image
-        file_url = video_url or audio_url or image_url
-        if not file_url:
-            return jsonify({'error': 'No asset found for this story'}), 404
-
-        # --- If Stored in Firebase Storage, Generate a Signed URL ---
-        if "firebasestorage.googleapis.com" in file_url:
-            bucket = storage.bucket()
-            blob_path = file_url.split('/o/')[-1].split('?')[0]  # Extract path
-            blob = bucket.blob(blob_path)
-            signed_url = blob.generate_signed_url(expiration=600)  # Valid for 10 minutes
-            return jsonify({"download_url": signed_url})
-
-        # --- If Stored Locally, Serve the File ---
-        local_path = os.path.join("uploads", file_url)  # Adjust path accordingly
-        if os.path.exists(local_path):
-            return send_file(local_path, as_attachment=True)
-
-        return jsonify({'error': 'File not found'}), 404
+        # --- Create ZIP File ---
+        zip_buffer = io.BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zipf:
+
+            # --- Generate a PDF with Story Content ---
+            pdf = FPDF()
+            pdf.set_auto_page_break(auto=True, margin=15)
+            pdf.add_page()
+            pdf.set_font("Arial", size=12)
+            pdf.multi_cell(0, 10, f"Story ID: {story_id}\n\n")
+            pdf.multi_cell(0, 10, f"Title: {story_record.get('title', 'Untitled')}\n\n")
+            pdf.multi_cell(0, 10, f"Full Story:\n\n{story_record.get('full_story', '')}\n\n")
+
+            # Add sections if available
+            sections = story_record.get("sections", [])
+            for idx, section in enumerate(sections):
+                pdf.add_page()
+                pdf.set_font("Arial", style='B', size=14)
+                pdf.multi_cell(0, 10, f"Section {idx + 1}\n")
+                pdf.set_font("Arial", size=12)
+                pdf.multi_cell(0, 10, section)
+
+            # Save PDF to buffer and add to ZIP
+            pdf_buffer = io.BytesIO()
+            pdf.output(pdf_buffer)
+            pdf_buffer.seek(0)
+            zipf.writestr(f"{story_id}.pdf", pdf_buffer.read())
+
+            # --- Generate a TXT file with Story Content ---
+            txt_content = f"Story ID: {story_id}\n\n"
+            txt_content += f"Title: {story_record.get('title', 'Untitled')}\n\n"
+            txt_content += f"Full Story:\n\n{story_record.get('full_story', '')}\n\n"
+
+            # Add sections if available
+            for idx, section in enumerate(sections):
+                txt_content += f"\n\nSection {idx + 1}:\n{section}"
+
+            # Add TXT file to ZIP
+            zipf.writestr(f"{story_id}.txt", txt_content)
+
+            # --- Download Images ---
+            image_urls = story_record.get("image_urls", [])  # Expecting a list of image URLs
+            for idx, img_url in enumerate(image_urls):
+                if "firebasestorage.googleapis.com" in img_url:
+                    bucket = storage.bucket()
+                    blob_path = img_url.split('/o/')[-1].split('?')[0]
+                    blob = bucket.blob(blob_path)
+                    image_data = blob.download_as_bytes()
+                    zipf.writestr(f"image_{idx + 1}.jpg", image_data)
+
+            # --- Download Audio Files ---
+            audio_urls = story_record.get("audio_urls", [])  # Expecting a list of audio URLs
+            for idx, audio_url in enumerate(audio_urls):
+                if "firebasestorage.googleapis.com" in audio_url:
+                    bucket = storage.bucket()
+                    blob_path = audio_url.split('/o/')[-1].split('?')[0]
+                    blob = bucket.blob(blob_path)
+                    audio_data = blob.download_as_bytes()
+                    zipf.writestr(f"audio_{idx + 1}.mp3", audio_data)
+
+        # --- Serve ZIP File ---
+        zip_buffer.seek(0)
+        return send_file(zip_buffer, mimetype='application/zip', as_attachment=True, download_name=f"story_{story_id}.zip")
 
     except Exception as e:
         return jsonify({'error': str(e)}), 500