Update app.py

app.py

@@ -23,55 +23,69 @@ from wordpress_xmlrpc import Client
 from wordpress_xmlrpc.compat import xmlrpc_client
 from wordpress_xmlrpc.methods import media
 
-def get_user_id():
-    url = 'https://songlabai.com/wp-json/custom-api/v1/current-user-id-session'
-    try:
-        response = requests.get(url)
-        if response.status_code == 200:
-            data = response.json()
-            if 'user_id' in data:
-                return data['user_id']
-            else:
-                return None
-        else:
-            return None
-    except Exception as e:
-        print(f"Error getting user ID: {e}")
-        return None
+# Define the same secret key used in WordPress
+SECRET_KEY = '722C3AFB1F60816066AFB6045098E7016D98D1FCD5AFF4D3DE132CC1C6A141C9'  # Must match the key used in WordPress
+
+def get_user_info_from_jwt():
+    # Get the token from the URL parameters
+    query_params = st.experimental_get_query_params()
+    token = query_params.get('token', [None])[0]
+
+    if token:
+        try:
+            # Decode the token
+            payload = jwt.decode(token, SECRET_KEY, algorithms=['HS256'])
+            user_id = payload['data']['user']['id']
+            user_email = payload['data']['user']['email']
+            username = payload['data']['user']['username']
+            return {'user_id': user_id, 'email': user_email, 'username': username, 'token': token}
+        except ExpiredSignatureError:
+            st.error("Session expired. Please refresh the page and log in again.")
+        except InvalidTokenError:
+            st.error("Invalid authentication token. Please log in again.")
+    else:
+        st.error("No authentication token found. Please log in.")
+    return None
 
-def get_subscription_info(user_id):
-    url = f'https://songlabai.com/wp-json/custom-api/subscription/{user_id}'
+def get_subscription_info(token):
+    url = 'https://songlabai.com/wp-json/custom-api/subscription'
+    headers = {'Authorization': f'Bearer {token}'}
     try:
-        response = requests.get(url)
+        response = requests.get(url, headers=headers)
         if response.status_code == 200:
             data = response.json()
             return data
         else:
+            st.error(f"Failed to get subscription info. Status code: {response.status_code}")
             return None
     except Exception as e:
-        print(f"Error getting subscription info: {e}")
+        st.error(f"Error getting subscription info: {e}")
         return None
 
-def update_user_limit(user_id):
-    url = f'https://songlabai.com/wp-json/custom-api/update-limit/{user_id}'
+def update_user_limit(token):
+    url = 'https://songlabai.com/wp-json/custom-api/update-limit'
+    headers = {'Authorization': f'Bearer {token}'}
     try:
-        response = requests.post(url)
+        response = requests.post(url, headers=headers)
         if response.status_code == 200:
             data = response.json()
             if data.get('success') == True:
                 print("User limit updated successfully.")
                 return True
             else:
-                print(f"Failed to update user limit: {data}")
+                st.error(f"Failed to update user limit: {data}")
                 return False
         else:
-            print(f"Failed to update user limit, status code: {response.status_code}")
+            st.error(f"Failed to update user limit, status code: {response.status_code}")
             return False
     except Exception as e:
-        print(f"Error updating user limit: {e}")
+        st.error(f"Error updating user limit: {e}")
         return False
 
 
+
+
+
 # Try to get API_URL from environment variables, if not found set to a default value
 try:
     API_URL = os.environ["API_URL"]
@@ -350,6 +364,46 @@ def generate_audio(genre, energy_level, tempo, description, duration):
         )
         return
 
+    # After successful response, perform API checks
+
+    # Get user info from JWT
+    user_info = get_user_info_from_jwt()
+    if user_info is None:
+        st.stop()
+    else:
+        user_id = user_info['user_id']
+        token = user_info['token']
+
+    # Check subscription via API
+    subscription_info = get_subscription_info(token)
+    if not subscription_info:
+        st.error("Could not retrieve subscription information.")
+        return
+
+    subscription_plan_id = subscription_info['subscription_plan_id']
+    has_exceeded_daily_limit = subscription_info['has_exceeded_daily_limit']
+
+    if subscription_plan_id == '576':
+        if has_exceeded_daily_limit == True:
+            st.error("You have reached your daily limit for audio generation.")
+            return
+        else:
+            # Allow the user to proceed
+            placeholder1.success(f"✔️ Audio Generated, Loading...")
+            load_and_play_generated_audio(response)
+            placeholder1.empty()
+            # After generating audio, update the limit via API
+            update_user_limit(token)
+    elif subscription_plan_id in ['304', '305', '306']:
+        # Unlimited generation
+        placeholder1.success(f"✔️ Audio Generated, Loading...")
+        load_and_play_generated_audio(response)
+        placeholder1.empty()
+    else:
+        st.error("Your subscription does not allow audio generation.")
+        return
+
+
     # After successful response (after retries if necessary), perform API checks
 
     # Get user ID