"""Logic for Cerbos."""

import logging
import uuid
from json import JSONDecodeError
from typing import Any, Dict, List, Optional

from cerbos.sdk.client import AsyncCerbosClient
from cerbos.sdk.model import CheckResourcesResponse as CerbosCheckResourcesResponse
from cerbos.sdk.model import CheckResourcesResult as CerbosCheckResourcesResult
from cerbos.sdk.model import Effect as CerbosEffect
from cerbos.sdk.model import Principal as CerbosPrincipal
from cerbos.sdk.model import Resource as CerbosResource
from cerbos.sdk.model import ResourceList as CerbosResourceList
from ddtrace.trace import tracer
from fastapi import HTTPException
from httpcore import ConnectError
from owscontext import get_correlation_id
from pydantic import UUID4

from pdp.config import CERBOS_BATCH_SIZE
from pdp.connectors.features import BooleanFeature, SplitioClient
from pdp.connectors.ows_account import OwsAccountClient
from pdp.connectors.ows_participant import OwsParticipantClient
from pdp.connectors.redis_client import RedisConnector
from pdp.constants.constants import AUTH_EFFECT_ALLOW, AUTH_EFFECT_DENY
from pdp.constants.features import FEATURE_FLAG_PP_VENDOR_FEATURES_LOOKUP
from pdp.fastapi.schemas.check_resources import CheckResourcesRequest
from pdp.fastapi.schemas.identity import (
    CheckResourceActionResult,
    CheckResourcesResponse,
    TenantRoles,
)
from pdp.fastapi.schemas.identity import Resource as PdpResource
from pdp.fastapi.schemas.principal import Principal
from pdp.fastapi.schemas.tenant import TenantWithMaybeHierarchy
from pdp.proxies.multi_tenant_proxy import MultiTenantProxy
from pdp.utils.cerbos_policy_metadata import get_cerbos_policy_metadata

CERBOS_EFFECT_MAPPING = {
    CerbosEffect.DENY: AUTH_EFFECT_DENY,
    CerbosEffect.ALLOW: AUTH_EFFECT_ALLOW,
}

logger = logging.getLogger(__name__)


class CerbosLookupError(Exception):
    """Cerbos lookup error."""

    pass


CERBOS_ATTRIBUTES_ERROR_MESSAGE = (
    "Schema validation error found in the request. "
    "Verify that the request contains the fields required by the resourceSchema."
)


@tracer.wrap()
def _merge_cerbos_responses(
    check_resources_request: CheckResourcesRequest,
    batch_responses: List[CerbosCheckResourcesResponse],
    include_resource_attributes: bool = False,
) -> CheckResourcesResponse:
    """Merge the Cerbos responses into a PDP CheckResourcesResponse."""
    result = CheckResourcesResponse(
        request_id=batch_responses[0].request_id, resources=[]
    )
    batch_results: List[CerbosCheckResourcesResult] = []
    # Put all of the CheckResourcesResults from each
    # CheckResourcesResponse into a single list.
    for batch_response in batch_responses:
        if batch_response.results:
            batch_results.extend(batch_response.results)
        else:
            logger.warning("batch_response.results is None")

    # Confirm that the request and response lists are the same length
    if len(check_resources_request.resources) != len(batch_results):
        raise HTTPException(
            status_code=500,
            detail="Batched Cerbos responses does not match request length.",
        )

    # Zip the list original request resources with the list of results.
    for resource_request, resource_result in zip(
        check_resources_request.resources, batch_results
    ):
        result.resources.append(
            CheckResourceActionResult(
                resource=PdpResource(
                    resource_id=resource_result.resource.id,
                    resource_type=resource_result.resource.kind,
                    attributes=resource_request.resource.attributes
                    if include_resource_attributes
                    else {},
                ),
                action=list(resource_result.actions.keys())[0],
                effect=CERBOS_EFFECT_MAPPING[list(resource_result.actions.values())[0]],
                errors={"validation_errors": resource_result.validation_errors},
            )
        )
    return result


@tracer.wrap()
async def _paginated_resource_check(
    check_resources_request: CheckResourcesRequest,
    cerbos_client: AsyncCerbosClient,
    principal: CerbosPrincipal,
    include_resource_attributes: bool = False,
) -> CheckResourcesResponse:
    """Build a paginated resource list."""
    batch_results: List[CerbosCheckResourcesResponse] = []

    for offset in range(0, len(check_resources_request.resources), CERBOS_BATCH_SIZE):
        resource_batch = check_resources_request.resources[
            offset : offset + CERBOS_BATCH_SIZE
        ]

        resource_list = await _build_resource_list(
            CheckResourcesRequest(
                resources=resource_batch,
                include_resource_attributes_in_response=include_resource_attributes,
            )
        )

        request_id = get_correlation_id()

        # Call Cerbos for batch
        try:
            _response = await cerbos_client.check_resources(
                principal, resource_list, request_id
            )
            if _response.results is None:
                # Return a 500 to the client.
                raise CerbosLookupError(f"Invalid response from cerbos. {_response}")

            # Store Cerbos response for each batch.
            batch_results.append(_response)
        except KeyError as e:
            logger.warning(
                "%s Invalid {%s}",
                CERBOS_ATTRIBUTES_ERROR_MESSAGE,
                str(e.args),
                exc_info=True,
            )
            raise HTTPException(
                status_code=400,
                detail=f"{CERBOS_ATTRIBUTES_ERROR_MESSAGE} Invalid {str(e.args)}",
            )
        except ConnectError:
            logger.error("Cerbos is unavailable.")
            # default_error_handler will log the trace
            raise
        except JSONDecodeError as e:
            # AWS WAF errors will trigger a JSONDecodeError because:
            #  - 503 errors from WAF return an HTML-formatted response.
            #  - The cerbos SDK assumes it will only receive a JSON-formatted response
            #    from the cerbos-server.
            raise HTTPException(status_code=500, detail="AWS WAF or ELB error.") from e
        except Exception:
            logger.error("Unhandled Cerbos error.")
            # default_error_handler will log the trace
            raise

    return _merge_cerbos_responses(
        check_resources_request,
        batch_results,
        include_resource_attributes=include_resource_attributes,
    )


@tracer.wrap()
async def check_resources(
    identity_uuid: str,
    check_resources_request: CheckResourcesRequest,
    cerbos_client: AsyncCerbosClient,
    pdp_tenant_roles: Dict[uuid.UUID, TenantRoles],
    ows_account_client: OwsAccountClient,
    ows_participant_client: OwsParticipantClient,
    redis_connector: RedisConnector,
    ows_permissions_tenant_roles: Optional[Dict[uuid.UUID, TenantRoles]] = None,
    include_resource_attributes: bool = False,
    authenticated_identity_uuid: Optional[UUID4] = None,
    principal: Optional[Principal] = None,
    splitio_client: Optional[SplitioClient] = None,
) -> CheckResourcesResponse:
    """Make cerbos request to check identity's authorization to resources.

    Args:
        identity_uuid: the Principal's identity uuid,
            whose authorization is being checked
        check_resources_request: the resources the Principal wants to perform actions on
        cerbos_client: Cerbos client connection
        pdp_tenant_roles: the Principal's TenantRoles, according to PDP
        ows_account_client: OwsAccount client connection
        ows_participant_client: OwsParticipant client connection
        redis_connector: Redis client connection
        ows_permissions_tenant_roles: the Principal's TenantRoles,
            according to ows-permissions. This is optional, and
            only expected when checking authorization on Role Administration endpoints
            WE EXPECT TO DEPRECATE THIS WHEN PDP IS SOURCE OF TRUTH
        include_resource_attributes: Flag to include resource attributes
            provided by the client in the response.
        authenticated_identity_uuid: Used identity-based feature flagging behavior
        principal: object describing the principal, whose authorization is being checked
            Future iterations will make this required to supplant the following params:
                - identity_uuid
                - pdp_tenant_roles
                - ows_permissions_tenant_roles
    """
    if principal:
        cerbos_principal = principal.get_cerbos_principal()
    else:
        cerbos_principal = _build_principal(
            identity_uuid,
            pdp_tenant_roles=pdp_tenant_roles,
            ows_permissions_tenant_roles=ows_permissions_tenant_roles,
        )

    if splitio_client is not None:
        vendor_features_flag = BooleanFeature(
            client=splitio_client,
            feature_name=FEATURE_FLAG_PP_VENDOR_FEATURES_LOOKUP,
        )
        if vendor_features_flag.is_enabled():
            await get_cerbos_policy_metadata(redis_connector)

    check_resources_request = await _hydrate_resources_with_hierarchy_as_needed(
        check_resources_request=check_resources_request,
        redis_connector=redis_connector,
        ows_account_client=ows_account_client,
        ows_participant_client=ows_participant_client,
    )

    result = await _paginated_resource_check(
        check_resources_request=check_resources_request,
        cerbos_client=cerbos_client,
        principal=cerbos_principal,
        include_resource_attributes=include_resource_attributes,
    )
    return result


@tracer.wrap()
def _build_principal(
    identity_uuid: str,
    pdp_tenant_roles: Dict[uuid.UUID, TenantRoles],
    ows_permissions_tenant_roles: Optional[Dict[uuid.UUID, TenantRoles]] = None,
) -> CerbosPrincipal:
    """Get all data to represent an identity_uuid as a principal.

    Paginate over all roles by identity to aggregate an identity's roles

    Args:
        identity_uuid (str): UUID to match
        identity_ddb_connector: dependency on DynamoDB
        ows_permissions_tenant_roles: principal attributes, only expected when
            we are checking authorization on Role Administration endpoints
    """
    tenants = pdp_tenant_roles

    cerbos_principal_tenants: Dict[str, Any] = {}
    for tenant in tenants.values():
        cerbos_principal_tenants[str(tenant.tenant_uuid)] = (
            tenant.as_cerbos_principal_tenants_attribute()
        )

    if ows_permissions_tenant_roles:
        for tenant_role in ows_permissions_tenant_roles.values():
            tenant_uuid = str(tenant_role.tenant_uuid)
            if tenant_uuid in cerbos_principal_tenants:
                # Use the roles
                cerbos_principal_tenants[tenant_uuid]["roles"].update(
                    tenant_role.as_cerbos_principal_tenants_attribute()["roles"]
                )
            else:
                cerbos_principal_tenants[tenant_uuid] = (
                    tenant_role.as_cerbos_principal_tenants_attribute()
                )

    return CerbosPrincipal(
        id=str(identity_uuid),
        roles={"user"},
        attr={
            "type": "human",
            "tenants": cerbos_principal_tenants,
        },
    )


@tracer.wrap()
async def _build_resource_list(
    check_resources_request: CheckResourcesRequest,
) -> CerbosResourceList:
    """Transform request to check resource actions for Cerbos SDK client."""
    resources = CerbosResourceList()
    for check_resource in check_resources_request.resources:
        resources.add(
            CerbosResource(
                id=str(check_resource.resource.resource_id),
                kind=check_resource.resource.resource_type,
                attr=check_resource.resource.attributes,
            ),
            actions={check_resource.action},
        )

    return resources


@tracer.wrap()
async def _hydrate_resources_with_hierarchy_as_needed(
    check_resources_request: CheckResourcesRequest,
    redis_connector: RedisConnector,
    ows_account_client: OwsAccountClient,
    ows_participant_client: OwsParticipantClient,
) -> CheckResourcesRequest:
    """Hydrate resources' attributes with tenant hierarchy if not already present."""
    # Get all the tenants from check_request's resources
    tenants: List[Optional[TenantWithMaybeHierarchy]] = []
    for check_resource_action in check_resources_request.resources:
        tenant = check_resource_action.get_tenant_with_maybe_hierarchy()
        tenants.append(tenant)

    # Create MultiTenantProxy and get the tenant hierarchies
    mtp = MultiTenantProxy(
        tenants=[tenant for tenant in tenants if tenant and tenant.needs_hierarchy()],
        redis_client=redis_connector,
        ows_account_client=ows_account_client,
        ows_participant_client=ows_participant_client,
    )
    await mtp.gather_tenant_hierarchies()

    # Iterate over check_request to update the hierarchy
    for tenant, check_resource_action in zip(
        tenants, check_resources_request.resources
    ):
        if tenant and tenant.needs_hierarchy():
            hierarchy = mtp.get_tenant_hierarchy(tenant.tenant_uuid)
            if hierarchy:
                check_resource_action.resource.attributes["tenant"][
                    "tenant_hierarchy"
                ] = hierarchy.to_array()

    return check_resources_request