#!/usr/bin/python
#
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging

from google.datacatalog_connectors.apache_atlas import scrape


class MetadataScraper:

    def __init__(self, connection_args):
        self._apache_atlas_facade = scrape.apache_atlas_facade.\
            ApacheAtlasFacade(connection_args)
        self.__metadata_enricher = scrape.metadata_enricher.\
            MetadataEnricher(self._apache_atlas_facade)

    def get_metadata(self, **kwargs):
        self._log_scrape_start('Scraping all Metadata...')
        classifications_dict = {}
        entity_types_dict = {}
        enum_types_dict = {}

        self._log_scrape_start('Scraping admin metrics...')
        admin_metrics = self._apache_atlas_facade.get_admin_metrics()
        logging.info(admin_metrics)
        self._log_single_object_scrape_result(admin_metrics)

        self._log_scrape_start('Scraping typedefs...')
        for typedef in self._apache_atlas_facade.get_typedefs():
            self._scrape_classification_types(classifications_dict, typedef)

            self._scrape_enum_types(enum_types_dict, typedef)

            self._scrape_entity_types(entity_types_dict, typedef)

        self.__metadata_enricher.enrich_entity_relationships(entity_types_dict)

        return {
            'classifications': classifications_dict,
            'enum_types': enum_types_dict,
            'entity_types': entity_types_dict
        }, None

    def _scrape_entity_types(self, entity_types_dict, typedef):
        self._log_scrape_start('Scraping EntityTypes...')
        for entity_type in typedef.entityDefs:
            entity_type_name = entity_type.name

            entity_type_dict = {
                'name': entity_type_name,
                'data': entity_type._data,
                'superTypes': entity_type.superTypes,
                'entities': {}
            }

            entities = self.__scrape_entity_type(entity_type)
            entity_type_dict['entities'] = entities

            entity_types_dict[entity_type_name] = entity_type_dict

    def _scrape_classification_types(self, classifications_dict, typedef):
        self._log_scrape_start('Scraping Classifications/Templates...')
        for classification_type in typedef.classificationDefs:
            classification_data = classification_type._data
            logging.info('Classification: %s', classification_type.name)
            logging.debug(classification_data)
            classifications_dict[classification_type.name] = {
                'name': classification_type.name,
                'guid': classification_type.guid,
                'data': classification_data
            }

    def _scrape_enum_types(self, enum_types_dict, typedef):
        self._log_scrape_start('Scraping Enum types...')
        for enum_type in typedef.enumDefs:
            enum_data = enum_type._data
            logging.info('Enum type: %s', enum_type.name)
            logging.debug(enum_data)
            enum_types_dict[enum_type.name] = {
                'name': enum_type.name,
                'guid': enum_type.guid,
                'data': enum_data
            }

    def __scrape_entity_type(self, entity_type):
        searched_entries = {}
        entity_type_name = entity_type.name

        logging.info('=> Entity Type: %s', entity_type_name)
        logging.debug(entity_type._data)

        search_results = self._apache_atlas_facade.\
            search_entities_from_entity_type(entity_type_name)

        guids = []
        for entity in search_results:
            # Collecting guids and storing entity to enricher data later on.
            guid = entity.guid
            guids.append(guid)
            searched_entries[guid] = {'guid': guid, 'data': entity._data}

        fetched_entities_dict = {}
        if guids:
            fetched_entities_dict = self._apache_atlas_facade.fetch_entities(
                guids)
            self.__metadata_enricher.enrich_entity_classifications(
                fetched_entities_dict, searched_entries)

        logging.info('Entity Type: %s scrapped!', entity_type_name)
        logging.info('')
        return fetched_entities_dict

    @classmethod
    def _log_scrape_start(cls, message, *args):
        logging.info('')
        logging.info(message, *args)
        logging.info('-------------------------------------------------')

    @classmethod
    def _log_single_object_scrape_result(cls, the_object):
        logging.info('Found!' if the_object else 'NOT found!')

[(25, 'google.datacatalog_connectors.apache_atlas.scrape.apache_atlas_facade.ApacheAtlasFacade', 'scrape.apache_atlas_facade.ApacheAtlasFacade', (['connection_args'], {}), False, 'from google.datacatalog_connectors.apache_atlas import scrape\n'), (27, 'google.datacatalog_connectors.apache_atlas.scrape.metadata_enricher.MetadataEnricher', 'scrape.metadata_enricher.MetadataEnricher', (['self._apache_atlas_facade'], {}), False, 'from google.datacatalog_connectors.apache_atlas import scrape\n'), (38, 'logging.info', 'logging.info', (['admin_metrics'], {}), False, 'import logging\n'), (102, 'logging.info', 'logging.info', (['"""=> Entity Type: %s"""', 'entity_type_name'], {}), False, 'import logging\n'), (103, 'logging.debug', 'logging.debug', (['entity_type._data'], {}), False, 'import logging\n'), (122, 'logging.info', 'logging.info', (['"""Entity Type: %s scrapped!"""', 'entity_type_name'], {}), False, 'import logging\n'), (123, 'logging.info', 'logging.info', (['""""""'], {}), False, 'import logging\n'), (128, 'logging.info', 'logging.info', (['""""""'], {}), False, 'import logging\n'), (129, 'logging.info', 'logging.info', (['message', '*args'], {}), False, 'import logging\n'), (130, 'logging.info', 'logging.info', (['"""-------------------------------------------------"""'], {}), False, 'import logging\n'), (134, 'logging.info', 'logging.info', (["('Found!' if the_object else 'NOT found!')"], {}), False, 'import logging\n'), (78, 'logging.info', 'logging.info', (['"""Classification: %s"""', 'classification_type.name'], {}), False, 'import logging\n'), (79, 'logging.debug', 'logging.debug', (['classification_data'], {}), False, 'import logging\n'), (90, 'logging.info', 'logging.info', (['"""Enum type: %s"""', 'enum_type.name'], {}), False, 'import logging\n'), (91, 'logging.debug', 'logging.debug', (['enum_data'], {}), False, 'import logging\n')]