import json
import os
from typing import Optional, Dict

import click
from tabulate import tabulate

from montecarlodata.common.data import ConnectionOptions, ValidationResult, OnboardingConfiguration
from montecarlodata.common.user import UserService
from montecarlodata.config import Config
from montecarlodata.errors import complain_and_abort, prompt_connection
from montecarlodata.integrations.onboarding.fields import S3_CERT_MECHANISM, GQL_TO_FRIENDLY_CONNECTION_MAP, \
    DATA_LAKE_WAREHOUSE_TYPE, EXPECTED_ADD_CONNECTION_RESPONSE_FIELD, CONFIRM_CONNECTION_VERBIAGE, \
    VALIDATIONS_FAILED_VERBIAGE, \
    SKIP_ADD_CONNECTION_VERBIAGE, CONNECTION_TEST_SUCCESS_VERBIAGE, CONNECTION_TEST_FAILED_VERBIAGE, \
    ADD_CONNECTION_FAILED_VERBIAGE, ADD_CONNECTION_SUCCESS_VERBIAGE, MAIN_CONNECTION_TYPES
from montecarlodata.queries.onboarding import ADD_CONNECTION_MUTATION
from montecarlodata.utils import GqlWrapper, AwsClientWrapper


class BaseOnboardingService:
    _VALIDATION_STATUS_HEADERS = ['Validation', 'Result', 'Details']
    _VALIDATION_PASSED = 'Passed'
    _VALIDATION_FAILED = 'Failed'

    def __init__(self, config: Config, user_service: Optional[UserService] = None,
                 request_wrapper: Optional[GqlWrapper] = None,
                 aws_wrapper: Optional[AwsClientWrapper] = None):

        self._abort_on_error = True  # Aborts methods with deco on unhandled error
        self._dc_outputs = None

        self._user_service = user_service or UserService(config=config)
        self._request_wrapper = request_wrapper or GqlWrapper(config)
        self._aws_wrapper = aws_wrapper or AwsClientWrapper(profile_name=config.aws_profile,
                                                            region_name=config.aws_region)

    def onboard(self, **kwargs):
        """
        Convenience wrapper to validate and add a connection.
        """
        onboarding_config = OnboardingConfiguration(**kwargs)  # not passed as arg to avoid having to change all clients
        onboarding_config.connection_options = self._build_connection_options(**onboarding_config.connection_options)

        onboarding_config.warehouse_name = onboarding_config.connection_options.monolith_base_payload.pop(
            'warehouseName',
            onboarding_config.connection_type
        )
        # Prefer explicitly set types.
        if not onboarding_config.warehouse_type:
            onboarding_config.warehouse_type = onboarding_config.connection_options.monolith_base_payload.pop(
                'warehouseType',
                DATA_LAKE_WAREHOUSE_TYPE
            )

        result = self._validate_connection(
            query=onboarding_config.validation_query,
            response_field=onboarding_config.validation_response,
            **onboarding_config.connection_options.monolith_base_payload
        )

        if not onboarding_config.connection_options.validate_only:
            prompt = VALIDATIONS_FAILED_VERBIAGE if result.has_warnings else CONFIRM_CONNECTION_VERBIAGE
            prompt_connection(message=prompt, skip_prompt=onboarding_config.connection_options.auto_yes)
            self._add_connection(temp_path=result.credentials_key, onboarding_config=onboarding_config)
        else:
            click.echo(SKIP_ADD_CONNECTION_VERBIAGE)

    def handle_cert(self, cert_prefix: str, options: Dict) -> None:
        """
        Handles cert payload from either an s3 path or file. Uploading the latter
        Options is updated if successful.
        """
        if options.get('cert_file') is not None:
            bucket_name = self._get_dc_property(prop='PrivateS3BucketArn').split(':')[5]  # get name from arn
            object_name = os.path.join(cert_prefix, os.path.basename(options['cert_file']))
            self._aws_wrapper.upload_file(bucket_name=bucket_name, object_name=object_name,
                                          file_path=options['cert_file'])

            click.echo(f"Uploaded '{options['cert_file']}' to s3://{bucket_name}/{object_name}")
            options['cert_s3'] = object_name

        ssl_options = {}
        if options.get('cert_s3') is not None:
            # reformat to generic options and specify a mechanism
            ssl_options = {'mechanism': S3_CERT_MECHANISM, 'cert': options.pop('cert_s3')}
        options.pop('cert_file', None)

        if options.get('skip_cert_verification') is not None:
            ssl_options['skip_verification'] = options.pop('skip_cert_verification')

        if ssl_options:
            options['ssl_options'] = ssl_options

    def _validate_connection(self, query: str, response_field: str, **kwargs) -> ValidationResult:
        """
        Validate connection before adding using the expected gql response_field (e.g. testPrestoCredentials)
        """
        response = self._request_wrapper.make_request_v2(
            query=query,
            operation=response_field,
            variables=kwargs
        )

        summary = self._build_validation_summary(response.data)
        temp_path = response.data.get('key')
        warnings = response.data.get('warnings')

        if temp_path is not None:
            click.echo(f'{CONNECTION_TEST_SUCCESS_VERBIAGE}{summary}')
            return ValidationResult(
                has_warnings=warnings is not None and len(warnings) > 0,
                credentials_key=temp_path
            )

        complain_and_abort(f'{CONNECTION_TEST_FAILED_VERBIAGE}{summary}')

    def _build_validation_summary(self, response: Dict, table_format: Optional[str] = 'fancy_grid') -> str:
        data = []
        if response.get('validations'):
            for passed in response.get('validations'):
                data.append([
                    passed.get('message'),
                    self._VALIDATION_PASSED,
                    None
                ])

        if response.get('warnings'):
            for failed in response.get('warnings', []):
                data.append([
                    failed.get('message'),
                    self._VALIDATION_FAILED,
                    failed.get('data', {}).get('error')
                ])

        if data:
            return '\n' + tabulate(data, headers=self._VALIDATION_STATUS_HEADERS, tablefmt=table_format)

        # do not provide a summary if no detailed response is available
        return ''

    def _add_connection(self, temp_path: str, onboarding_config: OnboardingConfiguration) -> bool:
        """
        Add connection and setup any associated jobs
        """
        connection_request = {'key': temp_path, 'connectionType': onboarding_config.connection_type}
        connection_request.update(self._disambiguate_warehouses(onboarding_config))

        # Set optional properties
        if onboarding_config.warehouse_name:
            connection_request['name'] = onboarding_config.warehouse_name
        if onboarding_config.job_types:
            connection_request['jobTypes'] = onboarding_config.job_types
        if onboarding_config.job_limits:
            connection_request['jobLimits'] = json.dumps(onboarding_config.job_limits)
        if onboarding_config.connection_options and onboarding_config.connection_options.dc_id:
            connection_request['dcId'] = onboarding_config.connection_options.dc_id

        response = self._request_wrapper.make_request_v2(
            query=onboarding_config.connection_query or ADD_CONNECTION_MUTATION,
            operation=onboarding_config.connection_response or EXPECTED_ADD_CONNECTION_RESPONSE_FIELD,
            variables=connection_request,
        )

        connection_id = response.data.get('connection', {}).get('uuid')
        if connection_id is not None:
            click.echo(f"{ADD_CONNECTION_SUCCESS_VERBIAGE}"
                       f"{GQL_TO_FRIENDLY_CONNECTION_MAP.get(onboarding_config.connection_type, onboarding_config.connection_type.capitalize())}.")
            return True
        complain_and_abort(ADD_CONNECTION_FAILED_VERBIAGE)

    def _disambiguate_warehouses(self, onboarding_config: OnboardingConfiguration) -> Dict:
        """
        Determine type of connection request to build using the following criteria -
            1) If it is a main connection type or a traditional warehouse (i.e. not a lake), create a new warehouse.
            2) If a name has been passed, create the connection in that warehouse.
            3) If no name has been passed but there is only one warehouse, use it.
            4) Otherwise fail.
        """
        warehouse_type = onboarding_config.warehouse_type or DATA_LAKE_WAREHOUSE_TYPE  # default to data-lake
        connection_type = onboarding_config.connection_type
        create_warehouse = onboarding_config.create_warehouse

        # Main connections require the creation of a new warehouse.
        if (connection_type in MAIN_CONNECTION_TYPES or warehouse_type != DATA_LAKE_WAREHOUSE_TYPE) and create_warehouse:
            return {'createWarehouseType': warehouse_type}

        name = onboarding_config.warehouse_name

        dw_id = self._get_warehouse_id_from_name(name)
        return {'dwId': dw_id}

    def _get_warehouse_id_from_name(self, name: Optional[str]) -> str:
        warehouses = self._user_service.warehouses

        # If it is not a new warehouse, find the warehouse with the given name to create the new connection belonging
        # to it. If the name is not found, it is an error.
        if name:
            warehouses = [warehouse for warehouse in warehouses if warehouse.get('name') == name]
            if warehouses:
                return warehouses[0]['uuid']
            complain_and_abort('Could not find warehouse with given name')
        # Name is optional if there is only one warehouse.
        elif len(warehouses) == 1:
            return warehouses[0]['uuid']
        # If more than one warehouse exists, name is required.
        else:
            complain_and_abort('Name is required to disambiguate lake warehouses')

    def _get_dc_property(self, prop: str) -> Optional[str]:
        """
        Retrieve property from DC stack outputs
        """
        self._dc_outputs = self._dc_outputs or self._aws_wrapper.get_stack_outputs(
            self._user_service.active_collector['stackArn'])  # cache lookup
        for output in self._dc_outputs:
            if output['OutputKey'] == prop:
                return output['OutputValue']

    @staticmethod
    def _build_connection_options(**kwargs) -> ConnectionOptions:
        """
        Create connection options from arguments
        """
        connection_options = ConnectionOptions(**kwargs)

        if connection_options.monolith_connection_payload:
            connection_options.monolith_base_payload['connectionOptions'] = connection_options.monolith_connection_payload
        return connection_options
