# Copyright (c) 2021 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 typing import Mapping, Optional

from czc.api.v2.constants import CLOUDZERO_TEST_KEY
from czc.constants import CLOUDZERO_TEST_KEY_DO_NOT_FORWARD_SUFFIX, CLOUDZERO_TEST_KEY_EVENT_PARAMETER
from czc.logging import get_logger

logger = get_logger()


def _set_test_key(test_key_value):
    os.environ[CLOUDZERO_TEST_KEY] = test_key_value
    logger.info(f'Test Key is present! Using {test_key_value}')


def _unset_test_key():
    if CLOUDZERO_TEST_KEY in os.environ:
        logger.info(f'Removing {CLOUDZERO_TEST_KEY} from environment.')
        os.environ.pop(CLOUDZERO_TEST_KEY)


class TestKeyEnvironment:
    def __init__(self, cloudzero_test_key_value: Optional[str]):
        self._cloudzero_test_key_value = cloudzero_test_key_value

    def __enter__(self):
        if self._cloudzero_test_key_value:
            _set_test_key(self._cloudzero_test_key_value)
        else:
            logger.info(f'Test Key is not present! Resetting {CLOUDZERO_TEST_KEY}.')
            _unset_test_key()

    def __exit__(self, exc_type, exc_val, exc_tb):
        _unset_test_key()

    @staticmethod
    def get():
        return os.environ.get(CLOUDZERO_TEST_KEY)

    @staticmethod
    def add_test_key_parameter(payload: Mapping, test_key_value: Optional[str] = None) -> Mapping:
        '''
        If a test key exists and it is able to be forwarded, then add it into the payload.

        Parameters:
            payload: Mapping - The data/parameters that are being used for the cross feature call or event
            test_key_value: Optional[str] - Normally this parameter should be omitted so the environment variable
                is used. However, there are some cases, due to the use of iterators, where the scope of the environment
                variable is lost, so the test key value can be passed manually in those cases.

        Returns:
            A new payload with the test key added or not depending on its presence and value.
        '''
        test_key_value = test_key_value or os.environ.get(CLOUDZERO_TEST_KEY)

        return {
            **payload,
            **({CLOUDZERO_TEST_KEY_EVENT_PARAMETER: test_key_value} if test_key_value and not test_key_value.endswith(CLOUDZERO_TEST_KEY_DO_NOT_FORWARD_SUFFIX) else {})
        }
