"""
Token Tracking Decorator
Simplifies adding token management to AI endpoints
"""

from functools import wraps
from flask import request, jsonify, g
import logging
from services.token_manager import TokenManager
from services.subscription_manager import SubscriptionManager

logger = logging.getLogger(__name__)


def require_tokens(
    estimate_tokens_from='message',  # Which request field to use for estimation
    model_field='model',  # Which request field contains model name
    endpoint_name=None  # Endpoint name for logging (auto-detected if None)
):
    """
    Decorator to add token checking and deduction to AI endpoints
    
    Usage:
        @require_tokens(estimate_tokens_from='prompt', model_field='ai_model')
        def my_ai_endpoint():
            # ... your AI logic ...
            # Return response with 'tokens_used' and 'model_used' in metadata
            return result
            
    The decorator will:
    1. Check user has sufficient tokens before processing
    2. Verify user has subscription access to the model
    3. Return 402 if insufficient tokens
    4. Return 403 if subscription doesn't allow model access
    5. Log token usage for analytics
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
 # Get endpoint name
            nonlocal endpoint_name
            if endpoint_name is None:
                endpoint_name = f"{request.endpoint}"
            
            # Get request data
            if request.is_json:
                data = request.get_json() or {}
            else:
                data = request.form.to_dict()
            
            # Get user_id (should be set by require_api_key middleware)
            user_id = data.get('user_id') or getattr(g, 'current_user_id', None)
            
            if not user_id:
                return jsonify({
                    "status": "error",
                    "message": "User ID is required for token tracking"
                }), 400
            
            # Estimate tokens needed
            content = data.get(estimate_tokens_from, '')
            if isinstance(content, str):
                estimated_tokens = max(100, len(content) / 4)
            elif isinstance(content, list):  # For batch operations
                estimated_tokens = max(100, sum(len(str(item)) / 4 for item in content))
            else:
                estimated_tokens = 500  # Default estimate
            
            # Check token balance
            balance_check = TokenManager.check_balance(user_id, estimated_tokens)
            if not balance_check['has_sufficient']:
                logger.warning(f"User {user_id} insufficient tokens for {endpoint_name}: {balance_check['balance']} < {estimated_tokens}")
                return jsonify({
                    "status": "error",
                    "error_code": "INSUFFICIENT_TOKENS",
                    "message": "Insufficient AI tokens",
                    "current_balance": balance_check['balance'],
                    "required_tokens": estimated_tokens,
                    "shortfall": balance_check.get('shortfall', 0),
                    "upgrade_url": "/pricing"
                }), 402
            
            # Get model to be used
            model_name = data.get(model_field, 'geminisearch')
            
            # Check subscription access
            access_check = SubscriptionManager.check_subscription_access(
                user_id=user_id,
                required_benefit={'type': 'ai_tokens', 'model': model_name}
            )
            
            if not access_check['has_access']:
                logger.warning(f"User {user_id} no access to {model_name} at {endpoint_name}")
                return jsonify({
                    "status": "error",
                    "error_code": "SUBSCRIPTION_REQUIRED",
                    "message": access_check.get('reason', 'Model not available in your plan'),
                    "required_model": model_name,
                    "suggested_plan": access_check.get('suggested_plan', 'ai-starter'),
                    "upgrade_url": f"/pricing?highlight={access_check.get('suggested_plan', 'ai-starter')}"
                }), 403
            
            logger.info(f"User {user_id} authorized for {model_name} at {endpoint_name}, balance: {balance_check['balance']}")
            
            # Call the original function
            result = func(*args, **kwargs)
            
            # Extract tokens_used from result (if available)
            # This assumes the endpoint returns (response_dict, status_code) or just response_dict
            if isinstance(result, tuple):
                response_data, status_code = result[0], result[1] if len(result) > 1 else 200
            else:
                response_data, status_code = result, 200
            
            # Only deduct tokens if request was successful
            if status_code < 400:
                try:
                    # Try to extract token usage from response
                    if hasattr(response_data, 'get_json'):
                        response_json = response_data.get_json()
                    elif isinstance(response_data, dict):
                        response_json = response_data
                    else:
                        response_json = {}
                    
                    # Get tokens used from response metadata
                    tokens_used = response_json.get('tokens_used', estimated_tokens)
                    model_used = response_json.get('model_used', model_name)
                    
                    # Deduct tokens
                    deduction_result = TokenManager.deduct_tokens(
                        user_id=user_id,
                        amount=tokens_used,
                        endpoint=endpoint_name,
                        ai_model=model_used,
                        metadata={'request_data': {estimate_tokens_from: content[:100]}}  # First 100 chars
                    )
                    
                    if deduction_result['success']:
                        logger.info(f"Deducted {tokens_used:.2f} tokens from user {user_id} for {endpoint_name}")
                    else:
                        logger.error(f"Token deduction failed for user {user_id}: {deduction_result.get('message')}")
                        
                except Exception as e:
                    logger.error(f"Error during token deduction for {endpoint_name}: {e}")
                    # Don't fail the request if token deduction fails
            
            return result
        
        return wrapper
    return decorator


# Simplified version for endpoints that handle token deduction manually
def check_tokens_only(estimate_tokens=500, model_field='model'):
    """
    Lighter decorator that only checks tokens/subscription, doesn't deduct
    Use this when endpoint needs custom token deduction logic
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            data = request.get_json() or request.form.to_dict()
            user_id = data.get('user_id') or getattr(g, 'current_user_id', None)
            
            if not user_id:
                return jsonify({"status": "error", "message": "User ID required"}), 400
            
            # Check balance
            balance_check = TokenManager.check_balance(user_id,estimate_tokens)
            if not balance_check['has_sufficient']:
                return jsonify({
                    "status": "error",
                    "error_code": "INSUFFICIENT_TOKENS",
                    "message": "Insufficient AI tokens",
                    "current_balance": balance_check['balance'],
                    "upgrade_url": "/pricing"
                }), 402
            
            # Check subscription
            model_name = data.get(model_field, 'geminisearch')
            access_check = SubscriptionManager.check_subscription_access(
                user_id=user_id,
                required_benefit={'type': 'ai_tokens', 'model': model_name}
            )
            
            if not access_check['has_access']:
                return jsonify({
                    "status": "error",
                    "error_code": "SUBSCRIPTION_REQUIRED",
                    "message": access_check.get('reason'),
                    "upgrade_url": "/pricing"
                }), 403
            
            return func(*args, **kwargs)
        
        return wrapper
    return decorator
