# Copyright (c) 2022 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
from functools import reduce
from typing import Dict, Any, List, Optional, Tuple, Generic, TypeVar, Type

from cachetools import TTLCache
from cachetools import cached
from toolz import merge_with, assoc_in

from czc import feature_data, type_util
from czc.helpers import coalesce
from czc.json import serializable
from czc.logging import get_logger

logger = get_logger()
get_config_cache = TTLCache(maxsize=128, ttl=2)

ORGANIZATION_DEFAULT = 'default'
ORGANIZATION_CONFIG_KEY = 'configuration'

T = TypeVar('T')


class OrganizationConfig(Generic[T]):
    """Singleton class for managing per-organization configuration settings.

    Create a model class with the desired configuration settings as attributes.
    This can be any kind of data class as long as it is "deserializable" (i.e. has a
    constructor that takes a dict of class's "serializable" form).  An easy way to do
    this is to use the `@deserializable()` decorator from czc.json.

    All the model class's attributes should have default values of either other data
    classes, empty dictionaries, or None.  This ensures that an empty configuration
    object can be constructed.  *Important*: Do not use other default values, or the models
    will not be merged correctly.

    Model data is stored in DynamoDB with a PK of `organization#<organization_id>` or
    `organization#default` for the default configuration.  The SK is `configuration` or
    a different key passed to the constructor.

    When the configuration is loaded, the default values will be merged with the
    organization-specific values.  The organization-specific values will override the
    defaults.  Dict and other data classes will be merged recursively.  Lists, sets, and
    other primitive types will be replaced.

    Example usage:

    class Visibility(Enum):
        show = 'show'
        hide = 'hide'

    class FeatureFlags(NamedTuple):
        turbo_mode: Optional[bool] = None
        low_cost_mode: Optional[bool] = None

    @deserializable()
    class OrgConfigModel(NamedTuple):
        feature_flags: FeatureFlags = FeatureFlags()        # initialize with empty FeatureFlags
        visibility_overrides: Dict[str, Visibility] = {}    # initialize with empty dict
        extra_strings: Optional[List[str]] = None           # must initialize with None, not [] or any other value
        param1: Optional[str] = None                        # must initialize with None, not '' or any other value
        num: Optional[int] = None                           # must initialize with None, not 0 or any other value

    # create singleton instance to manage configuration
    organization_config = OrganizationConfig(OrgConfigModel)

    # initialize default and org specific configuration
    organization_config.init(None)                                   # default
    organization_config.init('1a3266ab-05dd-479b-97bd-9865df60b003') # some organization ID

    # get configuration for an organization
    config = organization_config.get('1a3266ab-05dd-479b-97bd-9865df60b003')
    print(config.feature_flags.turbo_mode)
    """
    def __init__(self, model: Type[T], key: str = ORGANIZATION_CONFIG_KEY):
        self._model = model
        self._key = key
        self._feature_data = feature_data.for_type(model)

    @property
    def _conn(self):
        return feature_data.connect()

    def _coerce(self, data: dict) -> T:
        return type_util.coerce(self._model, data)

    def _merge(self, default: T, org: T):
        def merge_value(values: List[Any]):
            if all(isinstance(x, dict) for x in values):
                return merge_with(merge_value, values)

            return coalesce(*reversed(values))

        merged_data = merge_with(merge_value, serializable(default), serializable(org))
        return self._coerce(merged_data)

    def _get_config(self, organization_key: str) -> T:
        return self._feature_data.get(self._conn, organization_key, self._key) or self._model()

    @cached(cache=get_config_cache, key=lambda c, org_id: org_id)
    def get(self, organization_id: str) -> T:
        return self.get_uncached(organization_id)

    def get_uncached(self, organization_id: str) -> T:
        return self._merge(self._get_config(ORGANIZATION_DEFAULT),
                           self._get_config(organization_id))

    def update(self, organization_id: str, config_values: Dict[str, Any]):
        """Update the configuration for an organization.

        The config_values dict should be a flat dict of key/value pairs.  The keys should
        be dot-separated paths to the configuration values to update.  For example:
        {
            'feature_flags.turbo_mode': True,
            'visibility_overrides.cost_type': 'hide'
        }
        """
        def set_value(data: Dict[str, Any], update: Tuple[str, Any]):
            key, value = update
            return assoc_in(data, key.split('.'), value)

        def handler(data: Optional[T]):
            result = reduce(set_value, config_values.items(), serializable(data or self._model()))
            return self._coerce(result)

        self._feature_data.update(self._conn, organization_id, self._key, handler=handler)

    def init(self, organization_id: Optional[str], defaults: Optional[T] = None) -> None:
        """Initialize the configuration for an organization or the default configuration.

        If the organization_id is None, the default configuration will be initialized.
        Otherwise, the organization-specific configuration will be initialized.

        Default values can be passed in as a model instance.  These values will be merged
        with, but not override, the existing configuration.  This is useful for adding
        new configuration parameters.
        """
        def handler(data: Optional[T]):
            return self._merge(defaults or self._model(), data or self._model())

        self._feature_data.update(self._conn,
                                  organization_id or ORGANIZATION_DEFAULT,
                                  self._key,
                                  handler=handler)
