"""Shared fixtures for DynamoDB adapter tests."""

from __future__ import annotations

import os

import boto3
import pytest

TABLE_NAMES = ["users", "posts", "comments"]
CUSTOM_KEY_TABLE = "custom_keys"
PK_ATTR = "_pk"
SK_ATTR = "_sk"
CUSTOM_PK_ATTR = "PK"
CUSTOM_SK_ATTR = "SK"


@pytest.fixture(scope="session")
def dynamodb_resource():
    port = os.environ.get("DYNAMODB_PORT", "8000")
    resource = boto3.resource(
        "dynamodb",
        endpoint_url=f"http://localhost:{port}",
        region_name="us-east-1",
        aws_access_key_id="dummy",
        aws_secret_access_key="dummy",
    )

    for name in TABLE_NAMES:
        try:
            resource.create_table(
                TableName=name,
                KeySchema=[
                    {"AttributeName": PK_ATTR, "KeyType": "HASH"},
                    {"AttributeName": SK_ATTR, "KeyType": "RANGE"},
                ],
                AttributeDefinitions=[
                    {"AttributeName": PK_ATTR, "AttributeType": "S"},
                    {"AttributeName": SK_ATTR, "AttributeType": "S"},
                ],
                BillingMode="PAY_PER_REQUEST",
            )
        except resource.meta.client.exceptions.ResourceInUseException:
            pass

    try:
        resource.create_table(
            TableName=CUSTOM_KEY_TABLE,
            KeySchema=[
                {"AttributeName": CUSTOM_PK_ATTR, "KeyType": "HASH"},
                {"AttributeName": CUSTOM_SK_ATTR, "KeyType": "RANGE"},
            ],
            AttributeDefinitions=[
                {"AttributeName": CUSTOM_PK_ATTR, "AttributeType": "S"},
                {"AttributeName": CUSTOM_SK_ATTR, "AttributeType": "S"},
            ],
            BillingMode="PAY_PER_REQUEST",
        )
    except resource.meta.client.exceptions.ResourceInUseException:
        pass

    yield resource

    for name in [*TABLE_NAMES, CUSTOM_KEY_TABLE]:
        resource.Table(name).delete()


def _clear_table(table, pk_attr: str, sk_attr: str) -> None:
    scan_kwargs = {
        "ProjectionExpression": "#pk, #sk",
        "ExpressionAttributeNames": {"#pk": pk_attr, "#sk": sk_attr},
    }
    while True:
        response = table.scan(**scan_kwargs)
        with table.batch_writer() as batch:
            for item in response["Items"]:
                batch.delete_item(Key={pk_attr: item[pk_attr], sk_attr: item[sk_attr]})
        if "LastEvaluatedKey" not in response:
            break
        scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]


@pytest.fixture
def dynamodb(dynamodb_resource):
    yield dynamodb_resource
    for name in TABLE_NAMES:
        _clear_table(dynamodb_resource.Table(name), PK_ATTR, SK_ATTR)
    _clear_table(dynamodb_resource.Table(CUSTOM_KEY_TABLE), CUSTOM_PK_ATTR, CUSTOM_SK_ATTR)
