Update auth.py

auth.py

@@ -2,41 +2,109 @@ import json
 from huggingface_hub import HfApi, hf_hub_download
 from werkzeug.security import generate_password_hash, check_password_hash
 import gradio as gr
-from phone_verification import is_phone_available, send_dummy_otp, verify_otp, record_phone_in_user, record_cooldown, get_total_accounts
+from phone_verification import is_phone_available, verify_otp, record_phone_in_user, record_cooldown, get_total_accounts
 
+# ================== REGISTER ==================
 def register_user(username, password, phone, otp, logged_in_username, auth_globals):
     hf_token = auth_globals["hf_token"]
     user_repo = auth_globals["user_repo"]
     hf_api = HfApi(token=hf_token)
 
     if get_total_accounts(auth_globals) >= 100:
-        return [], None, [], gr.update(choices=[], value=None), logged_in_username, "❌ Maximum 100 accounts reached.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
+        return [], None, [], gr.update(choices=[], value=None), logged_in_username, "❌ Maximum 100 accounts reached for this project.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
 
-    if not username or not password or not phone:
+    if not username or not password or not phone or not otp:
         return [], None, [], gr.update(choices=[], value=None), logged_in_username, "Please fill all fields.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
 
     if not is_phone_available(phone, auth_globals):
-        return [], None, [], gr.update(choices=[], value=None), logged_in_username, "❌ This phone number is already used or in cooldown.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
+        return [], None, [], gr.update(choices=[], value=None), logged_in_username, "❌ This phone number is already used or in 3-hour cooldown.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
 
     if not verify_otp(phone, otp, auth_globals):
-        return [], None, [], gr.update(choices=[], value=None), logged_in_username, "❌ Wrong OTP.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
+        return [], None, [], gr.update(choices=[], value=None), logged_in_username, "❌ Wrong OTP. Use 123456 for testing.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
 
     try:
+        # Check if username exists
         try:
             hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token)
             return [], None, [], gr.update(choices=[], value=None), logged_in_username, "User already exists!", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
         except:
             pass
 
-        user_data = {"password": generate_password_hash(password), "chats": [], "phone": phone}
-        hf_api.upload_file(path_or_fileobj=json.dumps(user_data).encode('utf-8'), path_in_repo=f"{username}.json", repo_id=user_repo, repo_type="dataset", token=hf_token)
+        user_data = {
+            "password": generate_password_hash(password),
+            "chats": [],
+            "phone": phone
+        }
+        hf_api.upload_file(
+            path_or_fileobj=json.dumps(user_data).encode('utf-8'),
+            path_in_repo=f"{username}.json",
+            repo_id=user_repo,
+            repo_type="dataset",
+            token=hf_token
+        )
+
+        record_phone_in_user(username, phone, auth_globals)
 
         return [], None, [], gr.update(choices=[], value=None), username, "✅ Registration successful! Logged in.", gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
     except Exception as e:
         return [], None, [], gr.update(choices=[], value=None), logged_in_username, f"Registration failed: {str(e)}", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
 
-# ... (login_user, logout_user, save_chats, delete_account remain the same, but delete_account now calls record_cooldown if phone exists)
+# ================== LOGIN ==================
+def login_user(username, password, logged_in_username, auth_globals):
+    hf_token = auth_globals["hf_token"]
+    user_repo = auth_globals["user_repo"]
+    hf_api = HfApi(token=hf_token)
+
+    if not username or not password:
+        return [], None, [], gr.update(choices=[], value=None), logged_in_username, "Please enter username and password.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
+
+    try:
+        user_file = hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token)
+        with open(user_file, 'r') as f:
+            data = json.load(f)
+
+        if check_password_hash(data["password"], password):
+            chats = data.get("chats", [])
+            if chats:
+                current_title = chats[-1]["title"]
+                history = chats[-1]["history"]
+                choices = [chat["title"] for chat in chats]
+                value = current_title
+            else:
+                current_title = None
+                history = []
+                choices = []
+                value = None
+
+            dropdown = gr.update(choices=choices, value=value)
+            return chats, current_title, history, dropdown, username, "✅ Login successful!", gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
+        else:
+            return [], None, [], gr.update(choices=[], value=None), logged_in_username, "❌ Incorrect password.", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
+    except Exception as e:
+        return [], None, [], gr.update(choices=[], value=None), logged_in_username, f"Login failed: {str(e)}", gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False)
+
+# ================== LOGOUT ==================
+def logout_user():
+    return [], None, [], gr.update(choices=[], value=None), None, "Logged out successfully!", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+# ================== SAVE CHATS ==================
+def save_chats(username, chats, auth_globals):
+    if not username:
+        return "No user logged in!"
+    hf_token = auth_globals["hf_token"]
+    user_repo = auth_globals["user_repo"]
+    hf_api = HfApi(token=hf_token)
+    try:
+        user_file = hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token)
+        with open(user_file, 'r') as f:
+            data = json.load(f)
+        data["chats"] = chats
+        hf_api.upload_file(path_or_fileobj=json.dumps(data).encode('utf-8'), path_in_repo=f"{username}.json", repo_id=user_repo, repo_type="dataset", token=hf_token)
+        return "Chats saved automatically!"
+    except Exception as e:
+        return f"Failed to save chats: {str(e)}"
 
+# ================== DELETE ACCOUNT ==================
 def delete_account(confirm_username, logged_in_username, auth_globals):
     if not logged_in_username or confirm_username != logged_in_username:
         return "❌ Username does not match.", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None
@@ -45,7 +113,7 @@ def delete_account(confirm_username, logged_in_username, auth_globals):
     user_repo = auth_globals["user_repo"]
     hf_api = HfApi(token=hf_token)
     try:
-        # Get phone before delete
+        # Get phone before deleting
         try:
             file = hf_hub_download(repo_id=user_repo, filename=f"{logged_in_username}.json", repo_type="dataset", token=hf_token)
             with open(file, 'r') as f:
@@ -58,6 +126,6 @@ def delete_account(confirm_username, logged_in_username, auth_globals):
             pass
 
         hf_api.delete_file(path_in_repo=f"{logged_in_username}.json", repo_id=user_repo, repo_type="dataset", token=hf_token)
-        return "✅ Account and all data permanently deleted.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), [], None, [], gr.update(choices=[], value=None)
+        return "✅ Account permanently deleted.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), [], None, [], gr.update(choices=[], value=None)
     except Exception as e:
         return f"❌ Deletion failed: {str(e)}", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None