# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Token Refresh Service for PowerOn Gateway

This service handles automatic token refresh for OAuth connections
when they are accessed via API calls. It runs silently in the background
to ensure users don't experience token expiration issues.
"""

import logging
from typing import Dict, Any
from modules.datamodels.datamodelUam import UserConnection, AuthAuthority
from modules.shared.timeUtils import getUtcTimestamp
from modules.shared.auditLogger import audit_logger

logger = logging.getLogger(__name__)

class TokenRefreshService:
    """Service for automatic token refresh operations"""
    
    def __init__(self):
        self.rate_limit_map = {}  # Track refresh attempts per connection
        self.max_attempts_per_hour = 3
        self.refresh_window_minutes = 60
    
    def _is_rate_limited(self, connection_id: str) -> bool:
        """Check if connection is rate limited for refresh attempts"""
        now = getUtcTimestamp()
        if connection_id not in self.rate_limit_map:
            return False
        
        # Remove attempts older than 1 hour
        recent_attempts = [
            attempt_time for attempt_time in self.rate_limit_map[connection_id]
            if now - attempt_time < (self.refresh_window_minutes * 60)
        ]
        self.rate_limit_map[connection_id] = recent_attempts
        
        return len(recent_attempts) >= self.max_attempts_per_hour
    
    def _record_refresh_attempt(self, connection_id: str) -> None:
        """Record a refresh attempt for rate limiting"""
        now = getUtcTimestamp()
        if connection_id not in self.rate_limit_map:
            self.rate_limit_map[connection_id] = []
        self.rate_limit_map[connection_id].append(now)
    
    async def _refresh_google_token(self, interface, connection: UserConnection) -> bool:
        """Refresh Google OAuth token"""
        try:
            logger.debug(f"Refreshing Google token for connection {connection.id}")
            
            # Get current token (no refresh in interface layer)
            current_token = interface.getConnectionToken(connection.id)
            if not current_token:
                logger.warning(f"No Google token found for connection {connection.id}")
                return False
            
            # Import Google token refresh logic
            from modules.auth.tokenManager import TokenManager
            token_manager = TokenManager()
            
            # Attempt to refresh the token
            refreshedToken = token_manager.refreshToken(current_token)
            if refreshedToken:
                # Save the refreshed token
                interface.saveConnectionToken(refreshedToken)
                
                # Update connection status
                interface.db.recordModify(UserConnection, connection.id, {
                    "lastChecked": getUtcTimestamp(),
                    "expiresAt": refreshedToken.expiresAt
                })
                
                logger.info(f"Successfully refreshed Google token for connection {connection.id}")
                
                # Log audit event
                try:
                    audit_logger.logSecurityEvent(
                        userId=str(connection.userId),
                        mandateId="system",
                        action="token_refresh",
                        details=f"Google token refreshed for connection {connection.id}"
                    )
                except Exception:
                    pass
                
                return True
            else:
                logger.warning(f"Failed to refresh Google token for connection {connection.id}")
                return False
                
        except Exception as e:
            logger.error(f"Error refreshing Google token for connection {connection.id}: {str(e)}")
            return False
    
    async def _refresh_microsoft_token(self, interface, connection: UserConnection) -> bool:
        """Refresh Microsoft OAuth token"""
        try:
            logger.debug(f"Refreshing Microsoft token for connection {connection.id}")
            
            # Get current token (no refresh in interface layer)
            current_token = interface.getConnectionToken(connection.id)
            if not current_token:
                logger.warning(f"No Microsoft token found for connection {connection.id}")
                return False
            
            # Import Microsoft token refresh logic
            from modules.auth.tokenManager import TokenManager
            token_manager = TokenManager()
            
            # Attempt to refresh the token
            refreshedToken = token_manager.refreshToken(current_token)
            if refreshedToken:
                # Save the refreshed token
                interface.saveConnectionToken(refreshedToken)
                
                # Update connection status
                interface.db.recordModify(UserConnection, connection.id, {
                    "lastChecked": getUtcTimestamp(),
                    "expiresAt": refreshedToken.expiresAt
                })
                
                logger.info(f"Successfully refreshed Microsoft token for connection {connection.id}")
                
                # Log audit event
                try:
                    audit_logger.logSecurityEvent(
                        userId=str(connection.userId),
                        mandateId="system",
                        action="token_refresh",
                        details=f"Microsoft token refreshed for connection {connection.id}"
                    )
                except Exception:
                    pass
                
                return True
            else:
                logger.warning(f"Failed to refresh Microsoft token for connection {connection.id}")
                return False
                
        except Exception as e:
            logger.error(f"Error refreshing Microsoft token for connection {connection.id}: {str(e)}")
            return False
    
    async def refresh_expired_tokens(self, user_id: str) -> Dict[str, Any]:
        """
        Refresh expired OAuth tokens for a user
        
        Args:
            user_id: User ID to refresh tokens for
            
        Returns:
            Dict with refresh results
        """
        try:
            logger.debug(f"Starting silent token refresh for user {user_id}")
            
            # Get user interface
            from modules.security.rootAccess import getRootUser
            from modules.interfaces.interfaceDbApp import getInterface
            rootUser = getRootUser()
            root_interface = getInterface(rootUser)
            
            # Get user connections
            connections = root_interface.getUserConnections(user_id)
            if not connections:
                logger.debug(f"No connections found for user {user_id}")
                return {"refreshed": 0, "failed": 0, "rate_limited": 0}
            
            refreshed_count = 0
            failed_count = 0
            rate_limited_count = 0
            
            # Process each connection
            for connection in connections:
                # Only refresh expired OAuth connections
                if (connection.tokenStatus == 'expired' and 
                    connection.authority in [AuthAuthority.GOOGLE, AuthAuthority.MSFT]):
                    
                    # Check rate limiting
                    if self._is_rate_limited(connection.id):
                        logger.warning(f"Rate limited for connection {connection.id}")
                        rate_limited_count += 1
                        continue
                    
                    # Record attempt
                    self._record_refresh_attempt(connection.id)
                    
                    # Refresh based on authority
                    success = False
                    if connection.authority == AuthAuthority.GOOGLE:
                        success = await self._refresh_google_token(root_interface, connection)
                    elif connection.authority == AuthAuthority.MSFT:
                        success = await self._refresh_microsoft_token(root_interface, connection)
                    
                    if success:
                        refreshed_count += 1
                    else:
                        failed_count += 1
            
            result = {
                "refreshed": refreshed_count,
                "failed": failed_count,
                "rate_limited": rate_limited_count
            }
            
            logger.info(f"Silent token refresh completed for user {user_id}: {result}")
            return result
            
        except Exception as e:
            logger.error(f"Error during silent token refresh for user {user_id}: {str(e)}")
            return {"refreshed": 0, "failed": 0, "rate_limited": 0, "error": str(e)}
    
    async def proactive_refresh(self, user_id: str) -> Dict[str, Any]:
        """
        Proactively refresh tokens that expire within 5 minutes
        
        Args:
            user_id: User ID to check tokens for
            
        Returns:
            Dict with refresh results
        """
        try:
            logger.debug(f"Starting proactive token refresh for user {user_id}")
            
            # Get user interface
            from modules.security.rootAccess import getRootUser
            from modules.interfaces.interfaceDbApp import getInterface
            rootUser = getRootUser()
            root_interface = getInterface(rootUser)
            
            # Get user connections
            connections = root_interface.getUserConnections(user_id)
            if not connections:
                return {"refreshed": 0, "failed": 0, "rate_limited": 0}
            
            refreshed_count = 0
            failed_count = 0
            rate_limited_count = 0
            current_time = getUtcTimestamp()
            five_minutes = 5 * 60  # 5 minutes in seconds
            
            # Process each connection
            for connection in connections:
                # Only refresh active tokens that expire soon
                if (connection.tokenStatus == 'active' and 
                    connection.tokenExpiresAt and
                    connection.authority in [AuthAuthority.GOOGLE, AuthAuthority.MSFT]):
                    
                    # Check if token expires within 5 minutes
                    time_until_expiry = connection.tokenExpiresAt - current_time
                    if 0 < time_until_expiry <= five_minutes:
                        
                        # Check rate limiting
                        if self._is_rate_limited(connection.id):
                            logger.warning(f"Rate limited for proactive refresh of connection {connection.id}")
                            rate_limited_count += 1
                            continue
                        
                        # Record attempt
                        self._record_refresh_attempt(connection.id)
                        
                        # Refresh based on authority
                        success = False
                        if connection.authority == AuthAuthority.GOOGLE:
                            success = await self._refresh_google_token(root_interface, connection)
                        elif connection.authority == AuthAuthority.MSFT:
                            success = await self._refresh_microsoft_token(root_interface, connection)
                        
                        if success:
                            refreshed_count += 1
                            logger.info(f"Proactively refreshed {connection.authority} token for connection {connection.id}")
                        else:
                            failed_count += 1
            
            result = {
                "refreshed": refreshed_count,
                "failed": failed_count,
                "rate_limited": rate_limited_count
            }
            
            if refreshed_count > 0:
                logger.info(f"Proactive token refresh completed for user {user_id}: {result}")
            
            return result
            
        except Exception as e:
            logger.error(f"Error during proactive token refresh for user {user_id}: {str(e)}")
            return {"refreshed": 0, "failed": 0, "rate_limited": 0, "error": str(e)}

# Global service instance
token_refresh_service = TokenRefreshService()

