gateway/modules/serviceCenter/services/serviceBilling/mainServiceBilling.py

# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Billing Service - Central service for billing operations.

Handles:
- Balance checks before AI operations
- Cost recording after AI operations
- Provider permission checks via RBAC
- Price calculation with markup
"""

import logging
from typing import Dict, Any, List, Optional
from datetime import datetime

from modules.datamodels.datamodelUam import User
from modules.datamodels.datamodelBilling import (
    BillingModelEnum,
    BillingCheckResult,
    TransactionTypeEnum,
    ReferenceTypeEnum,
    BillingTransaction,
    BillingBalanceResponse,
)
from modules.interfaces.interfaceDbBilling import getInterface as getBillingInterface

logger = logging.getLogger(__name__)

# Markup percentage for internal pricing (+50% für Infrastruktur und Platform Service + 50% für Währungsrisiko ==> Faktor 2.0)
BILLING_MARKUP_PERCENT = 100

# Singleton cache
_billingServices: Dict[str, "BillingService"] = {}


def getService(currentUser: User, mandateId: str, featureInstanceId: str = None, featureCode: str = None) -> "BillingService":
    """
    Factory function to get or create a BillingService instance.
    
    Args:
        currentUser: Current user object
        mandateId: Mandate ID for context
        featureInstanceId: Optional feature instance ID
        featureCode: Optional feature code (e.g., 'chatplayground', 'automation')
        
    Returns:
        BillingService instance
    """
    cacheKey = f"{currentUser.id}_{mandateId}_{featureInstanceId}"
    
    if cacheKey not in _billingServices:
        _billingServices[cacheKey] = BillingService(currentUser, mandateId, featureInstanceId, featureCode)
    else:
        _billingServices[cacheKey].setContext(currentUser, mandateId, featureInstanceId, featureCode)
    
    return _billingServices[cacheKey]


def _get_feature_code_from_context(context) -> Optional[str]:
    """Extract featureCode from ServiceCenterContext."""
    if context.workflow and hasattr(context.workflow, "feature") and context.workflow.feature:
        return getattr(context.workflow.feature, "code", None)
    return getattr(context.workflow, "featureCode", None) if context.workflow else None


class BillingService:
    """
    Central billing service for AI operations.
    
    Responsibilities:
    - Check balance before operations
    - Record usage costs
    - Apply pricing markup
    - Check provider permissions via RBAC
    
    Supports both service center (context, get_service) and legacy (user, mandateId, ...) initialization.
    """

    def __init__(self, context_or_user, mandateId=None, featureInstanceId=None, featureCode=None, get_service=None):
        """
        Initialize the billing service.
        
        Service center: (context, get_service) - resolver passes exactly these two args
        Legacy: (currentUser, mandateId, featureInstanceId, featureCode) from getService() factory
        """
        # Detect service center: second arg is callable (get_service)
        if mandateId is not None and callable(mandateId):
            ctx = context_or_user
            get_service = mandateId
            self.currentUser = ctx.user
            self.mandateId = ctx.mandate_id or ""
            self.featureInstanceId = ctx.feature_instance_id
            self.featureCode = _get_feature_code_from_context(ctx)
        elif get_service is not None and hasattr(context_or_user, "user"):
            ctx = context_or_user
            self.currentUser = ctx.user
            self.mandateId = ctx.mandate_id or ""
            self.featureInstanceId = ctx.feature_instance_id
            self.featureCode = _get_feature_code_from_context(ctx)
        else:
            self.currentUser = context_or_user
            self.mandateId = mandateId or ""
            self.featureInstanceId = featureInstanceId
            self.featureCode = featureCode

        self._billingInterface = getBillingInterface(self.currentUser, self.mandateId)
        self._settingsCache = None
    
    def setContext(
        self, 
        currentUser: User, 
        mandateId: str, 
        featureInstanceId: str = None,
        featureCode: str = None
    ):
        """Update service context."""
        self.currentUser = currentUser
        self.mandateId = mandateId
        self.featureInstanceId = featureInstanceId
        self.featureCode = featureCode
        self._billingInterface = getBillingInterface(currentUser, mandateId)
        self._settingsCache = None
    
    def _getSettings(self) -> Optional[Dict[str, Any]]:
        """Get billing settings with caching."""
        if self._settingsCache is None:
            self._settingsCache = self._billingInterface.getSettings(self.mandateId)
        return self._settingsCache
    
    # =========================================================================
    # Price Calculation
    # =========================================================================
    
    def calculatePriceWithMarkup(self, basePriceCHF: float) -> float:
        """
        Calculate final price with markup.
        
        The AICore plugins return prices in their original currency (USD).
        This method applies the configured markup percentage.
        
        Args:
            basePriceCHF: Base price from AI model (actually USD from provider)
            
        Returns:
            Final price in CHF with markup applied
        """
        if basePriceCHF <= 0:
            return 0.0
        
        # Apply markup (50% = multiply by 1.5)
        markup_multiplier = 1 + (BILLING_MARKUP_PERCENT / 100)
        return round(basePriceCHF * markup_multiplier, 6)
    
    # =========================================================================
    # Balance Operations
    # =========================================================================
    
    def checkBalance(self, estimatedCost: float = 0.0) -> BillingCheckResult:
        """
        Check if the current user/mandate has sufficient balance.
        
        Args:
            estimatedCost: Estimated cost of the operation (with markup applied)
            
        Returns:
            BillingCheckResult indicating if operation is allowed
        """
        return self._billingInterface.checkBalance(
            self.mandateId,
            self.currentUser.id,
            estimatedCost
        )
    
    def hasBalance(self, estimatedCost: float = 0.0) -> bool:
        """
        Quick check if balance is sufficient.
        
        Args:
            estimatedCost: Estimated cost with markup
            
        Returns:
            True if operation is allowed
        """
        result = self.checkBalance(estimatedCost)
        return result.allowed
    
    def getCurrentBalance(self) -> float:
        """
        Get current balance for the user/mandate.
        
        Returns:
            Current balance in CHF
        """
        result = self.checkBalance(0.0)
        return result.currentBalance or 0.0
    
    # =========================================================================
    # Usage Recording
    # =========================================================================
    
    def recordUsage(
        self,
        priceCHF: float,
        workflowId: str = None,
        aicoreProvider: str = None,
        aicoreModel: str = None,
        description: str = None,
        processingTime: float = None,
        bytesSent: int = None,
        bytesReceived: int = None,
        errorCount: int = None
    ) -> Optional[Dict[str, Any]]:
        """Record AI usage cost as a billing transaction with markup applied."""
        if priceCHF <= 0:
            return None
        
        finalPrice = self.calculatePriceWithMarkup(priceCHF)
        if finalPrice <= 0:
            return None
        
        if not description:
            description = f"AI Usage: {aicoreModel or aicoreProvider or 'unknown'}"
        
        return self._billingInterface.recordUsage(
            mandateId=self.mandateId,
            userId=self.currentUser.id,
            priceCHF=finalPrice,
            workflowId=workflowId,
            featureInstanceId=self.featureInstanceId,
            featureCode=self.featureCode,
            aicoreProvider=aicoreProvider,
            aicoreModel=aicoreModel,
            description=description,
            processingTime=processingTime,
            bytesSent=bytesSent,
            bytesReceived=bytesReceived,
            errorCount=errorCount
        )
    
    def getWorkflowCost(self, workflowId: str) -> float:
        """Get total cost for a workflow from billing transactions."""
        return self._billingInterface.getWorkflowCost(workflowId)
    
    # =========================================================================
    # Provider Permission Check (via RBAC)
    # =========================================================================
    
    def isProviderAllowed(self, provider: str) -> bool:
        """
        Check if the user has permission to use an AICore provider.
        
        Uses RBAC to check for resource permission:
        resource.aicore.{provider}
        
        Args:
            provider: Provider name (e.g., 'anthropic', 'openai')
            
        Returns:
            True if provider is allowed
        """
        try:
            from modules.security.rbac import RbacClass
            from modules.datamodels.datamodelRbac import AccessRuleContext
            from modules.security.rootAccess import getRootDbAppConnector
            
            # Get database connector via established pattern
            dbApp = getRootDbAppConnector()
            
            rbac = RbacClass(dbApp, dbApp)
            resourceKey = f"resource.aicore.{provider}"
            
            # Check if user has view permission for this resource (view = use for RESOURCE context)
            permissions = rbac.getUserPermissions(
                self.currentUser,
                AccessRuleContext.RESOURCE,
                resourceKey,
                mandateId=self.mandateId
            )
            
            return permissions.view
        except Exception as e:
            logger.warning(f"Error checking provider permission: {e}")
            # Default to allowed if RBAC check fails
            return True
    
    def getallowedProviders(self) -> List[str]:
        """
        Get list of AICore providers the user is allowed to use.
        
        Returns:
            List of allowed provider names
        """
        try:
            from modules.aicore.aicoreModelRegistry import modelRegistry
            
            # Get all available providers
            connectors = modelRegistry.discoverConnectors()
            allProviders = [c.getConnectorType() for c in connectors]
            
            # Filter by RBAC permissions
            return [p for p in allProviders if self.isProviderAllowed(p)]
        except Exception as e:
            logger.warning(f"Error getting allowed providers: {e}")
            return []
    
    # =========================================================================
    # Admin Operations
    # =========================================================================
    
    def addCredit(
        self,
        amount: float,
        description: str = "Manual credit",
        referenceType: ReferenceTypeEnum = ReferenceTypeEnum.ADMIN
    ) -> Optional[Dict[str, Any]]:
        """
        Add credit to the account (admin operation).
        
        Args:
            amount: Amount to credit (positive)
            description: Transaction description
            referenceType: Reference type (ADMIN, PAYMENT, SYSTEM)
            
        Returns:
            Created transaction dict or None
        """
        if amount <= 0:
            return None
        
        settings = self._getSettings()
        if not settings:
            logger.warning(f"No billing settings for mandate {self.mandateId}")
            return None
        
        billingModel = BillingModelEnum(settings.get("billingModel", BillingModelEnum.UNLIMITED.value))
        
        # Get or create account
        if billingModel == BillingModelEnum.PREPAY_USER:
            account = self._billingInterface.getOrCreateUserAccount(
                self.mandateId, 
                self.currentUser.id,
                initialBalance=0.0
            )
        else:
            account = self._billingInterface.getOrCreateMandateAccount(
                self.mandateId,
                initialBalance=0.0
            )
        
        # Create credit transaction
        transaction = BillingTransaction(
            accountId=account["id"],
            transactionType=TransactionTypeEnum.CREDIT,
            amount=amount,
            description=description,
            referenceType=referenceType
        )
        
        return self._billingInterface.createTransaction(transaction)
    
    # =========================================================================
    # Statistics & Reporting
    # =========================================================================
    
    def getBalancesForUser(self) -> List[BillingBalanceResponse]:
        """
        Get all billing balances for the current user.
        
        Returns:
            List of balance responses for each mandate
        """
        return self._billingInterface.getBalancesForUser(self.currentUser.id)
    
    def getTransactionHistory(self, limit: int = 100) -> List[Dict[str, Any]]:
        """
        Get transaction history for the user across all mandates.
        
        Args:
            limit: Maximum number of transactions
            
        Returns:
            List of transactions
        """
        return self._billingInterface.getTransactionsForUser(self.currentUser.id, limit=limit)


# ============================================================================
# Exception Classes
# ============================================================================

class InsufficientBalanceException(Exception):
    """Raised when there's insufficient balance for an operation."""
    
    def __init__(self, currentBalance: float, requiredAmount: float, message: str = None):
        self.currentBalance = currentBalance
        self.requiredAmount = requiredAmount
        self.message = message or f"Insufficient balance. Current: {currentBalance:.2f} CHF, Required: {requiredAmount:.2f} CHF"
        super().__init__(self.message)


class ProviderNotAllowedException(Exception):
    """Raised when a user doesn't have permission to use an AI provider."""
    
    def __init__(self, provider: str, message: str = None):
        self.provider = provider
        self.message = message or f"Provider '{provider}' is not allowed for your role"
        super().__init__(self.message)


class BillingContextError(Exception):
    """Raised when billing context is incomplete (missing mandateId, user, etc.).
    
    This is a FAIL-SAFE error: AI calls MUST NOT proceed without valid billing context.
    Acts like a 0 CHF credit card pre-authorization check - validates that billing
    CAN be recorded before any expensive AI operation starts.
    """
    
    def __init__(self, message: str = None):
        self.message = message or "Billing context incomplete - AI call blocked"
        super().__init__(self.message)


# Expose exception classes on BillingService so consumers can use service.InsufficientBalanceException
# instead of importing from this module
BillingService.InsufficientBalanceException = InsufficientBalanceException
BillingService.ProviderNotAllowedException = ProviderNotAllowedException
BillingService.BillingContextError = BillingContextError