"""UuidToIdExchangeTenantProxy."""

import asyncio
import logging
from typing import Awaitable, Callable, Dict, List, Optional
from uuid import UUID

from ddtrace.trace import tracer

from pdp.connectors.ows_account import (
    LookupCompanyBrandsResponse,
    LookupSubaccountsResponse,
    LookupVendorFetchFlags,
    LookupVendorsResponse,
    OwsAccountClient,
)
from pdp.connectors.redis_client import PydanticSchemaSerializer, RedisConnector
from pdp.constants.constants import TenantType
from pdp.constants.error import ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN
from pdp.fastapi.schemas.tenant import (
    IdExchangeTenantHierarchy,
    Tenant,
    UuidToIdExchangeTenant,
)
from pdp.proxies.helpers import lookup_with_error_handling
from pdp.proxies.multi_tenant_proxy import TenantHierarchyLookupError
from pdp.utils.tenant_cache import (
    TENANT_HIERARCHY_CACHE_TTL_BY_TENANT_TYPE,
    TENANT_HIERARCHY_KEY_PREFIX,
)

logger = logging.getLogger(__name__)


class UuidToIdExchangeLookupError(Exception):
    """UUID to ID lookup error."""

    pass


SUPPORTED_TENANT_TYPES: List[TenantType] = [
    TenantType.TENANT_TYPE_ACCOUNT,
    TenantType.TENANT_TYPE_SUBACCOUNT,
    TenantType.TENANT_TYPE_COMPANY_BRAND,
]


class UuidToIdExchangeTenantProxy:
    """Handles UUID to ID lookups for PDP."""

    def __init__(
        self,
        tenants: List[UuidToIdExchangeTenant],
        redis_connector: RedisConnector,
        ows_account_client: OwsAccountClient,
    ):
        """UuidToIdExchangeTenantProxy constructor."""
        self._tenants = self._get_unique_tenants(tenants)
        self._redis_connector = redis_connector
        self._ows_account_client = ows_account_client
        self._tenants_by_type = self._init_tenants_by_type()
        self._gathered_tenant_exchanges: Optional[
            Dict[UUID, Optional[IdExchangeTenantHierarchy]]
        ] = None

    @property
    def gathered_tenant_exchanges(
        self,
    ) -> dict[UUID, Optional[IdExchangeTenantHierarchy]]:
        """Access gathered IdExchangeTenantHierarchy objects, if non-None.

        Raises UuidToIdExchangeLookupError if accessed before ows-* lookups.
        """
        if self._gathered_tenant_exchanges is None:
            raise UuidToIdExchangeLookupError(
                "Attempted to access uuid exchange objects before ows-* lookups."
            )

        return self._gathered_tenant_exchanges

    @tracer.wrap()
    def get_tenant_exchange(
        self,
        uuid: UUID,
    ) -> Optional[IdExchangeTenantHierarchy]:
        """Get tenant exchange for tenant uuid."""
        return self.gathered_tenant_exchanges.get(uuid, None)

    @tracer.wrap()
    def _get_unique_tenants(
        self,
        tenants: List[UuidToIdExchangeTenant],
    ) -> List[UuidToIdExchangeTenant]:
        """Return unique, consistently ordered list of UuidToIdExchangeTenant."""
        return sorted(set(tenants))

    @tracer.wrap()
    def _init_tenants_by_type(
        self,
    ) -> Dict[TenantType, List[UuidToIdExchangeTenant]]:
        """
        Build data structure to organize UuidToIdExchangeTenant objects by Tenant Type.
        """
        # initialize the mapping with empty lists
        _tenants_by_type: Dict[TenantType, List[UuidToIdExchangeTenant]] = {
            tenant_type: [] for tenant_type in SUPPORTED_TENANT_TYPES
        }

        for tenant in self._tenants:
            if tenant.tenant_type in SUPPORTED_TENANT_TYPES:
                _tenants_by_type[tenant.tenant_type].append(tenant)
        return _tenants_by_type

    @tracer.wrap()
    async def _exchange_accounts_from_ows_account(
        self, tenant_uuids: Optional[List[UUID]] = None
    ) -> Dict[UUID, IdExchangeTenantHierarchy]:
        """Lookup tenants with a tenant_type account from ows-account."""
        if not tenant_uuids:
            account_tenants = self._tenants_by_type.get(
                TenantType.TENANT_TYPE_ACCOUNT, []
            )
            tenant_uuids = [tenant.tenant_uuid for tenant in account_tenants]

        if not len(tenant_uuids):
            return {}

        response: LookupVendorsResponse = await lookup_with_error_handling(
            lookup_cb=self._ows_account_client.lookup_vendors_by_uuids(
                uuids=tenant_uuids,
                fetch_flags=[LookupVendorFetchFlags.TENANT_HIERARCHY],
            ),
            http_error_message_400_string="Bad request to ows-account for uuid-to-id-exchange tenant proxy.",  # noqa: E501
            http_error_message_500_string="Account lookup by tenant_uuids failed for uuid-to-id-exchange tenant proxy.",  # noqa: E501
            unhandled_exception_string="Server error: account lookup by tenant_uuids failed for uuid-to-id-exchange tenant proxy.",  # noqa: E501
            lookup_error_cls=TenantHierarchyLookupError,
        )

        if not response or not len(response.vendors):
            logger.warning(ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN)
            return {}

        results = list(zip(tenant_uuids, response.vendors))
        missing_tenant_ids = [int(vendor_id) for vendor_id, item in results if not item]
        if len(missing_tenant_ids):
            # lookup_vendors_by_uuids does not return results without ids
            logger.warning(
                "ows-account lookup did not return results for every requested tenant:"
                f" {','.join(map(str, missing_tenant_ids))}"
            )

        missing_company_brand_uuids = [
            str(uuid)
            for uuid, vendor in results
            if vendor and not vendor.company_brand_uuid
        ]
        if len(missing_company_brand_uuids):
            raise TenantHierarchyLookupError(
                "tenant_hierarchy missing company_brand_uuid for:"
                f" '{','.join(missing_company_brand_uuids)}'"
            )

        missing_parent_company_uuids = [
            str(uuid)
            for uuid, vendor in results
            if vendor and not vendor.parent_company_uuid
        ]
        if len(missing_parent_company_uuids):
            raise TenantHierarchyLookupError(
                "tenant_hierarchy missing parent_company_uuid for:"
                f" '{','.join(missing_parent_company_uuids)}'"
            )
        return {
            vendor.uuid: IdExchangeTenantHierarchy(
                tenant_id=vendor.vendor_id,
                tenant_type=TenantType.TENANT_TYPE_ACCOUNT,
                tenant_uuid=vendor.uuid,
                company_brand=Tenant(
                    tenant_uuid=vendor.company_brand_uuid,
                    tenant_type=TenantType.TENANT_TYPE_COMPANY_BRAND,
                ),
                parent_company=Tenant(
                    tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY,
                    tenant_uuid=vendor.parent_company_uuid,
                ),
            )
            for vendor in response.vendors
            if vendor
        }

    @tracer.wrap()
    async def _exchange_subaccounts_from_ows_account(
        self,
        tenant_uuids: Optional[List[UUID]] = None,
    ) -> Dict[UUID, IdExchangeTenantHierarchy]:
        """Lookup tenants with tenant_type subaccount from ows-account."""
        if not tenant_uuids:
            subaccount_tenants = self._tenants_by_type.get(
                TenantType.TENANT_TYPE_SUBACCOUNT, []
            )
            tenant_uuids = [tenant.tenant_uuid for tenant in subaccount_tenants]

        if not len(tenant_uuids):
            return {}

        response: LookupSubaccountsResponse = await lookup_with_error_handling(
            lookup_cb=self._ows_account_client.lookup_subaccounts_by_uuids(
                uuids=tenant_uuids,
                fetch_flags=[LookupVendorFetchFlags.TENANT_HIERARCHY],
            ),
            http_error_message_400_string="Bad request to ows-account for uuid-to-id-exchange tenant proxy (subaccounts).",  # noqa: E501
            http_error_message_500_string="Subaccount lookup by subaccount_uuids failed for uuid-to-id-exchange tenant proxy.",  # noqa: E501
            unhandled_exception_string="Server error: subaccount lookup by subaccount_uuids failed for uuid-to-id-exchange tenant proxy.",  # noqa: E501
            lookup_error_cls=TenantHierarchyLookupError,
        )

        if not response or not len(response.subaccounts):
            logger.warning(ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN)
            return {}

        assert len(tenant_uuids) == len(response.subaccounts)
        results = list(zip(tenant_uuids, response.subaccounts))
        missing_tenant_uuids = [int(uuid) for uuid, item in results if not item]
        if len(missing_tenant_uuids):
            logger.warning(
                "ows-account lookup did not return results for every requested tenant:"
                f" {','.join(map(str, missing_tenant_uuids))}"
            )

        missing_company_brand_uuids = [
            str(uuid)
            for uuid, vendor in results
            if vendor and not vendor.company_brand_uuid
        ]
        if len(missing_company_brand_uuids):
            raise TenantHierarchyLookupError(
                "tenant_hierarchy missing company_brand_uuid for:"
                f" '{','.join(missing_company_brand_uuids)}'"
            )
        missing_vendor_uuids = [
            str(uuid)
            for uuid, subaccount in results
            if subaccount and not subaccount.vendor_uuid
        ]
        if len(missing_vendor_uuids):
            raise TenantHierarchyLookupError(
                "tenant_hierarchy missing vendor_uuid for:"
                f" '{','.join(missing_vendor_uuids)}'"
            )

        missing_parent_company_uuids = [
            str(uuid)
            for uuid, subaccount in results
            if subaccount and not subaccount.parent_company_uuid
        ]
        if len(missing_parent_company_uuids):
            raise TenantHierarchyLookupError(
                "tenant_hierarchy missing parent_company_uuid for:"
                f" '{','.join(missing_parent_company_uuids)}'"
            )

        return {
            subaccount.uuid: IdExchangeTenantHierarchy(
                tenant_id=subaccount.subaccount_id,
                tenant_type=TenantType.TENANT_TYPE_SUBACCOUNT,
                tenant_uuid=subaccount.uuid,
                account=Tenant(
                    tenant_uuid=subaccount.vendor_uuid,
                    tenant_type=TenantType.TENANT_TYPE_ACCOUNT,
                ),
                company_brand=Tenant(
                    tenant_uuid=subaccount.company_brand_uuid,
                    tenant_type=TenantType.TENANT_TYPE_COMPANY_BRAND,
                ),
                parent_company=Tenant(
                    tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY,
                    tenant_uuid=subaccount.parent_company_uuid,
                ),
            )
            for subaccount in response.subaccounts
            if subaccount
        }

    @tracer.wrap()
    async def _exchange_company_brands_from_ows_account(
        self,
        tenant_uuids: Optional[List[UUID]] = None,
    ) -> Dict[UUID, IdExchangeTenantHierarchy]:
        """Lookup hierarchies for tenants with tenant_type company_brand."""

        if not tenant_uuids:
            # gather_tenant_exchange() will pass along tenant_uuids.
            # For other callers, load the company brand type ids here.
            company_brand_tenants = self._tenants_by_type.get(
                TenantType.TENANT_TYPE_COMPANY_BRAND, []
            )
            tenant_uuids = [tenant.tenant_uuid for tenant in company_brand_tenants]

        if not len(tenant_uuids):
            return {}

        response: LookupCompanyBrandsResponse = await lookup_with_error_handling(
            lookup_cb=self._ows_account_client.lookup_company_brands_by_uuids(
                uuids=tenant_uuids,
            ),
            http_error_message_400_string="Bad request to ows-account for multiple company_brand tenants.",  # noqa: E501
            http_error_message_500_string="tenant_hierarchy lookup failed for multiple company_brand tenants",  # noqa: E501
            unhandled_exception_string="Server error: tenant_hierarchy lookup failed for multiple company_brand tenants",  # noqa: E501
            lookup_error_cls=TenantHierarchyLookupError,
        )

        if not response or not len(response.company_brands):
            logger.warning(ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN)
            return {}

        # Figure out what was missing to log warning
        assert len(tenant_uuids) == len(response.company_brands)
        results = list(zip(tenant_uuids, response.company_brands))
        missing_tenant_uuids = [str(uuid) for uuid, item in results if not item]
        if len(missing_tenant_uuids):
            # lookup_company_brands_by_uuids does not return results for
            # company_brand uuids not found in the DB
            logger.warning(
                "ows-account lookup did not return results for every requested tenant:"  # noqa: E501
                f" {','.join(missing_tenant_uuids)}"
            )
        missing_parent_company_uuids = [
            str(uuid)
            for uuid, company_brand in results
            if company_brand and not company_brand.parent_company_uuid
        ]
        if len(missing_parent_company_uuids):
            raise TenantHierarchyLookupError(
                "tenant_hierarchy missing parent_company_uuid for:"
                f" '{','.join(missing_parent_company_uuids)}'"
            )

        return {
            company_brand.uuid: IdExchangeTenantHierarchy(
                tenant_id=company_brand.company_brand_id,
                tenant_type=TenantType.TENANT_TYPE_COMPANY_BRAND,
                tenant_uuid=company_brand.uuid,
                company_brand=None,
                parent_company=Tenant(
                    tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY,
                    tenant_uuid=company_brand.parent_company_uuid,
                ),
            )
            for company_brand in response.company_brands
            if company_brand
        }

    @tracer.wrap()
    async def _do_tenant_exchange(
        self,
        tenant_type: TenantType,
        tenant_exchange_lookup_cb: Callable[
            [List[UUID]],
            Awaitable[Dict[UUID, IdExchangeTenantHierarchy]],
        ],
    ) -> Dict[UUID, Optional[IdExchangeTenantHierarchy]]:
        """
        Fetch the tenant exchanges by type.

        Fetch tenant exchanges by tenant_type from cache
        or from the 'source of truth' service.

        Args:
            tenant_type: account, subaccount, etc..
            tenant_exchange_lookup_cb: callback function to perform \
                ows-* hierarchy lookup.
        """
        # Fetch tenant exchanges from cache if available
        tenant_uuids = [
            tenant.tenant_uuid for tenant in self._tenants_by_type[tenant_type]
        ]
        if not tenant_uuids:
            return {}

        tenant_exchanges = await self._get_tenant_exchange_from_cache(
            self._tenants_by_type[tenant_type]
        )

        missing_tenant_uuids_from_cache = [
            tenant_uuid
            for tenant_uuid in tenant_uuids
            if not tenant_exchanges.get(tenant_uuid, None)
        ]

        if not missing_tenant_uuids_from_cache:
            return tenant_exchanges

        # Fetch tenant exchange(s) from the source of truth
        tenants_from_source_of_truth = await tenant_exchange_lookup_cb(
            missing_tenant_uuids_from_cache
        )

        # If any, save fetched tenant exchange(s) to cache
        save_to_cache_items = [
            id_exchange_tenant_hierarchy
            for id_exchange_tenant_hierarchy in tenants_from_source_of_truth.values()
            if id_exchange_tenant_hierarchy
        ]
        if save_to_cache_items:
            save_result = await self._save_tenant_exchange_to_cache(
                entries=save_to_cache_items, tenant_type=tenant_type
            )
            if False in save_result.values():
                logger.error(
                    "Failed to save tenant exchanges to cache for %s",
                    save_result,
                )

        # Merge results to return all tenant exchanges
        for tenant in tenants_from_source_of_truth:
            tenant_exchanges[tenant] = tenants_from_source_of_truth[tenant]
        return tenant_exchanges

    @tracer.wrap()
    async def gather_tenant_exchange(
        self,
    ) -> Dict[UUID, Optional[IdExchangeTenantHierarchy]]:
        """Lookup IDs and hierarchies for all tenant types and merge the results."""
        self._gathered_tenant_exchanges = {}

        tenant_exchanges = await asyncio.gather(
            self._do_tenant_exchange(
                TenantType.TENANT_TYPE_ACCOUNT,
                self._exchange_accounts_from_ows_account,  # noqa: E501
            ),
            self._do_tenant_exchange(
                TenantType.TENANT_TYPE_SUBACCOUNT,
                self._exchange_subaccounts_from_ows_account,
            ),
            self._do_tenant_exchange(
                TenantType.TENANT_TYPE_COMPANY_BRAND,
                self._exchange_company_brands_from_ows_account,
            ),
        )
        for exchanges in tenant_exchanges:
            self._gathered_tenant_exchanges.update(exchanges)

        return self._gathered_tenant_exchanges

    @tracer.wrap()
    async def _get_tenant_exchange_from_cache(
        self, tenants: List[UuidToIdExchangeTenant]
    ) -> Dict[UUID, Optional[IdExchangeTenantHierarchy]]:
        """Get cached tenant exchange objects."""
        if not len(tenants):
            return {}

        cache_keys = []
        tenant_uuids = []
        for tenant in tenants:
            tenant_uuids.append(tenant.tenant_uuid)
            cache_keys.append(tenant.to_cache_key())

        cached_items: List[
            Optional[IdExchangeTenantHierarchy]
        ] = await self._redis_connector.mget(
            keys=cache_keys,
            serializer=PydanticSchemaSerializer(IdExchangeTenantHierarchy),
        )
        return dict(zip(tenant_uuids, cached_items))

    @tracer.wrap()
    async def _save_tenant_exchange_to_cache(
        self,
        entries: List[IdExchangeTenantHierarchy],
        tenant_type: TenantType,
        cache_prefix: str = TENANT_HIERARCHY_KEY_PREFIX,
    ) -> Dict[str, bool]:
        """Save Tenants' Exchange objects to cache."""
        if not entries:
            return {}
        # Form cache keys
        cache_entries = {}
        for entry in entries:
            cache_keys = entry.to_cache_keys()
            # Save the tenant hierarchy to both cache keys:
            # "tenant_hierarchy@tenant_id#123|tenant_type#account": IdExchangeTenantHierarchy #noqa E501
            # "tenant_hierarchy@tenant_type#account|tenant_uuid#61577c79-d788-4b88-bf6a-2def82e0b4ed": IdExchangeTenantHierarchy #noqa E501
            for cache_key in cache_keys:
                cache_entries[cache_key] = entry
        # Save cache entries
        mset_results = await self._redis_connector.mset_with_pipeline(
            cache_entries,
            serializer=PydanticSchemaSerializer(IdExchangeTenantHierarchy),
            ttl=TENANT_HIERARCHY_CACHE_TTL_BY_TENANT_TYPE[tenant_type],
        )

        # Set the mset response
        mset_response: Dict[str, bool] = {}
        for key in cache_entries.keys():
            mset_response[key] = mset_results.get(key, False)
        return mset_response