#  Copyright 2021 Collate
#  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.
"""Metadata source module"""

import logging
from typing import Iterable

from metadata.generated.schema.entity.data.dashboard import Dashboard
from metadata.generated.schema.entity.data.glossary import Glossary
from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
from metadata.generated.schema.entity.data.pipeline import Pipeline
from metadata.generated.schema.entity.data.table import Table
from metadata.generated.schema.entity.data.topic import Topic
from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
    OpenMetadataConnection,
)
from metadata.generated.schema.entity.services.pipelineService import PipelineService
from metadata.generated.schema.entity.teams.team import Team
from metadata.generated.schema.entity.teams.user import User
from metadata.generated.schema.metadataIngestion.workflow import (
    Source as WorkflowSource,
)
from metadata.ingestion.api.common import Entity
from metadata.ingestion.api.source import InvalidSourceException, SourceStatus
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.ingestion.source.metadata import MetadataSource

logger = logging.getLogger(__name__)


class PolicyWrapper:
    data_dict: dict

    def __init__(self, policy_dict) -> None:
        self.data_dict = policy_dict


class TableWrapper:
    data_dict: dict

    def __init__(self, table_dict) -> None:
        self.data_dict = table_dict


class TagWrapper:
    data_dict: dict

    def __init__(self, tag_dict) -> None:
        self.data_dict = tag_dict


class MessagingServiceWrapper:
    data_dict: dict

    def __init__(self, messaging_service_dict) -> None:
        self.data_dict = messaging_service_dict


class DatabaseServiceWrapper:
    data_dict: dict

    def __init__(self, database_service_dict) -> None:
        self.data_dict = database_service_dict


class MigrateSource(MetadataSource):
    """OpenmetadataSource class

    Args:
        config:
        metadata_config:

    Attributes:
        config:
        report:
        metadata_config:
        status:
        wrote_something:
        metadata:
        tables:
        topics:
    """

    config: WorkflowSource
    report: SourceStatus

    def __init__(
        self,
        config: WorkflowSource,
        metadata_config: OpenMetadataConnection,
    ):
        super().__init__(config, metadata_config)
        self.metadata = OpenMetadata(
            OpenMetadataConnection.parse_obj(self.service_connection)
        )

    @classmethod
    def create(cls, config_dict, metadata_config: OpenMetadataConnection):
        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
        connection: OpenMetadataConnection = config.serviceConnection.__root__.config
        if not isinstance(connection, OpenMetadataConnection):
            raise InvalidSourceException(
                f"Expected OpenMetadataConnection, but got {connection}"
            )
        return cls(config, metadata_config)

    def next_record(self) -> Iterable[Entity]:
        if self.service_connection.includeTables:
            yield from self.fetch_tables(
                fields=[
                    "columns",
                    "tableConstraints",
                    "usageSummary",
                    "owner",
                    "tags",
                    "followers",
                ],
            )
        if self.service_connection.includeTopics:
            yield from self.fetch_entities(
                entity_class=Topic,
                fields=["owner", "tags", "followers"],
            )
        if self.service_connection.includeDashboards:
            yield from self.fetch_entities(
                entity_class=Dashboard,
                fields=[
                    "owner",
                    "tags",
                    "followers",
                    "charts",
                    "usageSummary",
                ],
            )

        if self.service_connection.includePipelines:
            yield from self.fetch_entities(
                entity_class=Pipeline,
                fields=["owner", "tags", "followers", "tasks"],
            )
        if self.service_connection.includeUsers:
            yield from self.fetch_entities(
                entity_class=User,
                fields=["teams", "roles"],
            )

        if self.service_connection.includeTeams:
            yield from self.fetch_entities(
                entity_class=Team,
                fields=["users", "owns"],
            )

        if self.service_connection.includeGlossaryTerms:
            yield from self.fetch_entities(
                entity_class=GlossaryTerm,
                fields=[],
            )
            yield from self.fetch_entities(
                entity_class=Glossary,
                fields=["owner", "tags", "reviewers", "usageCount"],
            )

        if self.service_connection.includePolicy:
            yield from self.fetch_policy()

        if self.service_connection.includeTags:
            yield from self.fetch_tags()

        if self.service_connection.includeMessagingServices:
            yield from self.fetch_messaging_services()

        if self.service_connection.includeDatabaseServices:
            yield from self.fetch_database_services()

        if self.service_connection.includePipelineServices:
            yield from self.fetch_entities(
                entity_class=PipelineService,
                fields=["owner"],
            )

    def fetch_policy(self):
        after = None
        while True:
            after_param = f"&after={after}" if after else ""
            policy_entities = self.metadata.client.get(
                f"/policies?&limit={self.service_connection.limitRecords}{after_param}"
            )
            for policy in policy_entities.get("data"):
                yield PolicyWrapper(policy)
            after = policy_entities.get("paging").get("after")
            if not after:
                break

    def fetch_tables(self, fields):
        after = None
        while True:
            after_param = f"&after={after}" if after else ""
            table_entities = self.metadata.client.get(
                f"/tables?fields={','.join(fields)}&limit={self.service_connection.limitRecords}{after_param}"
            )
            for table in table_entities.get("data"):
                yield TableWrapper(table)
            after = table_entities.get("paging").get("after")
            if not after:
                break

    def fetch_tags(self):
        tag_entities = self.metadata.client.get("/tags")
        for tag in tag_entities.get("data"):
            tag_detailed_entity = self.metadata.client.get(f"/tags/{tag.get('name')}")
            yield TagWrapper(tag_detailed_entity)

    def fetch_messaging_services(self):
        after = None
        while True:
            after_param = f"&after={after}" if after else ""
            service_entities = self.metadata.client.get(
                f"/services/messagingServices?fields=owner&limit={self.service_connection.limitRecords}{after_param}"
            )
            for service in service_entities.get("data"):
                yield MessagingServiceWrapper(service)
            after = service_entities.get("paging").get("after")
            if not after:
                break

    def fetch_database_services(self):
        after = None
        while True:
            after_param = f"&after={after}" if after else ""
            service_entities = self.metadata.client.get(
                f"/services/databaseServices?fields=owner&limit={self.service_connection.limitRecords}{after_param}"
            )
            for service in service_entities.get("data"):
                yield DatabaseServiceWrapper(service)
            after = service_entities.get("paging").get("after")
            if not after:
                break
