# Copyright Contributors to the Amundsen project.
# SPDX-License-Identifier: Apache-2.0


import abc
from functools import reduce
from typing import (
    Any, Dict, Iterator, List, Tuple,
)

import requests
from dateutil import parser
from pyhocon import ConfigFactory, ConfigTree

from databuilder.extractor.base_extractor import Extractor

type_fields_mapping = List[Tuple[str, str, Any, Any]]


class ApacheSupersetBaseExtractor(Extractor):
    """
    Base class to create extractors pulling any dashboard metadata from Apache Superset.
    """
    APACHE_SUPERSET_PROTOCOL = 'apache_superset_protocol'
    APACHE_SUPERSET_HOST = 'apache_superset_host'
    APACHE_SUPERSET_PORT = 'apache_superset_port'
    APACHE_SUPERSET_SECURITY_SETTINGS_DICT = 'apache_superset_security_settings_dict'
    APACHE_SUPERSET_PAGE_SIZE = 'apache_superset_page_size'
    APACHE_SUPERSET_EXTRACT_PUBLISHED_ONLY = 'apache_superset_extract_published_only'
    APACHE_SUPERSET_SECURITY_PROVIDER = 'apache_superset_security_provider'

    DASHBOARD_GROUP_NAME = 'dashboard_group_name'
    DASHBOARD_GROUP_ID = 'dashboard_group_id'
    DASHBOARD_GROUP_DESCRIPTION = 'dashboard_group_description'

    PRODUCT = 'product'
    CLUSTER = 'cluster'

    DRIVER_TO_DATABASE_MAPPING = 'driver_to_database_mapping'

    DEFAULT_DRIVER_TO_DATABASE_MAPPING = {
        'postgresql': 'postgres',
        'mysql+pymysql': 'mysql'
    }

    DATABASE_TO_CLUSTER_MAPPING = 'database_to_cluster_mapping'  # map superset dbs to preferred clusters

    DEFAULT_CONFIG = ConfigFactory.from_dict({
        APACHE_SUPERSET_PROTOCOL: 'http',
        APACHE_SUPERSET_HOST: 'localhost',
        APACHE_SUPERSET_PORT: '8088',
        APACHE_SUPERSET_PAGE_SIZE: 20,
        APACHE_SUPERSET_EXTRACT_PUBLISHED_ONLY: False,
        PRODUCT: 'superset',
        DASHBOARD_GROUP_DESCRIPTION: '',
        DRIVER_TO_DATABASE_MAPPING: DEFAULT_DRIVER_TO_DATABASE_MAPPING,
        DATABASE_TO_CLUSTER_MAPPING: {}
    })

    def init(self, conf: ConfigTree) -> None:
        self.conf = conf.with_fallback(ApacheSupersetBaseExtractor.DEFAULT_CONFIG)
        self._extract_iter = self._get_extract_iter()

        self.authenticate()

    def get_scope(self) -> str:
        return 'extractor.apache_superset'

    def extract(self) -> Any:
        try:
            result = next(self._extract_iter)

            return result
        except StopIteration:
            return None

    def authenticate(self) -> None:
        security_settings = dict(self.conf.get(ApacheSupersetBaseExtractor.APACHE_SUPERSET_SECURITY_SETTINGS_DICT))

        token = requests.post(self.build_full_url('api/v1/security/login'),
                              json=security_settings)

        self.token = token.json()['access_token']

    def build_full_url(self, endpoint: str) -> str:
        return f'{self.base_url}/{endpoint}'

    def execute_query(self, url: str, params: dict = {}) -> Dict:
        try:
            data = requests.get(url, params=params, headers={'Authorization': f'Bearer {self.token}'})

            if data.status_code == 401:
                self.authenticate()

                return self.execute_query(url, params)
            else:
                return data.json()
        except Exception:
            return {}

    @property
    def base_url(self) -> str:
        _protocol = self.conf.get(ApacheSupersetBaseExtractor.APACHE_SUPERSET_PROTOCOL)
        _host = self.conf.get(ApacheSupersetBaseExtractor.APACHE_SUPERSET_HOST)
        _port = self.conf.get(ApacheSupersetBaseExtractor.APACHE_SUPERSET_PORT)

        return f'{_protocol}://{_host}:{_port}'

    @property
    def page_size(self) -> int:
        return self.conf.get_int(ApacheSupersetBaseExtractor.APACHE_SUPERSET_PAGE_SIZE)

    @property
    def product(self) -> str:
        return self.conf.get(ApacheSupersetBaseExtractor.PRODUCT)

    @property
    def cluster(self) -> str:
        return self.conf.get(ApacheSupersetBaseExtractor.CLUSTER)

    @property
    def extract_published_only(self) -> bool:
        return self.conf.get(ApacheSupersetBaseExtractor.APACHE_SUPERSET_EXTRACT_PUBLISHED_ONLY)

    @property
    def common_params(self) -> Dict[str, str]:
        return dict(dashboard_group=self.conf.get(ApacheSupersetBaseExtractor.DASHBOARD_GROUP_NAME),
                    dashboard_group_id=self.conf.get(ApacheSupersetBaseExtractor.DASHBOARD_GROUP_ID),
                    dashboard_group_url=self.base_url,
                    dashboard_group_description=self.conf.get(ApacheSupersetBaseExtractor.DASHBOARD_GROUP_DESCRIPTION),
                    product=self.product,
                    cluster=self.cluster)

    @staticmethod
    def parse_date(string_date: str) -> int:
        try:
            date_parsed = parser.parse(string_date)

            # date returned by superset api does not contain timezone so to be timezone safe we need to assume it's utc
            if not date_parsed.tzname():
                return ApacheSupersetBaseExtractor.parse_date(f'{string_date}+0000')

            return int(date_parsed.timestamp())
        except Exception:
            return 0

    @staticmethod
    def get_nested_field(input_dict: Dict, field: str) -> Any:
        return reduce(lambda x, y: x.get(y, dict()), field.split('.'), input_dict)

    @staticmethod
    def map_fields(data: Dict, mappings: type_fields_mapping) -> Dict:
        result = dict()

        for mapping in mappings:
            target_field, source_field, transform, default_value = mapping
            value = ApacheSupersetBaseExtractor.get_nested_field(data, source_field)

            if transform:
                value = transform(value)

            result[target_field] = value or default_value

        return result

    @abc.abstractmethod
    def _get_extract_iter(self) -> Iterator[Any]:
        pass

    def _get_resource_ids(self, resource: str) -> List[str]:
        i = 0
        result = []

        while True:
            url = self.build_full_url(f'api/v1/{resource}?q=(page_size:{self.page_size},page:{i},order_direction:desc)')

            data = self.execute_query(url)

            ids = data.get('ids', [])

            if ids:
                result += ids
                i += 1
            else:
                break

        return result
