Update main.py

main.py

@@ -2,10 +2,10 @@ import os
 import io
 import json
 import hashlib
-from datetime import datetime, time
+from datetime import datetime, time, timedelta
 from PIL import Image
 import pytz
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
 import google.generativeai as genai
 import firebase_admin
@@ -13,7 +13,6 @@ from firebase_admin import credentials, db, storage, auth
 import pandas as pd
 import requests
 from urllib.parse import urlparse, unquote
-from datetime import datetime, timedelta
 
 app = Flask(__name__)
 CORS(app)
@@ -29,8 +28,6 @@ bucket = storage.bucket()
 api_key = os.environ['Gemini']
 FIREBASE_API_KEY = os.environ.get('FIREBASE_API_KEY')
 
-
-
 # Helper functions
 def configure_gemini():
     genai.configure(api_key=api_key)
@@ -74,7 +71,7 @@ def process_receipt(model, image):
         response = model.generate_content([prompt, image])
         return response.text
     except Exception as e:
-        st.error(f"Gemini error: {str(e)}")
+        print(f"Gemini error: {str(e)}")
         return "{}"
 
 
@@ -86,22 +83,18 @@ def generate_report():
     Make sure the report is in plain text"""
     model = configure_gemini()
     try:
-        # Get the transaction data from the request body
-        
         transaction_data = request.get_json()
         transaction_json_string = json.dumps(transaction_data['transactions'])
-        
-        # Assuming model is defined/imported elsewhere in your Flask app
         response = model.generate_content([prompt, transaction_json_string])
         report = response.text
-        return jsonify({"report": report})  # Wrap the report in an object
-        
+        return jsonify({"report": report})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
+
 # ========================================
 # Authentication Endpoints
 # ========================================
-
+# (Any existing authentication endpoints remain unchanged)
 
 # ========================================
 # Receipt Processing Endpoint (unchanged)
@@ -161,7 +154,7 @@ def process_receipt_endpoint():
         )
 
     except Exception as e:
-        print(image)
+        print(e)
         return jsonify({'error': str(e)}), 500
 
 def handle_manual_entry(uid, user_ref, user_data):
@@ -190,15 +183,13 @@ def validate_and_save_transaction(uid, user_data, data, file_hash, image_bytes,
     receipt_number = data.get('receipt_number')
     items = data.get('items', [])
 
-    # Fetch existing transactions with the same receipt_number
+    # Check for duplicate transactions based on receipt number and items.
     existing_transactions_with_receipt = transactions_ref.order_by_child('receipt_number').equal_to(receipt_number).get()
 
     if existing_transactions_with_receipt:
         for transaction_id, existing_transaction_data in existing_transactions_with_receipt.items():
-            # Compare items lists - assuming order doesn't matter, so sorting for comparison
             existing_items = sorted(existing_transaction_data.get('items', []))
             current_items = sorted(items)
-
             if existing_items == current_items:
                 return jsonify({'error': f"Transaction with Receipt #{receipt_number} and identical items already exists"}), 400
 
@@ -218,7 +209,6 @@ def validate_and_save_transaction(uid, user_data, data, file_hash, image_bytes,
     new_remaining = user_data['remaining_cash'] - total
     db.reference(f'users/{uid}').update({'remaining_cash': new_remaining})
 
-    # Save transaction
     transaction_data = {
         'uid': uid,
         'total': total,
@@ -242,7 +232,6 @@ def validate_and_save_transaction(uid, user_data, data, file_hash, image_bytes,
 # Data Endpoints for Visualizations
 # ========================================
 
-
 @app.route('/api/user/spending-overview', methods=['GET'])
 def get_spending_overview():
     try:
@@ -256,17 +245,13 @@ def get_spending_overview():
                 'recent_transactions': []
             })
         
-        # Create dataframe
         df = pd.DataFrame.from_dict(transactions, orient='index')
-        
-        # Check if dataframe is empty
         if df.empty:
             return jsonify({
                 'daily_spending': [],
                 'recent_transactions': []
             })
         
-        # Check if all the data is here, otherwise we delete the data.
         df.dropna(subset=['uid','total','items','date', 'receipt_number', 'timestamp','hash', 'manual_entry'], inplace=True)
         if df.empty:
             return jsonify({
@@ -293,9 +278,6 @@ def get_spending_overview():
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
-
-# ... (previous imports remain the same)
-
 # ========================================
 # Modified verify_admin function (now checks database is_admin flag)
 # ========================================
@@ -313,25 +295,23 @@ def verify_admin(auth_header):
     if not user_data or not user_data.get('is_admin', False):
         raise PermissionError('Admin access required')
 
-    # **Set custom claim here AFTER admin verification is successful:**
     try:
         auth.set_custom_user_claims(uid, {"admin": True})
-        print(f"Custom admin claim set for user {uid}") # Optional log
+        print(f"Custom admin claim set for user {uid}")
     except Exception as e:
-        print(f"Error setting custom admin claim: {e}") # Log any errors
-        raise PermissionError('Error setting admin claim, but admin verified') # Or handle error as needed
+        print(f"Error setting custom admin claim: {e}")
+        raise PermissionError('Error setting admin claim, but admin verified')
 
-    return uid # Return uid as before
+    return uid
 
 # ========================================
-# New Admin Endpoints
+# Existing Admin Endpoints
 # ========================================
 @app.route('/api/admin/overview', methods=['GET'])
 def get_admin_overview():
     try:
         verify_admin(request.headers.get('Authorization', ''))
         
-        # Get all users
         users_ref = db.reference('users')
         all_users = users_ref.get() or {}
         users_list = []
@@ -350,7 +330,6 @@ def get_admin_overview():
                 'is_admin': user_data.get('is_admin', False)
             })
 
-        # Get all transactions
         transactions_ref = db.reference('transactions')
         all_transactions = transactions_ref.get() or {}
         transactions_list = [{'id': tid, **data} for tid, data in all_transactions.items()]
@@ -373,18 +352,16 @@ def create_user():
         verify_admin(request.headers.get('Authorization', ''))
         data = request.get_json()
         
-        # Create Firebase auth user
         user = auth.create_user(
             email=data['email'],
             password=data['password']
         )
         
-        # Create database user record
         user_ref = db.reference(f'users/{user.uid}')
         user_data = {
             'daily_cash': data.get('daily_cash', 100),
             'remaining_cash': data.get('daily_cash', 100),
-            'last_reset': '2000-01-01T00:00:00+00:00',  # Force reset on next check
+            'last_reset': '2000-01-01T00:00:00+00:00',
             'is_admin': data.get('is_admin', False)
         }
         user_ref.set(user_data)
@@ -414,8 +391,6 @@ def update_user_limit(uid):
             return jsonify({'error': 'User not found'}), 404
             
         updates = {'daily_cash': new_limit}
-        
-        # Adjust remaining cash if it exceeds new limit
         current_remaining = user_data.get('remaining_cash', new_limit)
         if current_remaining > new_limit:
             updates['remaining_cash'] = new_limit
@@ -430,9 +405,85 @@ def update_user_limit(uid):
     except Exception as e:
         return jsonify({'error': str(e)}), 400
 
+# ========================================
+# New Admin Endpoints for Updating Remaining Cash, Setting Cash Limits, and Resetting Password
+# ========================================
 
+@app.route('/api/admin/users/<string:uid>/remaining-cash', methods=['PUT'])
+def update_remaining_cash(uid):
+    try:
+        verify_admin(request.headers.get('Authorization', ''))
+        data = request.get_json()
+        if 'remaining_cash' not in data:
+            return jsonify({'error': 'remaining_cash is required'}), 400
+        new_remaining_cash = float(data['remaining_cash'])
+        user_ref = db.reference(f'users/{uid}')
+        user_data = user_ref.get()
+        if not user_data:
+            return jsonify({'error': 'User not found'}), 404
+
+        user_ref.update({'remaining_cash': new_remaining_cash})
+        return jsonify({'success': True, 'remaining_cash': new_remaining_cash})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 400
 
+@app.route('/api/admin/users/<string:uid>/cash-limit-period', methods=['PUT'])
+def update_cash_limit_period(uid):
+    try:
+        verify_admin(request.headers.get('Authorization', ''))
+        data = request.get_json()
+        start_date = data.get('start_date')
+        end_date = data.get('end_date')
+        cash_limit = data.get('cash_limit')
+        if not (start_date and end_date and cash_limit is not None):
+            return jsonify({'error': 'start_date, end_date, and cash_limit are required'}), 400
+        try:
+            datetime.strptime(start_date, '%Y-%m-%d')
+            datetime.strptime(end_date, '%Y-%m-%d')
+        except Exception as e:
+            return jsonify({'error': 'Invalid date format. Use YYYY-MM-DD for start_date and end_date'}), 400
+        
+        user_ref = db.reference(f'users/{uid}')
+        user_data = user_ref.get()
+        if not user_data:
+            return jsonify({'error': 'User not found'}), 404
+        
+        user_ref.update({
+            'cash_limit_period': {
+                'start_date': start_date,
+                'end_date': end_date,
+                'limit': float(cash_limit)
+            }
+        })
+
+        return jsonify({
+            'success': True,
+            'cash_limit_period': {
+                'start_date': start_date,
+                'end_date': end_date,
+                'limit': float(cash_limit)
+            }
+        })
+    except Exception as e:
+        return jsonify({'error': str(e)}), 400
+
+@app.route('/api/admin/users/<string:uid>/reset-password', methods=['PUT'])
+def admin_reset_password(uid):
+    try:
+        verify_admin(request.headers.get('Authorization', ''))
+        data = request.get_json()
+        new_password = data.get('new_password')
+        if not new_password:
+            return jsonify({'error': 'new_password is required'}), 400
+        auth.update_user(uid, password=new_password)
+        return jsonify({'success': True, 'message': 'Password reset successfully'})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 400
+
+# ========================================
 # User management endpoint for profile
+# ========================================
+
 @app.route('/api/user/profile', methods=['GET'])
 def get_user_profile():
     try:
@@ -454,26 +505,17 @@ def get_user_profile():
         user_data = db.reference(f'users/{uid}').get()
         return jsonify({'error': str(e)+ f'user data: {user_data}'}), 500
 
+# ========================================
 # Receipt media endpoints
-
+# ========================================
 
 def get_blob_from_image_url(image_url):
-    """
-    Extracts and returns the blob path from a Firebase Storage URL.
-    Supports URLs of the form:
-      https://storage.googleapis.com/<bucket>/<blob_path>
-    or the default Firebase URL format.
-    """
     parsed = urlparse(image_url)
-    # If URL is in the storage.googleapis.com format:
     if parsed.netloc == "storage.googleapis.com":
-        # Expected path: /<bucket>/<blob_path>
         parts = parsed.path.strip("/").split("/", 1)
         if len(parts) == 2:
             bucket_name, blob_path = parts
-            # (Optionally, check if bucket_name matches our bucket.name)
             return blob_path
-    # Otherwise, try the default format:
     prefix = f"/v0/b/{bucket.name}/o/"
     if parsed.path.startswith(prefix):
         encoded_blob_path = parsed.path[len(prefix):]
@@ -498,17 +540,14 @@ def view_receipt(transaction_id):
         if not blob_path:
             return jsonify({'error': 'Could not determine blob path from URL'}), 500
 
-        print(f"Blob path for view: {blob_path}")  # Debug log
-
+        print(f"Blob path for view: {blob_path}")
         blob = bucket.blob(blob_path)
         if not blob.exists():
             print("Blob does not exist at path:", blob_path)
             return jsonify({'error': 'Blob not found'}), 404
 
         signed_url = blob.generate_signed_url(expiration=timedelta(minutes=10))
-        print(f"Signed URL for view: {signed_url}")  # Debug log
         r = requests.get(signed_url)
-        print(f"View endpoint response status: {r.status_code}")  # Debug log
         if r.status_code != 200:
             return jsonify({'error': 'Unable to fetch image from storage'}), 500
 
@@ -534,17 +573,14 @@ def download_receipt(transaction_id):
         if not blob_path:
             return jsonify({'error': 'Could not determine blob path from URL'}), 500
 
-        print(f"Blob path for download: {blob_path}")  # Debug log
-
+        print(f"Blob path for download: {blob_path}")
         blob = bucket.blob(blob_path)
         if not blob.exists():
             print("Blob does not exist at path:", blob_path)
             return jsonify({'error': 'Blob not found'}), 404
 
         signed_url = blob.generate_signed_url(expiration=timedelta(minutes=10))
-        print(f"Signed URL for download: {signed_url}")  # Debug log
         r = requests.get(signed_url)
-        print(f"Download endpoint response status: {r.status_code}")  # Debug log
         if r.status_code != 200:
             return jsonify({'error': 'Unable to fetch image from storage'}), 500
 
@@ -552,34 +588,29 @@ def download_receipt(transaction_id):
             io.BytesIO(r.content),
             mimetype='image/jpeg',
             as_attachment=True,
-            attachment_filename='receipt.jpg'  # Use download_name if using Flask 2.x
+            attachment_filename='receipt.jpg'
         )
     except Exception as e:
         print(f"Download receipt error: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
-
-
-# delete users
+# ========================================
+# Delete user endpoint
+# ========================================
 
 @app.route('/api/admin/users/<string:uid>', methods=['DELETE'])
 def delete_user(uid):
     try:
         verify_admin(request.headers.get('Authorization', ''))
         
-        # Verify the user exists
         try:
             user = auth.get_user(uid)
         except auth.UserNotFoundError:
             return jsonify({'error': 'User not found'}), 404
 
-        # Delete authentication user
         auth.delete_user(uid)
-        
-        # Delete database record
         db.reference(f'users/{uid}').delete()
         
-        # Delete user's transactions
         transactions_ref = db.reference('transactions')
         user_transactions = transactions_ref.order_by_child('uid').equal_to(uid).get()
         if user_transactions:
@@ -593,6 +624,5 @@ def delete_user(uid):
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
-
 if __name__ == '__main__':
-    app.run(debug=True, host="0.0.0.0", port=7860)
+    app.run(debug=True, host="0.0.0.0", port=7860)