"""Module to create shared external connections for PDP endpoints."""

import logging
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, Dict

from cerbos.sdk.client import AsyncCerbosClient
from fastapi import FastAPI
from owsclient import AsyncOwsClient
from owscontext import get_correlation_id, get_request_context

from pdp import config
from pdp.connectors.cerbos import get_cerbos_client
from pdp.connectors.dynamo import DynamoDbConnector
from pdp.connectors.features import SplitioClient, splitio_client_factory
from pdp.connectors.ows_account import OwsAccountClient
from pdp.connectors.ows_participant import OwsParticipantClient
from pdp.connectors.ows_permissions import OwsPermissionsClient
from pdp.connectors.redis_client import (
    RedisConnector,
    RedisConnectorFactoryError,
    redis_connector_factory,
)
from pdp.constants.dynamo import LIFESPAN_TEST_IDENTITY_UUID

logger = logging.getLogger(__name__)

IDENTITY_DDB_CONNECTOR_KEY = "PP_IDENTITY_DYNAMODB_CONNECTOR"
SPLITIO_CLIENT_KEY = "SPLITIO_CLIENT"
ASYNC_OWS_CLIENT_KEY = "ASYNC_OWS_CLIENT"
OWS_PERMISSIONS_CLIENT_KEY = "OWS_PERMISSIONS_CLIENT"
ASYNC_CERBOS_CLIENT_KEY = "ASYNC_CERBOS_CLIENT"
OWS_ACCOUNT_CLIENT_KEY = "OWS_ACCOUNT_CLIENT"
OWS_PARTICIPANT_CLIENT_KEY = "OWS_PARTICIPANT_CLIENT"
REDIS_CONNECTOR_KEY = "REDIS_CONNECTOR"
DATA_SOURCES: Dict[str, Any] = {}


class PdpDatasourcesLifespanError(Exception):
    pass


@asynccontextmanager
async def datasources_lifespan(
    app: FastAPI,
) -> AsyncIterator[Dict[str, Any]]:
    """Fastapi Lifespan function to create shared external data source connections."""
    ddb_connection = DynamoDbConnector(
        config.DYNAMODB_TABLE_IDENTITY,
        config.IDENTITY_HASH_KEY,
        config.IDENTITY_RANGE_KEY,
    )
    # Create a DDB connection at startup by sending a query() for a non-existent key.
    # PDP handlers will reuse this connection.
    _ = ddb_connection.query_by_hash_key(hash_key=LIFESPAN_TEST_IDENTITY_UUID)
    logger.info(
        "[lifespan] Initialized %s DynamoDbConnector",
        config.DYNAMODB_TABLE_IDENTITY,
    )

    splitio_client = splitio_client_factory()
    logger.info("[lifespan] Initialized split.io client")

    async_ows_client = AsyncOwsClient(
        environment=config.ENVIRONMENT,
        service_name=config.SERVICE_NAME,
        correlation_id_getter=get_correlation_id,
        request_context_getter=get_request_context,
    )
    logger.info("[lifespan] Initialized async ows client")

    ows_permissions_client = OwsPermissionsClient(async_ows_client=async_ows_client)
    logger.info("[lifespan] Initialized ows-permissions client")

    async_cerbos_client = get_cerbos_client()
    logger.info("[lifespan] Initialized async_cerbos_client")

    ows_account_client = OwsAccountClient(async_ows_client=async_ows_client)
    logger.info("[lifespan] Initialized async ows-account client")

    ows_participant_client = OwsParticipantClient(async_ows_client=async_ows_client)
    logger.info("[lifespan] Initialized ows-participant client")

    try:
        redis_connector = await redis_connector_factory()
    except RedisConnectorFactoryError as exc:
        raise PdpDatasourcesLifespanError(
            "[lifespan] Failed to initialize RedisConnector", exc
        )

    DATA_SOURCES[IDENTITY_DDB_CONNECTOR_KEY] = ddb_connection
    DATA_SOURCES[SPLITIO_CLIENT_KEY] = splitio_client
    DATA_SOURCES[ASYNC_OWS_CLIENT_KEY] = async_ows_client
    DATA_SOURCES[OWS_PERMISSIONS_CLIENT_KEY] = ows_permissions_client
    DATA_SOURCES[ASYNC_CERBOS_CLIENT_KEY] = async_cerbos_client
    DATA_SOURCES[OWS_ACCOUNT_CLIENT_KEY] = ows_account_client
    DATA_SOURCES[OWS_PARTICIPANT_CLIENT_KEY] = ows_participant_client
    DATA_SOURCES[REDIS_CONNECTOR_KEY] = redis_connector

    try:
        # Lines before 'yield' will be executed before the application starts.
        yield DATA_SOURCES
        # Lines after the 'yield' will be executed after the application has finished.
        shutdown_ddb_connection(ddb_connection=ddb_connection)
        shutdown_splitio_client(splitio_client=splitio_client)
        await shutdown_redis_connector(redis_connection=redis_connector)
        await shutdown_cerbos_client(cerbos_client=async_cerbos_client)
        await shutdown_ows_client(async_ows_client=async_ows_client)
    finally:
        DATA_SOURCES.clear()


def get_boto_connector() -> DynamoDbConnector:
    """
    Dependency injector method for DynamoDbConnector.

    Returns the shared PP_IDENTITY_DYNAMODB_CONNECTOR boto connection.
    """
    connector = DATA_SOURCES[IDENTITY_DDB_CONNECTOR_KEY]
    assert isinstance(connector, DynamoDbConnector)
    return connector


def get_splitio_client() -> SplitioClient:
    """Dependency injector method for the split.io client.

    Returns the shared SPLITIO_CLIENT connection.
    """
    connector = DATA_SOURCES[SPLITIO_CLIENT_KEY]
    assert isinstance(connector, SplitioClient)
    return connector


def get_async_ows_client() -> AsyncOwsClient:
    """Dependency injector method for general async owsclient.

    Returns the shared ASYNC_OWS_CLIENT client.
    """
    connector = DATA_SOURCES[ASYNC_OWS_CLIENT_KEY]
    assert isinstance(connector, AsyncOwsClient)
    return connector


def get_ows_permissions_client() -> OwsPermissionsClient:
    """Dependency injector method for the ows-permissions client.

    Returns the shared OWS_PERMISSIONS_CLIENT client.
    """
    connector = DATA_SOURCES[OWS_PERMISSIONS_CLIENT_KEY]
    assert isinstance(connector, OwsPermissionsClient)
    return connector


def get_ows_account_client() -> OwsAccountClient:
    """Dependency injector method for the ows-account client.

    Returns the shared OWS_ACCOUNT_CLIENT_KEY client.
    """
    connector = DATA_SOURCES[OWS_ACCOUNT_CLIENT_KEY]
    assert isinstance(connector, OwsAccountClient)
    return connector


def get_ows_participant_client() -> OwsParticipantClient:
    """Dependency injector method for the ows-participant client.

    Returns the shared OWS_PARTICIPANT_CLIENT_KEY client.
    """
    connector = DATA_SOURCES[OWS_PARTICIPANT_CLIENT_KEY]
    assert isinstance(connector, OwsParticipantClient)
    return connector


def get_async_cerbos_client() -> AsyncCerbosClient:
    """Dependency injector method for the cerbos client.

    Returns the shared ASYNC_CERBOS_CLIENT client.
    """
    connector = DATA_SOURCES[ASYNC_CERBOS_CLIENT_KEY]
    assert isinstance(connector, AsyncCerbosClient)
    return connector


def get_redis_connector() -> RedisConnector:
    """Dependency injector method for the redis client.

    Returns the shared REDIS_CONNECTOR client.
    """
    connector = DATA_SOURCES[REDIS_CONNECTOR_KEY]
    assert isinstance(connector, RedisConnector)
    return connector


def shutdown_ddb_connection(ddb_connection: DynamoDbConnector) -> None:
    """Shut down the DDB connection."""
    try:
        ddb_connection.client.close()
    except Exception as ex:
        logger.warning("[lifespan] Failed to shutdown DynamoDbConnector: %s", str(ex))
    else:
        logger.info("[lifespan] Closed DynamoDbConnector")


def shutdown_splitio_client(splitio_client: SplitioClient) -> None:
    """Shut down the Split.io connection."""
    try:
        splitio_client.destroy()
    except Exception as ex:
        logger.warning("[lifespan] Failed to shutdown split.io client: %s", str(ex))
    else:
        logger.info("[lifespan] Closed split.io client")


async def shutdown_cerbos_client(cerbos_client: AsyncCerbosClient) -> None:
    """Shut down the Cerbos connection."""
    try:
        await cerbos_client.close()  # type: ignore[no-untyped-call]
    except Exception as ex:
        logger.warning("[lifespan] Failed to shutdown cerbos client: %s", str(ex))
    else:
        logger.info("[lifespan] Closed cerbos client")


async def shutdown_ows_client(async_ows_client: AsyncOwsClient) -> None:
    """Shut down the ows client connection.

    The singleton ows_client is used by multiple ows-service client connectors.
    """
    try:
        await async_ows_client.close()
    except Exception as ex:
        logger.warning("[lifespan] Failed to shutdown ows client: %s", str(ex))
    else:
        logger.info("[lifespan] Closed ows client")


async def shutdown_redis_connector(redis_connection: RedisConnector) -> None:
    """Shut down the Redis connector."""
    try:
        await redis_connection.client.aclose()
    except Exception as ex:
        logger.warning("[lifespan] Failed to shutdown RedisConnector: %s", str(ex))
    else:
        logger.info("[lifespan] Closed RedisConnector")