# Copyright (c) 2023 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
import os
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Generic, Optional, TypeVar

from aws_lambda_powertools.utilities.typing import LambdaContext
from toolz import get_in, keyfilter
from voluptuous import Schema

from czc import aws_environment, event_bridge, feature_data, s3, type_util
from czc import dynamodb as dyn
from czc.aws_environment import AccountInfo
from czc.dynamodb import DynamoDbTable
from czc.feature_data import FeatureDataConnection
from czc.helpers import coalesce
from czc.logging import get_logger

logger = get_logger()

T = TypeVar('T')
F = TypeVar('F', bound=Callable[..., Any])
ORGANIZATION_PARAMS = ['cz_organization_id', 'organization_id', 'organizationId']


@dataclass(frozen=True)
class HandlerContext:
    _organization_id: Optional[str]

    lambda_context: LambdaContext
    namespace: str
    account_info: AccountInfo
    feature_conn: FeatureDataConnection
    bucket: str
    dyn_table: DynamoDbTable
    event_bus: str

    @property
    def has_organization_id(self) -> bool:
        return bool(self._organization_id)

    @property
    def organization_id(self) -> str:
        return type_util.require(self._organization_id, 'No organization ID set in handler context')


@dataclass(frozen=True)
class HandlerContextWithEvent(HandlerContext, Generic[T]):
    lambda_event: T

    @classmethod
    def create(cls, event: T, context: LambdaContext, organization_id: Optional[str] = None) -> 'HandlerContextWithEvent[T]':
        return HandlerContextWithEvent(
            _organization_id=organization_id,
            lambda_event=event,
            lambda_context=context,
            namespace=os.environ['NAMESPACE'],
            account_info=aws_environment.resolve_account_info(),
            feature_conn=feature_data.connect(),
            bucket=s3.get_bucket_name(),
            dyn_table=dyn.table(dyn.get_table_name()),
            event_bus=event_bridge.get_event_bus_name(),
        )


def with_handler_context(event_params: Optional[Schema] = None):
    def _handler(handler):
        def handler_wrapper(event: T, context: LambdaContext, **kwargs):
            event_kwargs = event_params(keyfilter(lambda k: not k.startswith('__'), event)) if event_params else {}
            all_kwargs = {**event_kwargs, **kwargs}
            authorizer_params = get_in(['requestContext', 'authorizer'], event, default={})
            authorizer_org_id = coalesce(*(authorizer_params.get(x) for x in ORGANIZATION_PARAMS))
            organization_id = authorizer_org_id or coalesce(*(all_kwargs.get(x) for x in ORGANIZATION_PARAMS))

            logger.append_keys(cz_organization_id=organization_id)
            hc = HandlerContextWithEvent.create(event, context, organization_id)

            return handler(handler_context=hc, **all_kwargs)

        return handler_wrapper

    return _handler
