import json

import requests
from aws_lambda_powertools import Logger
from aws_lambda_powertools.shared.json_encoder import Encoder
from aws_lambda_powertools.utilities.data_classes import (
    EventBridgeEvent,
    event_source,
)
from aws_lambda_powertools.utilities.typing import LambdaContext
from layercake.dateutils import now
from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair

from boto3clients import dynamodb_client, s3_client
from config import BUCKET_NAME, ORDER_TABLE

WEASYPRINT_API = 'https://weasyprint.saladeaula.digital'
TEMPLATE_S3_URI = 's3://saladeaula.digital/billing/template.html'

logger = Logger(__name__)
order_layer = DynamoDBPersistenceLayer(ORDER_TABLE, dynamodb_client)


@event_source(data_class=EventBridgeEvent)
@logger.inject_lambda_context
def lambda_handler(event: EventBridgeEvent, context: LambdaContext) -> bool:
    keys = event.detail['keys']
    now_ = now()
    # Key pattern `BILLING#ORG#{org_id}`
    *_, org_id = keys['id'].split('#')
    # Key pattern `START#{start_date}#END#{end_date}#SCHEDULE#AUTO_CLOSE`
    _, start_date, _, end_date, *_ = keys['sk'].split('#')

    result = order_layer.collection.query(
        KeyPair(
            pk=keys['id'],
            sk=f'START#{start_date}#END#{end_date}#ENROLLMENT',
        ),
        limit=150,
    )

    json_data = json.dumps(
        {
            'template_s3_uri': TEMPLATE_S3_URI,
            'template_vars': {
                'start_date': start_date,
                'end_date': end_date,
                'items': result['items'],
            },
        },
        cls=Encoder,
    )
    # Send template URI and data to WeasyPrint API to generate a PDF
    r = requests.post(WEASYPRINT_API, data=json_data)
    r.raise_for_status()

    object_key = f'billing/{org_id}/{start_date}_{end_date}.pdf'
    s3_uri = f's3://{BUCKET_NAME}/{object_key}'

    try:
        s3_client.put_object(
            Bucket=BUCKET_NAME,
            Key=object_key,
            Body=r.content,
            ContentType='application/pdf',
        )
    except Exception as exc:
        logger.exception(exc)
        raise

    with order_layer.transact_writer() as transact:
        transact.update(
            key=KeyPair(
                pk=keys['id'],
                sk=f'START#{start_date}#END#{end_date}',
            ),
            update_expr='SET #status = :status, s3_uri = :s3_uri, \
                        updated_at = :updated_at',
            expr_attr_names={'#status': 'status'},
            expr_attr_values={
                ':status': 'CLOSED',
                ':s3_uri': s3_uri,
                ':updated_at': now_,
            },
            cond_expr='attribute_exists(sk)',
        )
        transact.put(
            item={
                'id': keys['id'],
                'sk': '{sk}#EXECUTED'.format(sk=keys['sk']),
                'created_at': now_,
            }
        )

    logger.info(f'PDF uploaded successfully to {s3_uri}')
    return True