models.py

"""
User authentication models for MongoDB integration
"""

import re
from datetime import datetime
from typing import Optional, Dict, Any
from bson import ObjectId
from flask_login import UserMixin
from werkzeug.security import generate_password_hash, check_password_hash
from email_validator import validate_email, EmailNotValidError
import logging

from database import get_users_collection, get_chat_sessions_collection

logger = logging.getLogger(__name__)


class User(UserMixin):
    """User model with MongoDB integration and Flask-Login compatibility"""
    
    def __init__(self, user_data: Dict[str, Any]):
        """Initialize User from MongoDB document"""
        self.id = str(user_data['_id'])
        self.email = user_data['email']
        self.password_hash = user_data['password_hash']
        self.created_at = user_data.get('created_at', datetime.utcnow())
        self._is_active = user_data.get('is_active', True)
    
    @property
    def is_active(self):
        """Return user active status for Flask-Login"""
        return self._is_active
    
    def get_id(self):
        """Return user ID as string for Flask-Login"""
        return self.id
    
    @staticmethod
    def find_by_email(email: str) -> Optional['User']:
        """Find user by email in MongoDB"""
        try:
            users_collection = get_users_collection()
            user_data = users_collection.find_one({"email": email.lower()})
            
            if user_data:
                return User(user_data)
            return None
            
        except Exception as e:
            logger.error(f"Error finding user by email {email}: {e}")
            return None
    
    @staticmethod
    def find_by_id(user_id: str) -> Optional['User']:
        """Find user by ID in MongoDB"""
        try:
            users_collection = get_users_collection()
            
            # Convert string ID to ObjectId
            if isinstance(user_id, str):
                user_id = ObjectId(user_id)
            
            user_data = users_collection.find_one({"_id": user_id})
            
            if user_data:
                return User(user_data)
            return None
            
        except Exception as e:
            logger.error(f"Error finding user by ID {user_id}: {e}")
            return None
    
    @staticmethod
    def create_user(email: str, password: str) -> Optional['User']:
        """Create new user in MongoDB"""
        try:
            # Validate email format
            if not User.validate_email_format(email):
                logger.error(f"Invalid email format: {email}")
                return None
            
            # Validate password strength
            if not User.validate_password_strength(password):
                logger.error("Password does not meet strength requirements")
                return None
            
            # Check if user already exists
            if User.find_by_email(email):
                logger.error(f"User with email {email} already exists")
                return None
            
            users_collection = get_users_collection()
            
            # Create user document
            user_doc = {
                "email": email.lower(),
                "password_hash": generate_password_hash(password),
                "created_at": datetime.utcnow(),
                "is_active": True
            }
            
            # Insert user into database
            result = users_collection.insert_one(user_doc)
            
            if result.inserted_id:
                # Retrieve the created user
                user_data = users_collection.find_one({"_id": result.inserted_id})
                logger.info(f"Successfully created user with email: {email}")
                return User(user_data)
            
            return None
            
        except Exception as e:
            logger.error(f"Error creating user with email {email}: {e}")
            return None
    
    def set_password(self, password: str) -> bool:
        """Hash and set user password"""
        try:
            if not self.validate_password_strength(password):
                return False
            
            users_collection = get_users_collection()
            new_password_hash = generate_password_hash(password)
            
            result = users_collection.update_one(
                {"_id": ObjectId(self.id)},
                {"$set": {"password_hash": new_password_hash}}
            )
            
            if result.modified_count > 0:
                self.password_hash = new_password_hash
                logger.info(f"Password updated for user: {self.email}")
                return True
            
            return False
            
        except Exception as e:
            logger.error(f"Error setting password for user {self.email}: {e}")
            return False
    
    def check_password(self, password: str) -> bool:
        """Verify password against hash"""
        try:
            return check_password_hash(self.password_hash, password)
        except Exception as e:
            logger.error(f"Error checking password for user {self.email}: {e}")
            return False
    
    @staticmethod
    def validate_email_format(email: str) -> bool:
        """Validate email format using email-validator"""
        try:
            # Use email-validator library for comprehensive validation
            # Disable deliverability check for testing
            validate_email(email, check_deliverability=False)
            return True
        except EmailNotValidError:
            return False
        except Exception as e:
            logger.error(f"Error validating email {email}: {e}")
            return False
    
    @staticmethod
    def validate_password_strength(password: str) -> bool:
        """Validate password strength requirements"""
        try:
            # Check minimum length (8 characters)
            if len(password) < 8:
                return False
            
            # Additional strength checks can be added here
            # For now, just checking minimum length as per requirements
            return True
            
        except Exception as e:
            logger.error(f"Error validating password strength: {e}")
            return False
    
    def get_chat_history(self, limit: int = 50) -> list:
        """Get user's chat history from MongoDB"""
        try:
            chat_sessions_collection = get_chat_sessions_collection()
            
            # Query chat sessions for this user, sorted by timestamp (newest first)
            cursor = chat_sessions_collection.find(
                {"user_id": ObjectId(self.id)}
            ).sort("timestamp", -1).limit(limit)
            
            history = []
            for session in cursor:
                history.append({
                    "id": str(session["_id"]),
                    "message": session.get("message", ""),
                    "response": session.get("response", ""),
                    "timestamp": session.get("timestamp", datetime.utcnow()),
                    "session_data": session.get("session_data", {})
                })
            
            # Return in chronological order (oldest first)
            return list(reversed(history))
            
        except Exception as e:
            logger.error(f"Error getting chat history for user {self.email}: {e}")
            return []
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert user to dictionary (excluding password hash)"""
        return {
            "id": self.id,
            "email": self.email,
            "created_at": self.created_at,
            "is_active": self._is_active
        }
    
    def __repr__(self):
        return f"<User {self.email}>"


class ChatSession:
    """Chat session model for API server integration"""
    
    @staticmethod
    def save_message(user_id: str, message: str, response: str, session_metadata: Optional[Dict] = None) -> bool:
        """Save chat message to MongoDB with user_id"""
        try:
            chat_sessions_collection = get_chat_sessions_collection()
            
            session_doc = {
                "user_id": ObjectId(user_id),
                "message": message,
                "response": response,
                "timestamp": datetime.utcnow(),
                "session_data": session_metadata or {}
            }
            
            result = chat_sessions_collection.insert_one(session_doc)
            
            if result.inserted_id:
                logger.info(f"Chat message saved for user_id: {user_id}")
                return True
            
            return False
            
        except Exception as e:
            logger.error(f"Error saving chat message for user_id {user_id}: {e}")
            return False
    
    @staticmethod
    def get_user_history(user_id: str, limit: int = 50) -> list:
        """Retrieve user's chat history from MongoDB"""
        try:
            chat_sessions_collection = get_chat_sessions_collection()
            
            cursor = chat_sessions_collection.find(
                {"user_id": ObjectId(user_id)}
            ).sort("timestamp", 1).limit(limit)  # Chronological order
            
            history = []
            for session in cursor:
                history.append({
                    "id": str(session["_id"]),
                    "message": session.get("message", ""),
                    "response": session.get("response", ""),
                    "timestamp": session.get("timestamp", datetime.utcnow()),
                    "session_data": session.get("session_data", {})
                })
            
            return history
            
        except Exception as e:
            logger.error(f"Error getting user history for user_id {user_id}: {e}")
            return []
    
    @staticmethod
    def get_session_context(user_id: str, limit: int = 10) -> list:
        """Get formatted chat history for API context"""
        try:
            history = ChatSession.get_user_history(user_id, limit)
            
            # Format for API context (last N exchanges)
            context = []
            for session in history[-limit:]:  # Get most recent exchanges
                context.append({
                    "user": session["message"],
                    "assistant": session["response"]
                })
            
            return context
            
        except Exception as e:
            logger.error(f"Error getting session context for user_id {user_id}: {e}")
            return []
    
    @staticmethod
    def get_conversation_history_for_atlas(user_id: str, limit: int = 10, max_tokens: int = 8000) -> list:
        """Get conversation history formatted for Atlas API with token management"""
        try:
            chat_sessions_collection = get_chat_sessions_collection()
            
            # Get recent chat sessions, sorted by timestamp (newest first)
            cursor = chat_sessions_collection.find(
                {"user_id": ObjectId(user_id)}
            ).sort("timestamp", -1).limit(limit * 2)  # Get more than needed for filtering
            
            sessions = list(cursor)
            
            # Format for Atlas API: [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
            history = []
            estimated_tokens = 0
            
            # Process sessions in reverse order (oldest first for conversation flow)
            for session in reversed(sessions):
                user_message = session.get("message", "").strip()
                assistant_response = session.get("response", "").strip()
                
                if not user_message or not assistant_response:
                    continue
                
                # Estimate tokens (rough approximation: 4 chars = 1 token)
                message_tokens = len(user_message) // 4 + len(assistant_response) // 4
                
                # Check if adding this exchange would exceed token limit
                if estimated_tokens + message_tokens > max_tokens:
                    break
                
                # Add user message first, then assistant response
                history.extend([
                    {"role": "user", "content": user_message},
                    {"role": "assistant", "content": assistant_response}
                ])
                
                estimated_tokens += message_tokens
                
                # Check if we've reached the exchange limit
                if len(history) >= limit * 2:  # limit exchanges = limit * 2 messages
                    break
            
            # Return most recent exchanges (keep chronological order)
            return history[-limit * 2:] if len(history) > limit * 2 else history
            
        except Exception as e:
            logger.error(f"Error getting conversation history for Atlas API, user_id {user_id}: {e}")
            return []
    
    @staticmethod
    def clear_user_history(user_id: str) -> bool:
        """Clear all chat history for a user"""
        try:
            chat_sessions_collection = get_chat_sessions_collection()
            
            # Delete all chat sessions for this user
            result = chat_sessions_collection.delete_many(
                {"user_id": ObjectId(user_id)}
            )
            
            logger.info(f"Cleared {result.deleted_count} chat sessions for user_id: {user_id}")
            return True
            
        except Exception as e:
            logger.error(f"Error clearing chat history for user_id {user_id}: {e}")
            return False