import json
import os
import boto3
import stripe
from datetime import datetime
from decimal import Decimal

# Initialize AWS clients
dynamodb = boto3.resource('dynamodb')
secrets_client = boto3.client('secretsmanager')

# Environment variables
STRIPE_API_KEY_SECRET = os.environ.get('STRIPE_API_KEY_SECRET')
DYNAMODB_TABLE = os.environ.get('DYNAMODB_TABLE')
COST_PER_1M_TOKENS = float(os.environ.get('COST_PER_1M_TOKENS', '2.00'))

def get_stripe_api_key():
    """Retrieve Stripe API key from Secrets Manager"""
    try:
        response = secrets_client.get_secret_value(
            SecretId=STRIPE_API_KEY_SECRET
        )
        return response['SecretString']
    except Exception as e:
        print(f"Error retrieving secret: {e}")
        raise

def calculate_cost(tokens_used):
    """Calculate cost based on tokens used"""
    return (tokens_used / 1_000_000) * COST_PER_1M_TOKENS

def lambda_handler(event, context):
    """
    Lambda handler for billing aggregation and Stripe integration
    """
    # Parse request body
    body = json.loads(event.get('body', '{}'))

    customer_id = body.get('customer_id')
    tokens_used = body.get('tokens_used', 0)
    action = body.get('action', 'record_usage')

    if not customer_id:
        return {
            'statusCode': 400,
            'body': json.dumps({'error': 'customer_id is required'})
        }

    # Initialize Stripe
    stripe.api_key = get_stripe_api_key()

    table = dynamodb.Table(DYNAMODB_TABLE)

    if action == 'record_usage':
        return record_usage(table, customer_id, tokens_used)
    elif action == 'get_invoice':
        return get_invoice(table, customer_id)
    elif action == 'process_payment':
        return process_payment(table, customer_id, tokens_used)
    else:
        return {
            'statusCode': 400,
            'body': json.dumps({'error': 'Invalid action'})
        }

def record_usage(table, customer_id, tokens_used):
    """Record usage to DynamoDB"""
    try:
        timestamp = datetime.utcnow().isoformat()
        cost = calculate_cost(tokens_used)

        item = {
            'customer_id': customer_id,
            'timestamp': timestamp,
            'tokens_used': Decimal(str(tokens_used)),
            'cost': Decimal(str(cost)),
            'billed': False
        }

        table.put_item(Item=item)

        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Usage recorded successfully',
                'customer_id': customer_id,
                'tokens_used': tokens_used,
                'cost': cost
            })
        }
    except Exception as e:
        print(f"Error recording usage: {e}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }

def get_invoice(table, customer_id):
    """Get customer invoice from DynamoDB"""
    try:
        response = table.query(
            KeyConditionExpression='customer_id = :cid',
            ExpressionAttributeValues={
                ':cid': customer_id
            }
        )

        items = response.get('Items', [])

        # Calculate totals
        total_tokens = sum(item.get('tokens_used', 0) for item in items)
        total_cost = sum(item.get('cost', 0) for item in items)

        return {
            'statusCode': 200,
            'body': json.dumps({
                'customer_id': customer_id,
                'total_tokens': float(total_tokens),
                'total_cost': float(total_cost),
                'usage_records': items
            })
        }
    except Exception as e:
        print(f"Error getting invoice: {e}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }

def process_payment(table, customer_id, tokens_used):
    """Process payment via Stripe and update DynamoDB"""
    try:
        cost = calculate_cost(tokens_used)

        # Create Stripe charge
        charge = stripe.Charge.create(
            amount=int(cost * 100),  # Convert to cents
            currency='usd',
            description=f'vLLM CAAE Usage - {tokens_used} tokens',
            metadata={
                'customer_id': customer_id,
                'tokens_used': str(tokens_used)
            }
        )

        # Update DynamoDB record
        timestamp = datetime.utcnow().isoformat()
        item = {
            'customer_id': customer_id,
            'timestamp': timestamp,
            'tokens_used': Decimal(str(tokens_used)),
            'cost': Decimal(str(cost)),
            'billed': True,
            'stripe_charge_id': charge.id
        }

        table.put_item(Item=item)

        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Payment processed successfully',
                'charge_id': charge.id,
                'amount': cost,
                'customer_id': customer_id
            })
        }
    except stripe.error.StripeError as e:
        print(f"Stripe error: {e}")
        return {
            'statusCode': 400,
            'body': json.dumps({'error': str(e)})
        }
    except Exception as e:
        print(f"Error processing payment: {e}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }
