Create main.py

main.py

@@ -0,0 +1,278 @@
+import os
+import io
+import json
+from datetime import datetime
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+import google.generativeai as genai
+import firebase_admin
+from firebase_admin import credentials, db, storage, auth
+
+# Initialize Flask app and CORS
+app = Flask(__name__)
+CORS(app)
+
+# Firebase initialization
+cred = credentials.Certificate('sozo-daac1-firebase-adminsdk-fbsvc-b5d5c23b6d.json')
+firebase_admin.initialize_app(cred, {
+    'databaseURL': 'https://sozo-daac1-default-rtdb.firebaseio.com',
+    'storageBucket': 'sozo-daac1.firebasestorage.app'
+})
+bucket = storage.bucket()
+
+# Gemini API initialization
+api_key = os.environ['Gemini']
+def configure_gemini():
+    genai.configure(api_key=api_key)
+    return genai.GenerativeModel('gemini-2.0-flash-thinking-exp')
+
+# Helper functions
+def verify_token(token):
+    try:
+        decoded_token = auth.verify_id_token(token)
+        return decoded_token['uid']
+    except Exception as e:
+        return None
+
+def verify_admin(auth_header):
+    if not auth_header or not auth_header.startswith('Bearer '):
+        raise ValueError('Invalid token')
+    token = auth_header.split(' ')[1]
+    uid = verify_token(token)
+    if not uid:
+        raise PermissionError('Invalid user')
+    user_ref = db.reference(f'users/{uid}')
+    user_data = user_ref.get()
+    if not user_data or not user_data.get('is_admin', False):
+        raise PermissionError('Admin access required')
+    return uid
+
+# ---------- Dummy Admin Creation on Startup ----------
+def create_dummy_admin():
+    admin_email = "rairo@test.com"
+    admin_password = "123456"
+    try:
+        # Try to get the user if it exists
+        admin_user = auth.get_user_by_email(admin_email)
+    except firebase_admin.auth.UserNotFoundError:
+        # Create the dummy admin if not found
+        admin_user = auth.create_user(email=admin_email, password=admin_password)
+    # Set or update admin record in the database
+    admin_ref = db.reference(f'users/{admin_user.uid}')
+    admin_data = admin_ref.get() or {}
+    if not admin_data.get('is_admin', False):
+        admin_ref.set({
+            'email': admin_email,
+            'credits': 9999,  # Optionally, give admin lots of credits
+            'is_admin': True,
+            'created_at': datetime.utcnow().isoformat()
+        })
+    print(f"Dummy admin ready: {admin_email}")
+
+# ---------- Authentication Endpoints ----------
+
+@app.route('/api/auth/signup', methods=['POST'])
+def signup():
+    try:
+        data = request.get_json()
+        email = data.get('email')
+        password = data.get('password')
+        if not email or not password:
+            return jsonify({'error': 'Email and password are required'}), 400
+
+        # Create user in Firebase Auth
+        user = auth.create_user(email=email, password=password)
+        # Set initial user data in the realtime database with 3 starting credits
+        user_ref = db.reference(f'users/{user.uid}')
+        user_data = {
+            'email': email,
+            'credits': 3,
+            'is_admin': False,
+            'created_at': datetime.utcnow().isoformat()
+        }
+        user_ref.set(user_data)
+        return jsonify({
+            'success': True,
+            'user': {
+                'uid': user.uid,
+                **user_data
+            }
+        }), 201
+    except Exception as e:
+        return jsonify({'error': str(e)}), 400
+
+# ---------- User Profile ----------
+
+@app.route('/api/user/profile', methods=['GET'])
+def get_user_profile():
+    try:
+        auth_header = request.headers.get('Authorization', '')
+        if not auth_header.startswith('Bearer '):
+            return jsonify({'error': 'Authorization header missing or malformed'}), 401
+        token = auth_header.split(' ')[1]
+        uid = verify_token(token)
+        if not uid:
+            return jsonify({'error': 'Invalid token'}), 401
+        user_record = auth.get_user(uid)
+        user_data = db.reference(f'users/{uid}').get()
+        return jsonify({
+            'uid': uid,
+            'email': user_record.email,
+            'credits': user_data.get('credits', 0),
+            'is_admin': user_data.get('is_admin', False),
+            'created_at': user_data.get('created_at')
+        })
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+# ---------- Video Generation Endpoint ----------
+
+@app.route('/api/video/generate', methods=['POST'])
+def generate_video():
+    try:
+        auth_header = request.headers.get('Authorization', '')
+        if not auth_header.startswith('Bearer '):
+            return jsonify({'error': 'Authorization header missing or malformed'}), 401
+        token = auth_header.split(' ')[1]
+        uid = verify_token(token)
+        if not uid:
+            return jsonify({'error': 'Invalid token'}), 401
+
+        # Get user data and check credits
+        user_ref = db.reference(f'users/{uid}')
+        user_data = user_ref.get()
+        if not user_data or user_data.get('credits', 0) < 1:
+            return jsonify({'error': 'Insufficient credits'}), 400
+
+        # Deduct one credit
+        new_credits = user_data.get('credits', 0) - 1
+        user_ref.update({'credits': new_credits})
+
+        # Get prompt and generate video (simulation using Gemini)
+        req_data = request.get_json()
+        prompt = req_data.get('prompt', '')
+        if not prompt:
+            return jsonify({'error': 'Prompt is required'}), 400
+
+        model = configure_gemini()
+        # Call to Gemini to generate video content (this is a simulation)
+        response = model.generate(prompt=prompt)
+        # For demonstration, we assume the response contains a 'result' field
+        generated_video = response.result if hasattr(response, 'result') else "Video content generated based on prompt."
+
+        return jsonify({
+            'success': True,
+            'video': generated_video,
+            'remaining_credits': new_credits
+        })
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+# ---------- Credit Request Endpoints ----------
+
+@app.route('/api/user/request-credits', methods=['POST'])
+def request_credits():
+    try:
+        auth_header = request.headers.get('Authorization', '')
+        if not auth_header.startswith('Bearer '):
+            return jsonify({'error': 'Authorization header missing or malformed'}), 401
+        token = auth_header.split(' ')[1]
+        uid = verify_token(token)
+        if not uid:
+            return jsonify({'error': 'Invalid token'}), 401
+
+        data = request.get_json()
+        requested_credits = data.get('requested_credits')
+        if requested_credits is None:
+            return jsonify({'error': 'requested_credits is required'}), 400
+
+        # Create a credit request entry
+        credit_request_ref = db.reference('credit_requests').push()
+        credit_request_ref.set({
+            'user_id': uid,
+            'requested_credits': requested_credits,
+            'status': 'pending',
+            'requested_at': datetime.utcnow().isoformat()
+        })
+        return jsonify({'success': True, 'request_id': credit_request_ref.key})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+# ---------- Admin Endpoints for Credit Requests ----------
+
+@app.route('/api/admin/credit_requests', methods=['GET'])
+def list_credit_requests():
+    try:
+        verify_admin(request.headers.get('Authorization', ''))
+        requests_ref = db.reference('credit_requests')
+        credit_requests = requests_ref.get() or {}
+        # Convert dict to list with id
+        requests_list = [{'id': req_id, **data} for req_id, data in credit_requests.items()]
+        return jsonify({'credit_requests': requests_list})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/api/admin/credit_requests/<string:request_id>', methods=['PUT'])
+def process_credit_request(request_id):
+    try:
+        admin_uid = verify_admin(request.headers.get('Authorization', ''))
+        req_ref = db.reference(f'credit_requests/{request_id}')
+        req_data = req_ref.get()
+        if not req_data:
+            return jsonify({'error': 'Credit request not found'}), 404
+
+        data = request.get_json()
+        decision = data.get('decision')
+        if decision not in ['approved', 'declined']:
+            return jsonify({'error': 'decision must be "approved" or "declined"'}), 400
+
+        # If approved, add credits to the user
+        if decision == 'approved':
+            user_ref = db.reference(f'users/{req_data["user_id"]}')
+            user_data = user_ref.get()
+            if not user_data:
+                return jsonify({'error': 'User not found'}), 404
+            new_total = user_data.get('credits', 0) + float(req_data.get('requested_credits', 0))
+            user_ref.update({'credits': new_total})
+            req_ref.update({
+                'status': 'approved',
+                'processed_by': admin_uid,
+                'processed_at': datetime.utcnow().isoformat()
+            })
+            return jsonify({'success': True, 'new_user_credits': new_total})
+        else:
+            req_ref.update({
+                'status': 'declined',
+                'processed_by': admin_uid,
+                'processed_at': datetime.utcnow().isoformat()
+            })
+            return jsonify({'success': True, 'message': 'Credit request declined'})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+# ---------- Admin Endpoint to Directly Update Credits ----------
+@app.route('/api/admin/users/<string:uid>/credits', methods=['PUT'])
+def admin_update_credits(uid):
+    try:
+        verify_admin(request.headers.get('Authorization', ''))
+        data = request.get_json()
+        add_credits = data.get('add_credits')
+        if add_credits is None:
+            return jsonify({'error': 'add_credits is required'}), 400
+
+        user_ref = db.reference(f'users/{uid}')
+        user_data = user_ref.get()
+        if not user_data:
+            return jsonify({'error': 'User not found'}), 404
+
+        new_total = user_data.get('credits', 0) + float(add_credits)
+        user_ref.update({'credits': new_total})
+        return jsonify({'success': True, 'new_total_credits': new_total})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+# ---------- Main ----------
+if __name__ == '__main__':
+    # Create dummy admin account if it doesn't exist
+    create_dummy_admin()
+    app.run(debug=True, host="0.0.0.0", port=7860)