api/services/auth_service.py

"""
Authentication service for user management with HF bucket storage.
Handles registration, login, password encryption, JWT tokens, and user CRUD.
"""

import json
import uuid
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
import jwt
import hashlib
from huggingface_hub import HfApi
from api.config import settings


class AuthService:
    """Manages user authentication and storage with HF buckets"""

    def __init__(self):
        self.api = HfApi()
        self.repo_id = settings.bucket.repo_id
        self.repo_type = settings.bucket.repo_type
        self.users_file = settings.bucket.users_file
        self.conversations_base = settings.bucket.conversations_base
        self.secret_key = settings.secret_key
        self.algorithm = settings.jwt_algorithm
        self.jwt_expiration_hours = settings.jwt_expiration_hours

    # ── Password Hashing ────────────────────────────────────────────────────────

    def hash_password(self, password: str) -> str:
        """Hash password using SHA-256 with salt (simple version)"""
        # In production, use bcrypt or argon2
        salt = settings.secret_key[:16]
        return hashlib.pbkdf2_hmac(
            'sha256',
            password.encode('utf-8'),
            salt.encode('utf-8'),
            100000
        ).hex()

    def verify_password(self, password: str, hash_value: str) -> bool:
        """Verify password against hash"""
        return self.hash_password(password) == hash_value

    # ── JWT Token Management ───────────────────────────────────────────────────

    def create_access_token(self, user_id: str, email: str) -> str:
        """Create JWT access token"""
        payload = {
            "sub": user_id,
            "email": email,
            "iat": datetime.utcnow(),
            "exp": datetime.utcnow() + timedelta(hours=self.jwt_expiration_hours)
        }
        return jwt.encode(payload, self.secret_key, algorithm=self.algorithm)

    def verify_token(self, token: str) -> Optional[Dict[str, Any]]:
        """Verify JWT token and return payload"""
        try:
            payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])
            return payload
        except jwt.InvalidTokenError:
            return None

    # ── User CRUD Operations ──────────────────────────────────────────────────

    def _read_users_file(self) -> list:
        """Read users.jsonl file from bucket"""
        try:
            content = self.api.read_file_content(
                repo_id=self.repo_id,
                repo_type=self.repo_type,
                filename=self.users_file
            )
            users = []
            for line in content.strip().split('\n'):
                if line:
                    users.append(json.loads(line))
            return users
        except Exception:
            # File doesn't exist yet, return empty list
            return []

    def _write_users_file(self, users: list) -> bool:
        """Write users.jsonl file to bucket"""
        try:
            content = '\n'.join(json.dumps(user) for user in users)
            self.api.upload_file(
                path_or_fileobj=content.encode('utf-8'),
                path_in_repo=self.users_file,
                repo_id=self.repo_id,
                repo_type=self.repo_type,
                commit_message="Update users database"
            )
            return True
        except Exception as e:
            print(f"Error writing users file: {e}")
            return False

    def register_user(
            self,
            email: str, 
            password: str,
            full_name: Optional[str] = None,
            dob: Optional[str] = "",
            sex: Optional[str] = "",
            medical_history: Optional[str] = "",
            geolocation: Optional[str] = ""
        ) -> tuple[bool, str, Optional[str]]:
        """
        Register new user
        Returns: (success, message, user_id)
        """
        users = self._read_users_file()

        # Check if user exists
        if any(u['email'] == email for u in users):
            return False, "Email already registered", None

        # Create new user
        user_id = str(uuid.uuid4())
        user = {
            "user_id": user_id,
            "email": email,
            "password_hash": self.hash_password(password),
            "full_name": full_name,
            "dob": dob,
            "sex": sex,
            "medical_history": medical_history,
            "geolocation": geolocation,
            "created_at": datetime.utcnow().isoformat(),
            "conversation_ids": []
        }

        users.append(user)
        if self._write_users_file(users):
            # Create user's conversation directory
            self._create_user_conversation_dir(user_id)
            return True, "User registered successfully", user_id
        else:
            return False, "Failed to save user to database", None

    def login_user(self, email: str, password: str) -> tuple[bool, str, Dict]:
        """
        Login user
        Returns: (success, message, user_data)
        """
        users = self._read_users_file()
        user = next((u for u in users if u['email'] == email), {})

        if not user:
            return False, "Email not found", {}

        if not self.verify_password(password, user['password_hash']):
            return False, "Invalid password", {}

        # Return user data (without password hash)
        user_data = {
            "user_id": user["user_id"],
            "email": user["email"],
            "full_name": user.get("full_name"),
            "created_at": user["created_at"],
            "dob": user.get("dob"),
            "sex": user.get("sex"),
            "medical_history": user.get("medical_history"),
            "geolocation": user.get("geolocation"),
            "conversation_count": len(user.get("conversation_ids", []))
        }
        return True, "Login successful", user_data

    def get_user_by_id(self, user_id: str) -> Optional[Dict]:
        """Get user information by ID"""
        users = self._read_users_file()
        user = next((u for u in users if u['user_id'] == user_id), None)

        if not user:
            return None

        return {
            "user_id": user["user_id"],
            "email": user["email"],
            "full_name": user.get("full_name"),
            "created_at": user["created_at"],
            "dob": user.get("dob"),
            "sex": user.get("sex"),
            "medical_history": user.get("medical_history"),
            "geolocation": user.get("geolocation"),
            "conversation_count": len(user.get("conversation_ids", []))
        }

    def update_user_profile(self, user_id: str, email: Optional[str] = None,
                            full_name: Optional[str] = None,
                            dob: Optional[str] = None,
                            sex: Optional[str] = None,
                            medical_history: Optional[str] = None,
                            geolocation: Optional[str] = None) -> tuple[bool, str]:
        """Update user profile information"""
        users = self._read_users_file()
        user = next((u for u in users if u['user_id'] == user_id), None)

        if not user:
            return False, "User not found"

        if email:
            # Check if email is already taken
            if any(u['email'] == email and u['user_id'] != user_id for u in users):
                return False, "Email already in use"
            user['email'] = email

        if full_name:
            user['full_name'] = full_name
        if dob:
            user['dob'] = dob
        if sex:
            user['sex'] = sex
        if medical_history:
            user['medical_history'] = medical_history
        if geolocation:
            user['geolocation'] = geolocation

        if self._write_users_file(users):
            return True, "Profile updated successfully"
        else:
            return False, "Failed to update profile"

    def change_password(self, user_id: str, old_password: str,
                        new_password: str) -> tuple[bool, str]:
        """Change user password"""
        users = self._read_users_file()
        user = next((u for u in users if u['user_id'] == user_id), None)

        if not user:
            return False, "User not found"

        if not self.verify_password(old_password, user['password_hash']):
            return False, "Current password is incorrect"

        user['password_hash'] = self.hash_password(new_password)

        if self._write_users_file(users):
            return True, "Password changed successfully"
        else:
            return False, "Failed to change password"

    def delete_user(self, user_id: str, password: str) -> tuple[bool, str]:
        """Delete user account and all associated data"""
        users = self._read_users_file()
        user = next((u for u in users if u['user_id'] == user_id), None)

        if not user:
            return False, "User not found"

        if not self.verify_password(password, user['password_hash']):
            return False, "Invalid password"

        # Remove user from users file
        users = [u for u in users if u['user_id'] != user_id]

        if self._write_users_file(users):
            # TODO: Delete all user's conversations from bucket
            return True, "Account deleted successfully"
        else:
            return False, "Failed to delete account"

    # ── Conversation Management ────────────────────────────────────────────────

    def _create_user_conversation_dir(self, user_id: str):
        """Create a directory for user conversations"""
        # HF buckets auto-create directories with file uploads, so no explicit creation needed
        pass

    def create_conversation(self, user_id: str, title: str = "New Conversation") -> tuple[bool, str, Optional[str]]:
        """
        Create new conversation with max 10 limit per user
        Returns: (success, message, conversation_id)
        """
        users = self._read_users_file()
        user = next((u for u in users if u['user_id'] == user_id), None)

        if not user:
            return False, "User not found", None

        conversation_ids = user.get("conversation_ids", [])

        # Check conversation limit
        if len(conversation_ids) >= 10:
            return False, "Maximum 10 conversations per user. Please delete an old conversation.", None

        conversation_id = str(uuid.uuid4())
        user["conversation_ids"].append(conversation_id)

        if self._write_users_file(users):
            # Create conversation JSONL file
            conversation_data = {
                "conversation_id": conversation_id,
                "user_id": user_id,
                "title": title,
                "created_at": datetime.utcnow().isoformat(),
                "messages": []
            }
            self._save_conversation(user_id, conversation_id, conversation_data)
            return True, "Conversation created successfully", conversation_id
        else:
            return False, "Failed to create conversation", None

    def get_conversations(self, user_id: str) -> tuple[bool, str, Optional[list]]:
        """Get list of all user conversations"""
        users = self._read_users_file()
        user = next((u for u in users if u['user_id'] == user_id), None)

        if not user:
            return False, "User not found", None

        conversation_ids = user.get("conversation_ids", [])
        conversations = []

        for conv_id in conversation_ids:
            conv_data = self._load_conversation(user_id, conv_id)
            if conv_data:
                conversations.append({
                    "conversation_id": conv_id,
                    "title": conv_data.get("title", "Untitled"),
                    "created_at": conv_data.get("created_at"),
                    "last_updated": conv_data.get("last_updated", conv_data.get("created_at")),
                    "message_count": len(conv_data.get("messages", [])),
                    "has_diagnosis": "diagnosis" in conv_data,
                    "last_message_preview": conv_data.get("messages", [{}])[-1].get("content", "")[:50] if conv_data.get("messages") else ""
                })

        return True, "Conversations retrieved", conversations

    def _save_conversation(self, user_id: str, conversation_id: str, data: Dict):
        """Save conversation to bucket"""
        try:
            filename = f"{self.conversations_base}{user_id}/{conversation_id}.json"
            self.api.upload_file(
                path_or_fileobj=json.dumps(data, indent=2).encode('utf-8'),
                path_in_repo=filename,
                repo_id=self.repo_id,
                repo_type=self.repo_type,
                commit_message=f"Update conversation {conversation_id}"
            )
        except Exception as e:
            print(f"Error saving conversation: {e}")

    def _load_conversation(self, user_id: str, conversation_id: str) -> Optional[Dict]:
        """Load conversation from bucket"""
        try:
            filename = f"{self.conversations_base}{user_id}/{conversation_id}.json"
            content = self.api.read_file_content(
                repo_id=self.repo_id,
                repo_type=self.repo_type,
                filename=filename
            )
            return json.loads(content)
        except Exception:
            return None

    def delete_conversation(self, user_id: str, conversation_id: str) -> tuple[bool, str]:
        """Delete conversation from user's list"""
        users = self._read_users_file()
        user = next((u for u in users if u['user_id'] == user_id), None)

        if not user:
            return False, "User not found"

        if conversation_id not in user.get("conversation_ids", []):
            return False, "Conversation not found"

        user["conversation_ids"].remove(conversation_id)

        if self._write_users_file(users):
            # TODO: Delete conversation file from bucket
            return True, "Conversation deleted successfully"
        else:
            return False, "Failed to delete conversation"


# Singleton instance
auth_service = AuthService()