content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  from typing import (
      Iterator, List, Optional, Union,
  )
  
  from amundsen_common.utils.atlas import AtlasCommonParams, AtlasCommonTypes
  from amundsen_rds.models import RDSModel
  from amundsen_rds.models.badge import Badge as RDSBadge
  from amundsen_rds.models.column import ColumnBadge as RDSColumnBadge
  from amundsen_rds.models.dashboard import DashboardBadge as RDSDashboardBadge
  from amundsen_rds.models.table import TableBadge as RDSTableBadge
  
  from databuilder.models.atlas_entity import AtlasEntity
  from databuilder.models.atlas_relationship import AtlasRelationship
  from databuilder.models.atlas_serializable import AtlasSerializable
  from databuilder.models.graph_node import GraphNode
  from databuilder.models.graph_relationship import GraphRelationship
  from databuilder.models.graph_serializable import GraphSerializable
  from databuilder.models.table_serializable import TableSerializable
  from databuilder.serializers.atlas_serializer import get_entity_attrs
  from databuilder.utils.atlas import AtlasRelationshipTypes, AtlasSerializedEntityOperation
  
  
  class Badge:
      def __init__(self, name: str, category: str):
          # Amundsen UI always formats badge display with first letter capitalized while other letters are lowercase.
          # Clicking table badges in UI always results in searching lower cases badges
          # https://github.com/amundsen-io/amundsen/blob/6ec9b398634264e52089bb9e1b7d76a6fb6a35a4/frontend/amundsen_application/static/js/components/BadgeList/index.tsx#L56
          # If badges stored in neo4j are not lowercase, they won't be searchable in UI.
          self.name = name.lower()
          self.category = category.lower()
  
      def __repr__(self) -> str:
          return f'Badge({self.name!r}, {self.category!r})'
  
      def __eq__(self, other: object) -> bool:
          if not isinstance(other, Badge):
              return NotImplemented
          return self.name == other.name and self.category == other.category
  
  
  class BadgeMetadata(GraphSerializable, TableSerializable, AtlasSerializable):
      """
      Badge model.
      """
      BADGE_NODE_LABEL = 'Badge'
      BADGE_KEY_FORMAT = '{badge}'
      BADGE_CATEGORY = 'category'
  
      # Relation between entity and badge
      BADGE_RELATION_TYPE = 'HAS_BADGE'
      INVERSE_BADGE_RELATION_TYPE = 'BADGE_FOR'
  
      LABELS_PERMITTED_TO_HAVE_BADGE = ['Table', 'Dashboard', 'Column', 'Feature', 'Type_Metadata']
  
      def __init__(self,
                   start_label: str,
                   start_key: str,
                   badges: List[Badge],
                   ):
          if start_label not in BadgeMetadata.LABELS_PERMITTED_TO_HAVE_BADGE:
              raise Exception(f'badges for {start_label} are not supported')
          self.start_label = start_label
          self.start_key = start_key
          self.badges = badges
  
          self._node_iter = self._create_node_iterator()
          self._relation_iter = self._create_relation_iterator()
          self._record_iter = self._create_record_iterator()
          self._atlas_entity_iterator = self._create_next_atlas_entity()
          self._atlas_relation_iterator = self._create_atlas_relation_iterator()
  
      def __repr__(self) -> str:
          return f'BadgeMetadata({self.start_label!r}, {self.start_key!r}, {self.badges!r})'
  
      def create_next_node(self) -> Optional[GraphNode]:
          try:
              return next(self._node_iter)
          except StopIteration:
              return None
  
      def create_next_relation(self) -> Optional[GraphRelationship]:
          try:
              return next(self._relation_iter)
          except StopIteration:
              return None
  
      def create_next_record(self) -> Union[RDSModel, None]:
          try:
              return next(self._record_iter)
          except StopIteration:
              return None
  
      @staticmethod
      def get_badge_key(name: str) -> str:
          if not name:
              return ''
          return BadgeMetadata.BADGE_KEY_FORMAT.format(badge=name)
  
      def get_badge_nodes(self) -> List[GraphNode]:
          nodes = []
          for badge in self.badges:
              if badge:
                  node = GraphNode(
                      key=self.get_badge_key(badge.name),
                      label=self.BADGE_NODE_LABEL,
                      attributes={
                          self.BADGE_CATEGORY: badge.category
                      }
                  )
                  nodes.append(node)
          return nodes
  
      def get_badge_relations(self) -> List[GraphRelationship]:
          relations = []
          for badge in self.badges:
              relation = GraphRelationship(
                  start_label=self.start_label,
                  end_label=self.BADGE_NODE_LABEL,
                  start_key=self.start_key,
                  end_key=self.get_badge_key(badge.name),
                  type=self.BADGE_RELATION_TYPE,
                  reverse_type=self.INVERSE_BADGE_RELATION_TYPE,
                  attributes={}
              )
              relations.append(relation)
          return relations
  
      def get_badge_records(self) -> List[RDSModel]:
          records = []
          for badge in self.badges:
              if badge:
                  record = RDSBadge(
                      rk=self.get_badge_key(badge.name),
                      category=badge.category
                  )
                  records.append(record)
  
          return records
  
      def _create_node_iterator(self) -> Iterator[GraphNode]:
          """
          Create badge nodes
          :return:
          """
          nodes = self.get_badge_nodes()
          for node in nodes:
              yield node
  
      def _create_relation_iterator(self) -> Iterator[GraphRelationship]:
          relations = self.get_badge_relations()
          for relation in relations:
              yield relation
  
      def _create_record_iterator(self) -> Iterator[RDSModel]:
          records = self.get_badge_records()
          for record in records:
              yield record
  
              if self.start_label == 'Table':
                  table_badge_record = RDSTableBadge(table_rk=self.start_key, badge_rk=record.rk)
                  yield table_badge_record
              elif self.start_label == 'Column':
                  column_badge_record = RDSColumnBadge(column_rk=self.start_key, badge_rk=record.rk)
                  yield column_badge_record
              elif self.start_label == 'Dashboard':
                  dashboard_badge_record = RDSDashboardBadge(dashboard_rk=self.start_key, badge_rk=record.rk)
                  yield dashboard_badge_record
  
      def _create_atlas_classification_entity(self, badge: Badge) -> AtlasEntity:
          attrs_mapping = [
              (AtlasCommonParams.qualified_name, badge.name),
              ('category', badge.category),
              ('name', badge.name)
          ]
  
          entity_attrs = get_entity_attrs(attrs_mapping)
  
          entity = AtlasEntity(
              typeName=AtlasCommonTypes.badge,
              operation=AtlasSerializedEntityOperation.CREATE,
              attributes=entity_attrs,
              relationships=None
          )
  
          return entity
  
      def _create_atlas_classification_relation(self, badge: Badge) -> AtlasRelationship:
          table_relationship = AtlasRelationship(
              relationshipType=AtlasRelationshipTypes.badge,
              entityType1=AtlasCommonTypes.data_set,
              entityQualifiedName1=self.start_key,
              entityType2=AtlasRelationshipTypes.badge,
              entityQualifiedName2=badge.name,
              attributes={}
          )
  
          return table_relationship
  
      def _create_atlas_relation_iterator(self) -> Iterator[AtlasRelationship]:
          for badge in self.badges:
              yield self._create_atlas_classification_relation(badge)
  
      def create_next_atlas_relation(self) -> Union[AtlasRelationship, None]:
          try:
              return next(self._atlas_relation_iterator)
          except StopIteration:
              return None
  
      def _create_next_atlas_entity(self) -> Iterator[AtlasEntity]:
          for badge in self.badges:
              yield self._create_atlas_classification_entity(badge)
  
      def create_next_atlas_entity(self) -> Union[AtlasEntity, None]:
          try:
              return next(self._atlas_entity_iterator)
          except StopIteration:
              return None
  