from montecarlodata.common.common import read_as_base64
from montecarlodata.config import Config
from montecarlodata.errors import manage_errors
from montecarlodata.integrations.onboarding.base import BaseOnboardingService
from montecarlodata.integrations.onboarding.fields import EXPECTED_GENERIC_DB_GQL_RESPONSE_FIELD, \
    REDSHIFT_CONNECTION_TYPE, REDSHIFT_WAREHOUSE_TYPE, SNOWFLAKE_CONNECTION_TYPE, SNOWFLAKE_WAREHOUSE_TYPE, \
    EXPECTED_SNOWFLAKE_GQL_RESPONSE_FIELD, BQ_WAREHOUSE_TYPE, EXPECTED_BQ_GQL_RESPONSE_FIELD, BQ_CONNECTION_TYPE
from montecarlodata.queries.onboarding import TEST_DATABASE_CRED_MUTATION, TEST_SNOWFLAKE_CRED_MUTATION, \
    TEST_BQ_CRED_MUTATION


class WarehouseOnboardingService(BaseOnboardingService):

    def __init__(self, config: Config, **kwargs):
        super().__init__(config, **kwargs)

    @manage_errors
    def onboard_redshift(self, **kwargs) -> None:
        """
        Onboard a redshift connection by validating and adding a connection.
        """
        kwargs['connectionType'] = REDSHIFT_CONNECTION_TYPE
        kwargs['warehouseType'] = REDSHIFT_WAREHOUSE_TYPE
        self.onboard(validation_query=TEST_DATABASE_CRED_MUTATION,
                     validation_response=EXPECTED_GENERIC_DB_GQL_RESPONSE_FIELD,
                     connection_type=REDSHIFT_CONNECTION_TYPE, **kwargs)

    @manage_errors
    def onboard_snowflake(self, **kwargs) -> None:
        """
        Onboard a snowflake connection by validating and adding a connection.
        """
        kwargs['warehouseType'] = SNOWFLAKE_WAREHOUSE_TYPE
        self.onboard(validation_query=TEST_SNOWFLAKE_CRED_MUTATION,
                     validation_response=EXPECTED_SNOWFLAKE_GQL_RESPONSE_FIELD,
                     connection_type=SNOWFLAKE_CONNECTION_TYPE, **kwargs)

    @manage_errors
    def onboard_bq(self, **kwargs) -> None:
        """
        Onboard a BigQuery connection by validating and adding a connection.

        Reads and encodes service file as base64.
        """
        kwargs['warehouseType'] = BQ_WAREHOUSE_TYPE
        kwargs['serviceJson'] = read_as_base64(kwargs.pop('ServiceFile')).decode('utf-8')
        self.onboard(validation_query=TEST_BQ_CRED_MUTATION,
                     validation_response=EXPECTED_BQ_GQL_RESPONSE_FIELD,
                     connection_type=BQ_CONNECTION_TYPE, **kwargs)
