# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Token Manager Service
Handles all token operations including automatic refresh for backend services.
"""

import logging
import httpx
from typing import Optional, Dict, Any, Callable

from modules.datamodels.datamodelSecurity import Token
from modules.datamodels.datamodelUam import AuthAuthority
from modules.shared.configuration import APP_CONFIG
from modules.shared.timeUtils import getUtcTimestamp, createExpirationTimestamp, parseTimestamp

logger = logging.getLogger(__name__)

class TokenManager:
    """Centralized token management service"""
    
    def __init__(self):
        # Microsoft OAuth configuration
        self.msft_client_id = APP_CONFIG.get("Service_MSFT_CLIENT_ID")
        self.msft_client_secret = APP_CONFIG.get("Service_MSFT_CLIENT_SECRET")
        self.msft_tenant_id = APP_CONFIG.get("Service_MSFT_TENANT_ID", "common")
        
        # Google OAuth configuration
        self.google_client_id = APP_CONFIG.get("Service_GOOGLE_CLIENT_ID")
        self.google_client_secret = APP_CONFIG.get("Service_GOOGLE_CLIENT_SECRET")
    
    def refreshMicrosoftToken(self, refreshToken: str, userId: str, oldToken: Token) -> Optional[Token]:
        """Refresh Microsoft OAuth token using refresh token"""
        try:
            logger.debug(f"refreshMicrosoftToken: Starting Microsoft token refresh for user {userId}")
            logger.debug(f"refreshMicrosoftToken: Configuration check - client_id: {bool(self.msft_client_id)}, client_secret: {bool(self.msft_client_secret)}")
            
            if not self.msft_client_id or not self.msft_client_secret:
                logger.error("Microsoft OAuth configuration not found")
                return None
            
            # Microsoft token refresh endpoint
            tokenUrl = f"https://login.microsoftonline.com/{self.msft_tenant_id}/oauth2/v2.0/token"
            logger.debug(f"refreshMicrosoftToken: Using token URL: {tokenUrl}")
            
            # Prepare refresh request
            data = {
                "client_id": self.msft_client_id,
                "client_secret": self.msft_client_secret,
                "grant_type": "refresh_token",
                "refresh_token": refreshToken,
                "scope": "Mail.ReadWrite Mail.Send Mail.ReadWrite.Shared User.Read"
            }
            logger.debug(f"refreshMicrosoftToken: Refresh request data prepared (refreshToken length: {len(refreshToken) if refreshToken else 0})")
            
            # Make refresh request
            with httpx.Client(timeout=30.0) as client:
                logger.debug(f"refreshMicrosoftToken: Making HTTP request to Microsoft OAuth endpoint")
                response = client.post(tokenUrl, data=data)
                logger.debug(f"refreshMicrosoftToken: HTTP response status: {response.status_code}")
                
                if response.status_code == 200:
                    tokenData = response.json()
                    logger.debug(f"refreshMicrosoftToken: Token refresh successful, creating new token")
                    
                    # Create new token
                    newToken = Token(
                        userId=userId,
                        authority=AuthAuthority.MSFT,
                        connectionId=oldToken.connectionId,  # Preserve connection ID
                        tokenAccess=tokenData["access_token"],
                        tokenRefresh=tokenData.get("refresh_token", refreshToken),  # Keep old refresh token if new one not provided
                        tokenType=tokenData.get("token_type", "bearer"),
                        expiresAt=createExpirationTimestamp(tokenData.get("expires_in", 3600)),
                        createdAt=getUtcTimestamp()
                    )
                    
                    logger.debug(f"refreshMicrosoftToken: New token created with ID: {newToken.id}")
                    return newToken
                else:
                    logger.error(f"Failed to refresh Microsoft token: {response.status_code} - {response.text}")
                    return None
                    
        except Exception as e:
            logger.error(f"Error refreshing Microsoft token: {str(e)}")
            return None
    
    def refreshGoogleToken(self, refreshToken: str, userId: str, oldToken: Token) -> Optional[Token]:
        """Refresh Google OAuth token using refresh token"""
        try:
            logger.debug(f"refreshGoogleToken: Starting Google token refresh for user {userId}")
            logger.debug(f"refreshGoogleToken: Configuration check - client_id: {bool(self.google_client_id)}, client_secret: {bool(self.google_client_secret)}")
            
            if not self.google_client_id or not self.google_client_secret:
                logger.error("Google OAuth configuration not found")
                return None
            
            # Google token refresh endpoint
            tokenUrl = "https://oauth2.googleapis.com/token"
            logger.debug(f"refreshGoogleToken: Using token URL: {tokenUrl}")
            
            # Prepare refresh request
            data = {
                "client_id": self.google_client_id,
                "client_secret": self.google_client_secret,
                "grant_type": "refresh_token",
                "refresh_token": refreshToken
            }
            logger.debug(f"refreshGoogleToken: Refresh request data prepared (refreshToken length: {len(refreshToken) if refreshToken else 0})")
            
            # Make refresh request
            with httpx.Client(timeout=30.0) as client:
                logger.debug(f"refreshGoogleToken: Making HTTP request to Google OAuth endpoint")
                response = client.post(tokenUrl, data=data)
                logger.debug(f"refreshGoogleToken: HTTP response status: {response.status_code}")
                
                if response.status_code == 200:
                    tokenData = response.json()
                    logger.debug(f"refreshGoogleToken: Token refresh successful, creating new token")
                    
                    # Validate the response contains required fields
                    if "access_token" not in tokenData:
                        logger.error("Google token refresh response missing access_token")
                        return None
                    
                    # Create new token
                    newToken = Token(
                        userId=userId,
                        authority=AuthAuthority.GOOGLE,
                        connectionId=oldToken.connectionId,  # Preserve connection ID
                        tokenAccess=tokenData["access_token"],
                        tokenRefresh=tokenData.get("refresh_token", refreshToken),  # Use new refresh token if provided
                        tokenType=tokenData.get("token_type", "bearer"),
                        expiresAt=createExpirationTimestamp(tokenData.get("expires_in", 3600)),
                        createdAt=getUtcTimestamp()
                    )
                    
                    logger.debug(f"refreshGoogleToken: New token created with ID: {newToken.id}")
                    return newToken
                else:
                    errorDetails = response.text
                    logger.error(f"Failed to refresh Google token: {response.status_code} - {errorDetails}")
                    
                    # Handle specific error cases
                    if response.status_code == 400:
                        try:
                            errorData = response.json()
                            errorCode = errorData.get("error")
                            if errorCode == "invalid_grant":
                                logger.warning("Google refresh token is invalid or expired - user needs to re-authenticate")
                            elif errorCode == "invalid_client":
                                logger.error("Google OAuth client configuration is invalid")
                        except:
                            pass
                    
                    return None
                    
        except Exception as e:
            logger.error(f"Error refreshing Google token: {str(e)}")
            return None
    
    def refreshToken(self, oldToken: Token) -> Optional[Token]:
        """Refresh an expired token using the appropriate OAuth service"""
        try:
            logger.debug(f"refreshToken: Starting refresh for token {oldToken.id}, authority: {oldToken.authority}")
            logger.debug(f"refreshToken: Token details: userId={oldToken.userId}, connectionId={oldToken.connectionId}, hasRefreshToken={bool(oldToken.tokenRefresh)}")
            
            # Cooldown: avoid refreshing too frequently if a workflow triggers refresh repeatedly
            # Only allow a new refresh if at least 10 minutes passed since the token was created/refreshed
            try:
                nowTs = getUtcTimestamp()
                createdTs = parseTimestamp(oldToken.createdAt, default=0.0)
                secondsSinceLastRefresh = nowTs - createdTs
                if secondsSinceLastRefresh < 10 * 60:
                    logger.info(
                        f"refreshToken: Skipping refresh for connection {oldToken.connectionId} due to cooldown. "
                        f"Last refresh {int(secondsSinceLastRefresh)}s ago (< 600s)."
                    )
                    # Return the existing token to avoid caller errors while preventing provider rate limits
                    return oldToken
            except Exception:
                # If any issue reading timestamps, proceed with normal refresh to be safe
                pass
            
            if not oldToken.tokenRefresh:
                logger.warning(f"No refresh token available for {oldToken.authority}")
                return None
            
            # Route to appropriate refresh method
            if oldToken.authority == AuthAuthority.MSFT:
                logger.debug(f"refreshToken: Refreshing Microsoft token")
                return self.refreshMicrosoftToken(oldToken.tokenRefresh, oldToken.userId, oldToken)
            elif oldToken.authority == AuthAuthority.GOOGLE:
                logger.debug(f"refreshToken: Refreshing Google token")
                return self.refreshGoogleToken(oldToken.tokenRefresh, oldToken.userId, oldToken)
            else:
                logger.warning(f"Unknown authority for token refresh: {oldToken.authority}")
                return None
                
        except Exception as e:
            logger.error(f"Error refreshing token: {str(e)}")
            return None
    
    def ensureFreshToken(self, token: Token, *, secondsBeforeExpiry: int = 30 * 60, saveCallback: Optional[Callable[[Token], None]] = None) -> Optional[Token]:
        """Ensure a token is fresh; refresh if expiring within threshold.

        Args:
            token: Existing token to validate/refresh.
            secondsBeforeExpiry: Threshold window to proactively refresh.
            saveCallback: Optional function to persist a refreshed token.

        Returns:
            A fresh token (refreshed or original) or None if refresh failed.
        """
        try:
            if token is None:
                return None

            nowTs = getUtcTimestamp()
            expiresAt = token.expiresAt or 0

            # If token expires within the threshold, try to refresh
            if expiresAt and expiresAt < (nowTs + secondsBeforeExpiry):
                logger.info(
                    f"ensureFreshToken: Token for connection {token.connectionId} expiring soon "
                    f"(in {max(0, expiresAt - nowTs)}s). Attempting proactive refresh."
                )
                refreshed = self.refreshToken(token)
                if refreshed:
                    if saveCallback is not None:
                        try:
                            saveCallback(refreshed)
                        except Exception as e:
                            logger.warning(f"ensureFreshToken: Failed to persist refreshed token: {e}")
                    return refreshed
                else:
                    logger.warning("ensureFreshToken: Token refresh failed")
                    return None

            # Token is sufficiently fresh
            return token
        except Exception as e:
            logger.error(f"ensureFreshToken: Error ensuring fresh token: {e}")
            return None

    # Convenience wrapper to fetch and ensure fresh token for a connection via interface layer
    def getFreshToken(self, connectionId: str, secondsBeforeExpiry: int = 30 * 60, interface=None) -> Optional[Token]:
        """Return a fresh token for a connection, refreshing when expiring soon.

        Reads the latest stored token via interface layer, then
        uses ensure_fresh_token to refresh if needed and persists the refreshed
        token via interface layer.
        
        Args:
            connectionId: Connection ID to get token for
            secondsBeforeExpiry: Seconds before expiry to refresh
            interface: Optional interface instance (if None, uses root interface)
        """
        try:
            if interface is None:
                from modules.security.rootAccess import getRootUser
                from modules.interfaces.interfaceDbApp import getInterface
                rootUser = getRootUser()
                interface = getInterface(rootUser)
            
            token = interface.getConnectionToken(connectionId)
            if not token:
                return None
            return self.ensureFreshToken(
                token,
                secondsBeforeExpiry=secondsBeforeExpiry,
                saveCallback=lambda t: interface.saveConnectionToken(t)
            )
        except Exception as e:
            logger.error(f"getFreshToken: Error fetching or refreshing token for connection {connectionId}: {e}")
            return None

