from typing import Any, cast, TYPE_CHECKING

import boto3

from aws_testing_utils import config

from .logger import log

if TYPE_CHECKING:
    from mypy_boto3_dynamodb import DynamoDBServiceResource
    from mypy_boto3_dynamodb.service_resource import Table


class DynamoDBHandler:
    """Handles interactions with AWS DynamoDB."""

    dynamodb_resource: 'DynamoDBServiceResource'

    def __init__(self) -> None:
        self.dynamodb_resource = boto3.resource('dynamodb', config.AWS_REGION)

    def table(self, table_name: str) -> 'Table':
        """Returns a DynamoDB Table resource for the given table name."""
        return self.dynamodb_resource.Table(table_name)

    def get(self, table_name: str, key_name: str, key_value: str) -> dict[str, Any]:
        """Returns an item from a table by primary key. Raises if not found.

        Args:
            table_name: DynamoDB table name.
            key_name: Primary key attribute name.
            key_value: Primary key attribute value.
        """
        log.info(f'Getting {key_name}={key_value} from {table_name}')
        response = self.table(table_name).get_item(Key={key_name: f'{key_value}'})
        try:
            item = response['Item']
            log.info(f'item found: {item}')
        except KeyError:
            raise LookupError(
                f'No results found for {key_name}={key_value} in {table_name}'
            ) from None
        return item

    def put(self, data: dict[str, Any], table_name: str) -> dict[str, Any]:
        """Puts an item into a table."""
        log.info(f'Putting {data} into {table_name}')
        return cast(dict[str, Any], self.table(table_name).put_item(Item=data))