Update main.py

main.py

@@ -21,6 +21,7 @@ from fpdf import FPDF
 import tempfile
 import urllib.parse
 from stories import generateResponse
+from styled_video_gen import create_styled_video, DEFAULT_WIDTH, DEFAULT_HEIGHT, DEFAULT_FPS, DEFAULT_TRANSITION_DURATION, DEFAULT_FONT, DEFAULT_LOGO_PATH
 
 
 # Initialize Flask app and CORS
@@ -651,6 +652,366 @@ def generate_video_endpoint():
         }), 500
 
 
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+
+# --- Styled Video Generation Endpoint ---
+@app.route('/api/styled_video/generate', methods=['POST']) # Uncomment when deploying in Flask app
+def generate_styled_video_endpoint(): # Use this if testing standalone
+    """
+    Generates a video based on story data from Firebase, applying client-specified options.
+    """
+    # --- Temporary file list for cleanup ---
+    temp_files_to_clean = []
+    video_output_path = None # Define here for broader scope in finally block
+
+    try:
+        logging.info("➡️ Received video generation request...")
+
+        # --- Authentication & Authorization ---
+        auth_header = request.headers.get('Authorization', '')
+        if not auth_header.startswith('Bearer '):
+            logging.error("❌ Auth Error: Missing or invalid token format")
+            return jsonify({'error': 'Missing or invalid token', 'log_url': upload_log()}), 401
+
+        token = auth_header.split(' ')[1]
+        uid = verify_token(token) # Assumes verify_token returns UID or None
+        if not uid:
+            logging.error("❌ Auth Error: Invalid or expired token")
+            return jsonify({'error': 'Invalid or expired token', 'log_url': upload_log()}), 401
+        logging.info(f"Authenticated user: {uid}")
+
+        # --- Check User Credits ---
+        user_ref = db.reference(f"users/{uid}") # Assumes this DB path structure
+        user_data = user_ref.get()
+        if not user_data: user_data = {} # Handle case where user node might not exist yet
+        current_credits = user_data.get("credits", 0)
+        video_cost = 10 # Define video cost
+
+        if current_credits < video_cost:
+            logging.warning(f"Insufficient credits for user {uid} (has {current_credits}, needs {video_cost})")
+            return jsonify({'error': f'Insufficient credits. You need at least {video_cost} credits.', 'log_url': upload_log()}), 403
+
+        # --- Get Request Data ---
+        data = request.get_json()
+        if not data:
+             logging.error("❌ Request Error: Invalid or missing JSON payload")
+             return jsonify({'error': 'Invalid JSON payload', 'log_url': upload_log()}), 400
+
+        story_id = data.get('story_id')
+        if not story_id:
+            logging.error("❌ Request Error: story_id is required")
+            return jsonify({'error': 'story_id is required', 'log_url': upload_log()}), 400
+
+        # Get video customization options from the request (or use empty dict for defaults)
+        video_options = data.get('video_options', {})
+        logging.info(f"Received video options: {video_options}")
+
+        # --- Fetch Story Data ---
+        logging.info(f"Fetching story '{story_id}' for user '{uid}' from Firebase...")
+        # Consider structuring story data under UID: db.reference(f"stories/{uid}/{story_id}")
+        # Using the path from your original code for now:
+        story_ref = db.reference(f"stories/{story_id}")
+        story_data = story_ref.get()
+
+        if not story_data:
+            logging.error(f"❌ Firebase Error: Story '{story_id}' not found.")
+            return jsonify({'error': 'Story not found', 'log_url': upload_log()}), 404
+
+        sections = story_data.get("sections", [])
+        if not sections or not isinstance(sections, list):
+            logging.error(f"❌ Data Error: No valid 'sections' array found in story '{story_id}'.")
+            return jsonify({'error': 'No sections found in the story', 'log_url': upload_log()}), 404
+
+        # --- Download Assets and Prepare Data ---
+        image_pil_list = []       # Stores downloaded PIL Images
+        audio_file_paths = []     # Stores paths to downloaded audio files
+        section_texts_list = []   # Stores text for each section
+        valid_section_indices = [] # Keep track of sections processed successfully
+
+        logging.info(f"Processing {len(sections)} sections for assets...")
+        download_errors = False
+        for i, section in enumerate(sections):
+            if not isinstance(section, dict):
+                logging.warning(f"Skipping section {i+1}, expected a dictionary, got {type(section)}")
+                continue # Skip malformed section
+
+            image_url = section.get("image_url")
+            audio_url = section.get("audio_url")
+            section_text = section.get("section_text") # Get the text
+
+            logging.info(f"--- Processing Section {i+1}/{len(sections)} ---")
+            logging.info(f"  Image URL: {image_url}")
+            logging.info(f"  Audio URL: {audio_url}")
+            logging.info(f"  Text: {str(section_text)[:50] if section_text else 'None'}")
+
+            # Download Image (Required for a section to be valid)
+            img_object = None
+            if image_url:
+                try:
+                    img_resp = requests.get(image_url, timeout=30)
+                    img_resp.raise_for_status()
+                    img_object = Image.open(io.BytesIO(img_resp.content))
+                    # Convert to RGB early to prevent potential palette issues later
+                    img_object = img_object.convert("RGB")
+                    logging.info("  ✅ Image downloaded and opened.")
+                except requests.exceptions.RequestException as e:
+                    logging.error(f"  ❌ ERROR downloading image {image_url}: {e}")
+                    download_errors = True
+                except UnidentifiedImageError: # Catch PIL errors
+                     logging.error(f"  ❌ ERROR: Cannot identify image file from {image_url}. Invalid format or corrupt?")
+                     download_errors = True
+                except Exception as e:
+                    logging.error(f"  ❌ ERROR processing image {image_url}: {e}")
+                    download_errors = True
+            else:
+                 logging.warning(f"  ⚠️ No image_url for section {i+1}. Skipping section.")
+                 # Don't add placeholders if image fails, just skip the index
+
+            # If image succeeded, process audio and text for this section index
+            if img_object:
+                 image_pil_list.append(img_object)
+                 section_texts_list.append(section_text) # Add text (can be None)
+                 valid_section_indices.append(i) # Mark this index as valid
+
+                 # Download Audio (Optional, will use silence if fails)
+                 audio_path = None
+                 if audio_url:
+                     try:
+                         aud_resp = requests.get(audio_url, timeout=60)
+                         aud_resp.raise_for_status()
+                         # Use a descriptive temp file name in system's temp dir
+                         temp_dir = tempfile.gettempdir()
+                         aud_filename = f"story_{story_id}_sec_{i}_audio_{uuid.uuid4().hex}.mp3" # Assume mp3, adjust if needed
+                         audio_path = os.path.join(temp_dir, aud_filename)
+                         with open(audio_path, "wb") as f:
+                             f.write(aud_resp.content)
+                         temp_files_to_clean.append(audio_path) # Add to cleanup list
+                         logging.info(f"  ✅ Audio downloaded to {audio_path}")
+                     except requests.exceptions.RequestException as e:
+                         logging.error(f"  ❌ ERROR downloading audio {audio_url}: {e}. Will use silence.")
+                         # download_errors = True # Don't mark as overall error if only audio fails
+                         audio_path = None # Ensure path is None on failure
+                     except Exception as e:
+                         logging.error(f"  ❌ ERROR saving audio {audio_url}: {e}. Will use silence.")
+                         # download_errors = True
+                         audio_path = None
+                 else:
+                     logging.info("  No audio_url for this section. Will use silence.")
+
+                 audio_file_paths.append(audio_path) # Add path or None
+
+            else: # Image failed, so skip adding audio/text for this section index
+                 logging.warning(f"Skipping audio/text for section {i+1} due to image failure.")
+
+
+        # Check if any valid sections remain
+        if not image_pil_list:
+             logging.error("❌ ERROR: No valid images could be downloaded or processed for any section.")
+             # upload_log() # Upload log before returning
+             return jsonify({'error': 'No images available for video generation', 'log_url': upload_log()}), 500
+
+        logging.info(f"Successfully processed {len(image_pil_list)} sections with images.")
+        if download_errors:
+             logging.warning("⚠️ Some assets encountered download/processing errors.")
+
+
+        # --- Handle Custom Logo/Watermark Download ---
+        custom_logo_path = None # Path to downloaded custom logo
+        watermark_opts_from_client = video_options.get("watermark_options", {})
+        watermark_final_config = watermark_opts_from_client.copy() # Start with client options
+        watermark_final_config["enabled"] = False # Default to disabled unless successfully set up
+
+        custom_logo_url = watermark_opts_from_client.get("custom_logo_url")
+        if watermark_opts_from_client.get("enabled") and custom_logo_url:
+            logging.info(f"➡️ Downloading custom logo/watermark: {custom_logo_url}")
+            try:
+                logo_resp = requests.get(custom_logo_url, timeout=20)
+                logo_resp.raise_for_status()
+                # Save custom logo to a temp file, try to get extension
+                file_ext = os.path.splitext(urllib.parse.urlparse(custom_logo_url).path)[1] or '.png' # Default to png
+                custom_logo_filename = f"custom_logo_{uid}_{uuid.uuid4().hex}{file_ext}"
+                custom_logo_path = os.path.join(tempfile.gettempdir(), custom_logo_filename)
+
+                with open(custom_logo_path, "wb") as f:
+                    f.write(logo_resp.content)
+                # Verify it's a valid image
+                try:
+                    Image.open(custom_logo_path).verify() # Quick check
+                    watermark_final_config["path"] = custom_logo_path # Update config with temp path
+                    watermark_final_config["enabled"] = True # Enable watermark
+                    temp_files_to_clean.append(custom_logo_path)
+                    logging.info(f"✅ Custom logo downloaded and verified: {custom_logo_path}")
+                except Exception as img_err:
+                     logging.error(f"❌ Custom logo file from {custom_logo_url} is not a valid image: {img_err}. Disabling watermark.")
+                     # Clean up invalid downloaded file immediately
+                     if os.path.exists(custom_logo_path): os.remove(custom_logo_path)
+
+            except requests.exceptions.RequestException as e:
+                logging.error(f"❌ ERROR downloading custom logo {custom_logo_url}: {e}. Watermark disabled.")
+            except Exception as e:
+                logging.error(f"❌ ERROR processing custom logo {custom_logo_url}: {e}. Watermark disabled.")
+        elif watermark_opts_from_client.get("enabled"):
+             logging.warning("⚠️ Watermark enabled in options, but no 'custom_logo_url' provided. Watermark disabled.")
+
+
+        # --- Prepare Final Config for create_video ---
+        final_video_config = {
+            "width": video_options.get("width", DEFAULT_WIDTH),
+            "height": video_options.get("height", DEFAULT_HEIGHT),
+            "fps": video_options.get("fps", DEFAULT_FPS),
+            "transition": video_options.get("transition", "fade"),
+            "transition_duration": video_options.get("transition_duration", DEFAULT_TRANSITION_DURATION),
+            "font_path": video_options.get("font_path", DEFAULT_FONT), # Allow overriding default font
+
+            "subtitle_options": video_options.get("subtitle_options", {"enabled": True}), # Default enabled
+
+            "particle_options": video_options.get("particle_options", {"enabled": False}), # Default disabled
+
+            "watermark_options": watermark_final_config, # Use the processed watermark config
+
+             # Use default Sozo logo unless custom one provided? (Simplifying: only supporting default end logo for now)
+            "end_logo_options": {
+                 "enabled": video_options.get("use_end_logo", True), # Control if end logo is used
+                 "path": video_options.get("end_logo_path", DEFAULT_LOGO_PATH), # Allow overriding default logo path via options
+                 "duration": video_options.get("end_logo_duration", 3.0)
+             },
+        }
+
+        # Ensure particle types list matches the number of *final* valid sections
+        particle_opts_config = final_video_config["particle_options"]
+        if particle_opts_config.get("enabled"):
+            particle_types_list_orig = particle_opts_config.get("types_per_section", [])
+            if isinstance(particle_types_list_orig, list):
+                 # Filter the original particle list based on the indices of sections that were successfully processed
+                 filtered_particle_types = [particle_types_list_orig[i] for i in valid_section_indices if i < len(particle_types_list_orig)]
+                 # Pad with None if the original list was too short
+                 if len(filtered_particle_types) < len(image_pil_list):
+                      filtered_particle_types.extend([None] * (len(image_pil_list) - len(filtered_particle_types)))
+                 particle_opts_config["types_per_section"] = filtered_particle_types
+                 logging.info(f"Aligned particle types for {len(image_pil_list)} sections: {particle_opts_config['types_per_section']}")
+            else:
+                logging.warning("particle_options.types_per_section was not a list. Disabling particles.")
+                particle_opts_config["enabled"] = False
+
+
+        # --- Create the Video ---
+        # Define output path in temp directory
+        video_output_filename = f"final_video_{uid}_{story_id}_{uuid.uuid4().hex}.mp4"
+        video_output_path = os.path.join(tempfile.gettempdir(), video_output_filename)
+        # Don't add to cleanup list immediately - only if creation fails or after successful upload
+
+        logging.info("🚀 Starting video creation with MoviePy...")
+        logging.info(f"Video Config Passed to create_video: {final_video_config}")
+        logging.info(f"Output Path: {video_output_path}")
+        logging.info(f"Number of image inputs: {len(image_pil_list)}")
+        logging.info(f"Number of audio inputs: {len(audio_file_paths)}")
+        logging.info(f"Number of text inputs: {len(section_texts_list)}")
+
+
+        # Call the MoviePy function from video_gen.py
+        generated_video_path = create_styled_video(
+            images=image_pil_list,          # List of PIL images
+            audio_files=audio_file_paths,   # List of paths (or None)
+            section_texts=section_texts_list,# List of strings (or None)
+            output_path=video_output_path,
+            config=final_video_config       # The dictionary of options
+        )
+
+        # --- Handle Video Creation Result ---
+        if not generated_video_path or not os.path.exists(generated_video_path):
+            logging.error("❌ ERROR: Video generation failed (create_video returned None or file missing).")
+            # Add the intended output path to cleanup just in case a partial file exists
+            temp_files_to_clean.append(video_output_path)
+            # upload_log() # Upload log before returning
+            return jsonify({'error': 'Video generation failed', 'log_url': upload_log()}), 500
+
+        logging.info(f"✅ Video generated successfully: {generated_video_path}")
+        # Add the successfully generated video path for cleanup after upload
+        temp_files_to_clean.append(generated_video_path)
+
+
+        # --- Upload Video to Firebase Storage ---
+        logging.info(f"☁️ Uploading video '{os.path.basename(generated_video_path)}' to Firebase Storage...")
+        # Make blob name more descriptive
+        video_blob_name = f"stories/{uid}/{story_id}/video_{uuid.uuid4().hex}.mp4"
+        try:
+             # Assuming upload_to_storage handles the upload and returns public URL
+             video_url = upload_to_storage(generated_video_path, video_blob_name)
+             if not video_url:
+                  raise Exception("Upload function returned no URL")
+             logging.info(f"✅ Video uploaded successfully to: {video_url}")
+        except Exception as upload_err:
+             logging.error(f"❌ Firebase Storage Error: Failed to upload video: {upload_err}")
+             # upload_log() # Upload log before returning
+             # Don't deduct credits if upload fails
+             return jsonify({'error': 'Video generated but failed to upload to storage.', 'log_url': upload_log()}), 500
+
+
+        # --- Update Firebase Realtime Database ---
+        try:
+            story_ref.update({"video_url": video_url, "last_generated": time.time()}) # Add timestamp
+            logging.info(f"✅ Updated story '{story_id}' record with video URL.")
+        except Exception as db_err:
+             logging.error(f"❌ Firebase DB Error: Failed to update story record: {db_err}")
+             # Decide if this is critical. Maybe log and continue, but don't deduct credits?
+             # For now, let's return an error as the client won't see the video URL in the story record.
+             # upload_log()
+             return jsonify({'error': 'Video generated and uploaded, but failed to update story record.', 'log_url': upload_log()}), 500
+
+
+        # --- Deduct Credits (Only after successful generation, upload, and DB update) ---
+        try:
+            new_credits = max(0, current_credits - video_cost)
+            user_ref.update({"credits": new_credits})
+            logging.info(f"✅ Deducted {video_cost} credits for user {uid}. New balance: {new_credits}")
+        except Exception as credit_err:
+             logging.error(f"❌ Firebase DB Error: Failed to update user credits: {credit_err}")
+             # This is less critical, log it but still return success to user
+             # upload_log() # Upload log
+
+        # --- Success Response ---
+        # final_log_url = upload_log()
+        return jsonify({
+            "message": "Video generated and uploaded successfully!",
+            "video_url": video_url,
+            "new_credits": new_credits,
+            # "log_url": final_log_url
+             "log_url": "Log upload function placeholder" # Replace with actual call if needed
+        }), 200 # Use 200 OK for success
+
+    except Exception as e:
+        # --- Generic Error Handler ---
+        trace = traceback.format_exc()
+        logging.error(f"❌ UNHANDLED EXCEPTION in generate_video_endpoint: {str(e)}\n{trace}")
+        # log_url = upload_log() # Upload log before returning
+        return jsonify({
+            'error': f"An unexpected error occurred: {str(e)}",
+            # 'log_url': log_url
+            'log_url': "Log upload function placeholder"
+        }), 500
+
+    finally:
+        # --- Cleanup Temporary Files ---
+        logging.info(f"🧹 Cleaning up {len(temp_files_to_clean)} temporary files...")
+        cleaned_count = 0
+        failed_count = 0
+        for file_path in temp_files_to_clean:
+             if file_path and os.path.exists(file_path):
+                 try:
+                     os.remove(file_path)
+                     # logging.debug(f"  - Removed: {file_path}")
+                     cleaned_count += 1
+                 except Exception as e:
+                     logging.error(f"  - Failed to remove temp file {file_path}: {e}")
+                     failed_count += 1
+             #else:
+                 # logging.debug(f"  - Skipping non-existent path: {file_path}")
+
+        logging.info(f"✅ Cleanup complete. Removed {cleaned_count} files, failed to remove {failed_count}.")
+        
+
+
 #----------Image Editing Endpoint ----------
 @app.route('/api/story/<string:story_id>/sections/<int:section_idx>/edit-image', methods=['POST'])
 def edit_section_image_endpoint(story_id, section_idx):