from datetime import date
from typing import Annotated

from aws_lambda_powertools.event_handler.api_gateway import Router
from aws_lambda_powertools.event_handler.exceptions import NotFoundError
from aws_lambda_powertools.event_handler.openapi.params import Query
from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair

from boto3clients import dynamodb_client
from config import ORDER_TABLE

router = Router()
dyn = DynamoDBPersistenceLayer(ORDER_TABLE, dynamodb_client)


class BillingNotFoundError(NotFoundError): ...


@router.get('/<org_id>/billing')
def billing(
    org_id: str,
    start_date: Annotated[date, Query()],
    end_date: Annotated[date, Query()],
):
    pk = f'BILLING#ORG#{org_id}'
    sk = f'START#{start_date}#END#{end_date}'
    billing = dyn.collection.get_item(
        KeyPair(pk, sk),
        exc_cls=BillingNotFoundError,
    )

    return billing | dyn.collection.query(
        KeyPair(
            pk=pk,
            sk=f'{sk}#ENROLLMENT',
        ),
        limit=150,
    )