Update main.py

main.py

@@ -639,12 +639,95 @@ def generate_sozo_report(project_id):
         return jsonify({'error': str(e)}), 500
 
 
+def send_video_generation_notification(user_id, user_email, project_name, video_url, send_email=False):
+    """
+    Send notification when video generation is completed successfully.
+    
+    Args:
+        user_id (str): Firebase user ID
+        user_email (str): User's email address
+        project_name (str): Name of the project
+        video_url (str): URL of the generated video
+        send_email (bool): Whether to send email notification
+    
+    Returns:
+        bool: True if successful, False otherwise
+    """
+    try:
+        # Create the notification message
+        message_content = f"Your video for project '{project_name}' has been generated successfully!"
+        
+        # Create email subject and body if email notification is requested
+        email_subject = None
+        email_body = None
+        
+        if send_email:
+            email_subject = f"🎥 Video Generation Complete - {project_name}"
+            email_body = f"""
+            <div style="font-family: Arial, sans-serif; max-width: 600px; margin: 0 auto; padding: 20px; background-color: #f9f9f9;">
+                <div style="background-color: white; padding: 30px; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);">
+                    <div style="text-align: center; margin-bottom: 30px;">
+                        <h1 style="color: #2563eb; margin: 0; font-size: 28px;">🎉 Video Generation Complete!</h1>
+                    </div>
+                    
+                    <div style="background-color: #f0f9ff; padding: 20px; border-radius: 8px; border-left: 4px solid #2563eb; margin-bottom: 25px;">
+                        <h2 style="color: #1e40af; margin: 0 0 10px 0; font-size: 20px;">Your video is ready!</h2>
+                        <p style="color: #374151; margin: 0; font-size: 16px;">
+                            Great news! Your video for project <strong>"{project_name}"</strong> has been successfully generated and is ready for viewing.
+                        </p>
+                    </div>
+                    
+                    <div style="text-align: center; margin: 30px 0;">
+                        <a href="{video_url}" 
+                           style="display: inline-block; background-color: #2563eb; color: white; padding: 15px 30px; 
+                                  text-decoration: none; border-radius: 8px; font-weight: bold; font-size: 16px;
+                                  box-shadow: 0 2px 4px rgba(37, 99, 235, 0.3);">
+                            🎬 Watch Your Video
+                        </a>
+                    </div>
+                    
+                    <div style="background-color: #f8fafc; padding: 20px; border-radius: 8px; margin-top: 25px;">
+                        <h3 style="color: #374151; margin: 0 0 15px 0; font-size: 18px;">What's Next?</h3>
+                        <ul style="color: #6b7280; margin: 0; padding-left: 20px; line-height: 1.6;">
+                            <li>Click the button above to view your generated video</li>
+                            <li>Share your video with your team or clients</li>
+                            <li>Download the video for offline use</li>
+                            <li>Create more videos with your remaining credits</li>
+                        </ul>
+                    </div>
+                    
+                    <div style="text-align: center; margin-top: 30px; padding-top: 20px; border-top: 1px solid #e5e7eb;">
+                        <p style="color: #9ca3af; font-size: 14px; margin: 0;">
+                            This video was generated by <strong>Sozo Business Studio</strong><br>
+                            Need help? Contact us at <a href="mailto:support@sozofix.tech" style="color: #2563eb;">support@sozofix.tech</a>
+                        </p>
+                    </div>
+                </div>
+            </div>
+            """
+        
+        # Use the existing _send_notification function
+        return _send_notification(
+            user_id=user_id,
+            user_email=user_email,
+            message_content=message_content,
+            send_email=send_email,
+            email_subject=email_subject,
+            email_body=email_body
+        )
+        
+    except Exception as e:
+        logger.error(f"Failed to send video generation notification: {e}")
+        return False
+
+
 @app.route('/api/sozo/projects/<string:project_id>/generate-video', methods=['POST'])
 def generate_sozo_video(project_id):
     try:
         token = request.headers.get('Authorization', '').split(' ')[1]
         uid = verify_token(token)
-        if not uid: return jsonify({'error': 'Unauthorized'}), 401
+        if not uid: 
+            return jsonify({'error': 'Unauthorized'}), 401
         
         # Check user credits first
         user_ref = db.reference(f'users/{uid}')
@@ -658,11 +741,15 @@ def generate_sozo_video(project_id):
         
         project_ref = db.reference(f'sozo_projects/{project_id}')
         project_data = project_ref.get()
-        if not project_data or project_data.get('uid') != uid: return jsonify({'error': 'Project not found or unauthorized'}), 404
+        if not project_data or project_data.get('uid') != uid: 
+            return jsonify({'error': 'Project not found or unauthorized'}), 404
         
         data = request.get_json()
         voice_model = data.get('voice_model', 'aura-2-andromeda-en')
         
+        # NEW: Get notification preference from request
+        send_email_notification = data.get('send_email_notification', False)
+        
         project_ref.update({'status': 'generating_video'})
         
         blob_path = f"sozo_projects/{uid}/{project_id}/data{Path(project_data['originalFilename']).suffix}"
@@ -670,28 +757,49 @@ def generate_sozo_video(project_id):
         file_bytes = blob.download_as_bytes()
         df = load_dataframe_safely(io.BytesIO(file_bytes), project_data['originalFilename'])
         
-        # **CORRECTION HERE: Pass the required 'data_context' argument**
+        # Generate the video
         video_url = generate_video_from_project(
             df,
             project_data.get('rawMarkdown', ''),
-            project_data.get('dataContext', {}), # Load the context from the project
+            project_data.get('dataContext', {}),
             uid,
             project_id,
             voice_model,
-            bucket                              # The missing argument from your error message
+            bucket
         )
         
-        if not video_url: raise Exception("Video generation failed in core function.")
+        if not video_url: 
+            raise Exception("Video generation failed in core function.")
         
         project_ref.update({'status': 'video_complete', 'videoUrl': video_url})
         
         # Deduct credits ONLY after successful generation
         user_ref.update({'credits': current_credits - 5})
         
+        # NEW: Send notification if requested
+        if send_email_notification:
+            project_name = project_data.get('name', 'Unnamed Project')
+            user_email = user_data.get('email')
+            
+            if user_email:
+                notification_sent = send_video_generation_notification(
+                    user_id=uid,
+                    user_email=user_email,
+                    project_name=project_name,
+                    video_url=video_url,
+                    send_email=True
+                )
+                
+                if not notification_sent:
+                    logger.warning(f"Failed to send notification for video generation completion to user {uid}")
+            else:
+                logger.warning(f"No email found for user {uid}, skipping email notification")
+        
         return jsonify({
             'success': True, 
             'video_url': video_url,
-            'credits_remaining': current_credits - 5
+            'credits_remaining': current_credits - 5,
+            'notification_sent': send_email_notification and user_data.get('email') is not None
         }), 200
         
     except Exception as e: