"""IdToUuidExchangeTenantProxy."""

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

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,
    IdToUuidExchangeTenant,
    Tenant,
)
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__)

# These are the tenant types for which an ID to UUID exchange can occur
SUPPORTED_TENANT_TYPES: List[TenantType] = [
    TenantType.TENANT_TYPE_ACCOUNT,
    TenantType.TENANT_TYPE_SUBACCOUNT,
    TenantType.TENANT_TYPE_COMPANY_BRAND,
]


class IdToUuidExchangeLookupError(Exception):
    """Id To UUID Exchange Lookup error."""

    pass


class IdToUuidExchangeTenantProxy:
    """Handles multiple tenants to proxy requests to redis or relevant microservices."""

    def __init__(
        self,
        tenants: List[IdToUuidExchangeTenant],
        redis_connector: RedisConnector,
        ows_account_client: OwsAccountClient,
    ):
        """Multi-Tenant Proxy 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[TenantType, Dict[str, Optional[IdExchangeTenantHierarchy]]]
        ] = None

    @property
    def gathered_tenant_exchanges(
        self,
    ) -> Dict[TenantType, Dict[str, Optional[IdExchangeTenantHierarchy]]]:
        """Access gathered hierarchies, if non-None.

        Example gathered_tenant_exchanges object:
        {
            TenantType.TENANT_TYPE_ACCOUNT: [
                {
                    11117: IdExchangeTenantHierarchy
                }
            ],
            TenantType.TENANT_TYPE_SUBACCOUNT: [
                {
                    11117: IdExchangeTenantHierarchy
                }
            ]
        }

        Raises IdToUuidExchangeLookupError if accessed before ows-* lookups.
        """
        if self._gathered_tenant_exchanges is None:
            raise IdToUuidExchangeLookupError(
                "Attempted to access gathered_hierarchies before ows-* lookups."
            )

        return self._gathered_tenant_exchanges

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

    @tracer.wrap()
    def _init_tenants_by_type(
        self,
    ) -> Dict[TenantType, List[IdToUuidExchangeTenant]]:
        """
        Build data structure to organize IdToUuidExchangeTenant objects by Tenant Type.

        Example object:
            {
                TenantType.TENANT_TYPE_ACCOUNT: [
                    IdToUuidExchangeTenant(
                        tenant_type=TenantType.TENANT_TYPE_ACCOUNT,
                        tenant_id=11117
                    )
                ],
                TenantType.TENANT_TYPE_SUBACCOUNT: [
                    IdToUuidExchangeTenant(
                        tenant_type=TenantType.TENANT_TYPE_SUBACCOUNT,
                        tenant_id=668894
                    )
                ],
            }
        """
        supported_tenant_types = SUPPORTED_TENANT_TYPES[:]

        # initialize the mapping with empty lists
        _tenants_by_type: Dict[TenantType, List[IdToUuidExchangeTenant]] = {
            tenant_type: [] for tenant_type in supported_tenant_types
        }
        for tenant in self._tenants:
            if tenant.tenant_type not in supported_tenant_types:
                logger.info(
                    "Unsupported tenant type %s. Skipping tenant %s",
                    tenant.tenant_type,
                    tenant,
                )
                continue
            _tenants_by_type[tenant.tenant_type].append(tenant)

        return _tenants_by_type

    @tracer.wrap()
    def get_tenant_exchange(
        self,
        tenant_type: TenantType,
        id: str | int,
    ) -> Optional[IdExchangeTenantHierarchy]:
        """Get tenant exchange for tenant type / id."""
        if not self.gathered_tenant_exchanges.get(tenant_type):
            return None
        return self.gathered_tenant_exchanges[tenant_type].get(str(id))

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

        cache_keys = []
        tenant_ids = []
        for tenant in tenants:
            tenant_ids.append(str(tenant.tenant_id))
            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_ids, 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

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

        if not len(tenant_ids):
            return {}

        response: LookupVendorsResponse = await lookup_with_error_handling(
            lookup_cb=self._ows_account_client.lookup_vendors_by_ids(
                vendor_ids=[int(tenant_id) for tenant_id in tenant_ids],
                fetch_flags=[LookupVendorFetchFlags.TENANT_HIERARCHY],
            ),
            http_error_message_400_string="Bad request to ows-account for id-to-uuid-exchange tenant proxy.",  # noqa: E501
            http_error_message_500_string="Account lookup by vendor_ids failed for id-to-uuid-exchange tenant proxy.",  # noqa: E501
            unhandled_exception_string="Server error: account lookup "
            "by vendor_ids failed for id-to-uuid-exchange tenant proxy.",
            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_ids, response.vendors))
        missing_tenant_uuids = [str(uuid) for uuid, item in results if not item]
        if len(missing_tenant_uuids):
            # lookup_vendors_by_ids does not return results without uuids
            logger.warning(
                "ows-account lookup did not return results for every requested tenant:"
                f" {','.join(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_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 {
            str(vendor.vendor_id): 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_uuid=vendor.parent_company_uuid,
                    tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY,
                ),
            )
            for vendor in response.vendors
            if vendor
        }

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

        if not len(tenant_ids):
            return {}

        response: LookupSubaccountsResponse = await lookup_with_error_handling(
            lookup_cb=self._ows_account_client.lookup_subaccounts_by_ids(
                subaccount_ids=[int(subaccount_id) for subaccount_id in tenant_ids],
                fetch_flags=[LookupVendorFetchFlags.TENANT_HIERARCHY],
            ),
            http_error_message_400_string="Bad request to ows-account for id-to-uuid-exchange tenant proxy (subaccounts).",  # noqa: E501
            http_error_message_500_string="Subaccount lookup by subaccount_ids failed for id-to-uuid-exchange tenant proxy.",  # noqa: E501
            unhandled_exception_string="Server error: subaccount lookup by subaccount_ids failed for id-to-uuid-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 {}

        results = list(zip(tenant_ids, response.subaccounts))
        missing_tenant_uuids = [str(uuid) for uuid, item in results if not item]
        if len(missing_tenant_uuids):
            logger.warning(
                "ows-account lookup did not return results for requested subaccounts:"  # noqa E501
                f" {','.join(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 {
            str(subaccount.subaccount_id): 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_uuid=subaccount.parent_company_uuid,
                    tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY,
                ),
            )
            for subaccount in response.subaccounts
            if subaccount
        }

    @tracer.wrap()
    async def _exchange_company_brands_from_ows_account(
        self,
        tenant_ids: Optional[List[int | str]] = None,
    ) -> Dict[str, IdExchangeTenantHierarchy]:
        """Lookup hierarchies for tenants with tenant_type company_brand."""
        if not tenant_ids:
            # gather_tenant_exchange() will pass along tenant_ids.
            # For other callers, load the company brand type ids here.
            account_tenants = self._tenants_by_type.get(
                TenantType.TENANT_TYPE_COMPANY_BRAND, []
            )
            tenant_ids = [tenant.tenant_id for tenant in account_tenants]

        if not len(tenant_ids):
            return {}

        response: LookupCompanyBrandsResponse = await lookup_with_error_handling(
            lookup_cb=self._ows_account_client.lookup_company_brands_by_ids(
                company_brand_ids=[int(tenant_id) for tenant_id in tenant_ids],
            ),
            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
        results = list(zip(tenant_ids, 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 {
            str(company_brand.company_brand_id): 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[int | str]],
            Awaitable[
                Dict[
                    str,
                    IdExchangeTenantHierarchy,
                ]
            ],
        ],
    ) -> Dict[str, IdExchangeTenantHierarchy | None]:
        """
        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_ids = [tenant.tenant_id for tenant in self._tenants_by_type[tenant_type]]
        if not tenant_ids:
            return {}

        tenant_exchanges = await self._get_tenant_exchange_from_cache(
            self._tenants_by_type[tenant_type]
        )
        missing_tenant_ids_from_cache = [
            tenant_id
            for tenant_id in tenant_ids
            if not tenant_exchanges.get(str(tenant_id), None)
        ]
        if not missing_tenant_ids_from_cache:
            return tenant_exchanges

        # Fetch tenant exchange(s) from source of truth
        tenants_from_source_of_truth = await tenant_exchange_lookup_cb(
            missing_tenant_ids_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.info(
                    "Failed to save some tenant exchanges to cache: %s",
                    save_result,
                )

        # Merge results to return all tenant exchanges
        tenant_exchanges.update(**tenants_from_source_of_truth)
        return tenant_exchanges

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

        tenant_exchanges = await asyncio.gather(
            self._do_tenant_exchange(
                TenantType.TENANT_TYPE_ACCOUNT,
                self._exchange_accounts_from_ows_account,
            ),
            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,
            ),
        )

        # Merge results
        # Ensure order of tenant types matches the order in asyncio.gather above
        self._gathered_tenant_exchanges[TenantType.TENANT_TYPE_ACCOUNT] = (
            tenant_exchanges[0]
        )
        self._gathered_tenant_exchanges[TenantType.TENANT_TYPE_SUBACCOUNT] = (
            tenant_exchanges[1]
        )
        self._gathered_tenant_exchanges[TenantType.TENANT_TYPE_COMPANY_BRAND] = (
            tenant_exchanges[2]
        )
        return self._gathered_tenant_exchanges