class Entity(BaseModel):
    """
    An entity (node) in a knowledge graph.

    Attributes:
        id (Optional[str]): the unique id for the entity
        label (Optional[str]): an optional label for the entity
        metadata (Optional[dict]): optional information about the entity
    """
    id: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4()))
    label: Optional[str] = "entity"
    metadata: Optional[dict] = Field(default_factory=dict)

class GraphDocument(BaseModel):
    """
    A graph document is a collection of entities and relationships that are part of a knowledge graph.

    Attributes:
        entities (List[Entity]): the entities in the graph document
        relationships (List[Relationship]): the relationships in the graph document
        source (Source): the source that contains the entities and relationships
    """
    entities: List[Entity]
    relationships: List[Relationship]
    source: Source

class LangChainWiseAgentGraphDB(WiseAgentGraphDB):
    """
    An abstract class that makes use of a LangChain graph database.
    """

    def __init__(self, embedding_model_name: Optional[str] = DEFAULT_EMBEDDING_MODEL_NAME):
        """
        Initialize a new instance of LangChainWiseAgentGraphDB.


        Args:
            embedding_model_name (Optional[str]): the optional name of the embedding model to use 
        """
        self._embedding_model_name = embedding_model_name
        self._embedding_function = HuggingFaceEmbeddings(model_name=self.embedding_model_name)

    @property
    def embedding_model_name(self):
        """Get the name of the embedding model."""
        return self._embedding_model_name

    def convert_to_lang_chain_node(self, entity: Entity) -> Node:
        return Node(id=entity.id, type=entity.label, properties=entity.metadata)

    def convert_to_lang_chain_relationship(self, relationship: Relationship) -> LangChainRelationship:
        return LangChainRelationship(source=self.convert_to_lang_chain_node(relationship.source),
                                     target=self.convert_to_lang_chain_node(relationship.target),
                                     type=relationship.label,
                                     properties=relationship.metadata)

    def convert_to_lang_chain_graph_document(self, graph_document: GraphDocument) -> LangChainGraphDocument:
        return LangChainGraphDocument(nodes=[self.convert_to_lang_chain_node(entity)
                                             for entity in graph_document.entities],
                                      relationships=[self.convert_to_lang_chain_relationship(relationship)
                                                     for relationship in graph_document.relationships],
                                      source=self.convert_to_lang_chain_document(graph_document.source))

    def convert_to_lang_chain_document(self, source: Source) -> LangChainDocument:
        return LangChainDocument(id=source.id, page_content=source.content, metadata=source.metadata)

    @abstractmethod
    def get_schema(self) -> str:
        ...

    @abstractmethod
    def refresh_schema(self):
        ...

    @abstractmethod
    def query(self, query: str, params: Optional[dict] = None):
        ...

    @abstractmethod
    def insert_entity(self, entity: Entity, source: Source):
        ...

    @abstractmethod
    def insert_relationship(self, relationship: Relationship, source: Source):
        ...

    @abstractmethod
    def insert_graph_documents(self, graph_documents: List[GraphDocument]):
        ...

    @abstractmethod
    def create_vector_db_from_graph_db(self, retrieval_query: str = ""):
        ...

class Neo4jLangChainWiseAgentGraphDB(LangChainWiseAgentGraphDB):
    """
    A LangChainWiseAgentGraphDB implementation that makes use of a LangChain Neo4j graph database
    and a corresponding Neo4j vector database.
    """
    yaml_tag = u'!Neo4jLangChainWiseAgentGraphDB'

    def __init__(self, properties: List[str], collection_name: str, url: Optional[str] = None,
                 refresh_graph_schema: Optional[bool] = True,
                 embedding_model_name: Optional[str] = DEFAULT_EMBEDDING_MODEL_NAME,
                 entity_label: Optional[str] = "entity"):
        """
        Initialize a new instance of Neo4jLangChainWiseAgent


        Args:
            properties (List[str]): the properties to be used as text node properties for the graph database
            collection_name (str): the collection name to use for the vector database
            url (Optional[str]): the URL of the Neo4j database (the username, password, and database name to be used
            will be obtained from the NEO4J_USERNAME, NEO4J_PASSWORD, and NEO4J_DATABASE environment variables)
            refresh_graph_schema (Optional[bool]): whether to refresh the graph schema
            embedding_model_name (Optional[str]): the optional name of the embedding model to use
            entity_label (Optional[str]): the label to use for entities in the graph database
        """
        super().__init__(embedding_model_name)
        self._properties = properties
        self._collection_name = collection_name
        self._url = url
        self._refresh_graph_schema = refresh_graph_schema
        self._entity_label = entity_label
        self._neo4j_graph_db = None
        self._neo4j_vector_db = None

    def __repr__(self):
        """Return a string representation of the graph DB."""
        return (f"{self.__class__.__name__}(properties={self.properties}, url={self.url}, refresh_schema={self.refresh_graph_schema},"
                f"embedding_model_name={self.embedding_model_name}, collection_name={self.collection_name},"
                f"entity_label={self._entity_label}")

    def __getstate__(self) -> object:
        """Return the state of the graph DB. Removing the instance variable neo4j_graph_db to avoid it being serialized/deserialized by pyyaml."""
        state = self.__dict__.copy()
        del state['_neo4j_graph_db']
        del state['_neo4j_vector_db']
        del state['_embedding_function']
        return state

    @classmethod
    def from_yaml(cls, loader, node):
        data = loader.construct_mapping(node, deep=True)
        url = data.get('_url', None)
        refresh_graph_schema = data.get('_refresh_graph_schema', True)
        embedding_model_name = data.get('_embedding_model_name', DEFAULT_EMBEDDING_MODEL_NAME)
        entity_label = data.get('_entity_label', "entity")
        return cls(properties=data.get('_properties'), collection_name=data.get('_collection_name'),
                   url=url, refresh_graph_schema=refresh_graph_schema, embedding_model_name=embedding_model_name,
                   entity_label=entity_label)

    @property
    def properties(self):
        """Get the properties to be used as text node properties for the graph database."""
        return self._properties

    @property
    def collection_name(self):
        """Get the collection name to use for the vector database."""
        return self._collection_name

    @property
    def entity_label(self):
        """Get the label to use for entities in the graph database."""
        return self._entity_label

    @property
    def url(self):
        """Get the URL of the Neo4j database."""
        return self._url

    @property
    def refresh_graph_schema(self):
        """Get whether to refresh the graph schema."""
        return self._refresh_graph_schema

    def connect(self):
        if self._neo4j_graph_db is None:
            self._neo4j_graph_db = Neo4jGraph(url=self.url, refresh_schema=self.refresh_graph_schema)

    def get_schema(self) -> str:
        self.connect()
        return self._neo4j_graph_db.get_schema

    def refresh_schema(self):
        self.connect()
        self._neo4j_graph_db.refresh_schema()

    def query(self, query: str, params: Optional[dict] = None):
        self.connect()
        return self._neo4j_graph_db.query(query=query, params=params)

    def insert_entity(self, entity: Entity, source: Source):
        self.connect()
        self.insert_graph_documents([GraphDocument(entities=[entity],
                                                   relationships=[],
                                                   source=source)])

    def insert_relationship(self, relationship: Relationship, source: Source):
        self.connect()
        self.insert_graph_documents([GraphDocument(entities=[],
                                                   relationships=[relationship],
                                                   source=source)])

    def insert_graph_documents(self, graph_documents: List[GraphDocument]):
        self.connect()
        self._neo4j_graph_db.add_graph_documents([self.convert_to_lang_chain_graph_document(graph_document)
                                                  for graph_document in graph_documents])

    def create_vector_db_from_graph_db(self, retrieval_query: str = ""):
        """
        Create a vector database from the graph database.


        Args:
            retrieval_query (str): the retrieval query to use for the vector database
        """
        self.connect()
        self._neo4j_vector_db = Neo4jVector.from_existing_graph(embedding=self._embedding_function,
                                                                node_label=self.entity_label,
                                                                embedding_node_property="embedding",
                                                                text_node_properties=self.properties,
                                                                url=self.url,
                                                                index_name=self.collection_name,
                                                                retrieval_query=retrieval_query)

    def query_with_embeddings(self, query: str, k: int, retrieval_query: str = "") -> List[Document]:
        """
        Query the vector database that corresponds to this graph database using the given query and
        retrieve the top k documents.


        Args:
            query (str): the query to execute
            k (int): the number of documents to retrieve
            retrieval_query (str): the retrieval query to use for the vector database

        Returns:
            List[Document]: the list of documents retrieved from the vector database
        """
        if self._neo4j_vector_db is None:
            # this assumes that the vector DB has already been created prior to attempting to query it
            # and we are simply retrieving the existing index here
            self.get_vector_db_from_existing_index(embedding_node_property="embedding", retrieval_query=retrieval_query)
        return [Document(content=doc.page_content, metadata=doc.metadata)
                for doc in self._neo4j_vector_db.similarity_search(query, k)]

    def get_vector_db_from_existing_index(self, embedding_node_property: str,
                                          retrieval_query: str = ""):
        # Neo4jVector.from_existing_graph doesn't currently work when using HuggingFaceEmbeddings
        # and trying to retrieve an existing graph (see https://github.com/langchain-ai/langchain/issues/24401).
        # We'll work around this using Neo4jVector.from_existing_index instead.
        self.connect()
        if not retrieval_query:
            retrieval_query = (
                    f"RETURN reduce(str='', k IN {self.properties} |"
                    " str + '\\n' + k + ': ' + coalesce(node[k], '')) AS text, "
                    "node {.*, `"
                    + embedding_node_property
                    + "`: Null, id: Null, "
                    + ", ".join([f"`{prop}`: Null" for prop in self.properties])
                    + "} AS metadata, score"
            )
        self._neo4j_vector_db = Neo4jVector.from_existing_index(embedding=self._embedding_function,
                                                                node_label=self.entity_label,
                                                                embedding_node_property="embedding",
                                                                url=self.url,
                                                                index_name=self.collection_name,
                                                                retrieval_query=retrieval_query,
                                                                keyword_index_name="keyword")

    def delete_vector_db(self):
        """
        Delete the vector database that corresponds to this graph database.
        """
        if self._neo4j_vector_db is not None:
            self._neo4j_vector_db.delete_index()
            self._neo4j_vector_db = None

    def close(self):
        """
        Close the Neo4j driver.
        """
        if self._neo4j_graph_db is not None:
            self._neo4j_graph_db._driver.close()
        if self._neo4j_vector_db is not None:
            self._neo4j_vector_db._driver.close()

class Relationship(BaseModel):
    """
    A relationship (edge) in a knowledge graph.

    Attributes:
        label (str): a description of the relationship
        source (Entity): the source entity
        target (Entity): the target entity
        metadata (Optional[dict]): optional information about the relationship
    """
    label: str
    source: Entity
    target: Entity
    metadata: Optional[dict] = Field(default_factory=dict)

class Source(BaseModel):
    """
    Information about a source from which entities and relationships have been derived from.

    Attributes:
        content (str): the content of the source
        id (str): the optional id associated with the source
        metadata (Optional[dict]): optional information about the source
    """
    content: str
    id: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4()))
    metadata: Optional[dict] = {}

class WiseAgentGraphDB(yaml.YAMLObject):
    """Abstract class to define the interface for a WiseAgentGraphDB."""

    yaml_tag = u'!WiseAgentGraphDB'

    @abstractmethod
    def get_schema(self) -> str:
        """
        Get the schema of the graph DB.

        Returns:
            str: the schema of the graph DB
        """
        ...

    @abstractmethod
    def refresh_schema(self):
        """
        Refresh the schema of the graph DB.
        """
        ...

    @abstractmethod
    def query(self, query: str, params: Optional[dict] = None) -> Any:
        """
        Query the graph DB.


        Args:
            query (str): the query to execute
            params (dict): the optional parameters for the query

        Returns:
            Any: the result of the query
        """
        ...

    @abstractmethod
    def insert_entity(self, entity: Entity, source: Source):
        """
        Insert an entity into the graph DB.


        Args:
            entity (Entity): the entity to insert
            source (Source): information about the source from which the entity has been derived from
        """
        ...

    @abstractmethod
    def insert_relationship(self, relationship: Relationship, source: Source):
        """
        Insert a relationship into the graph DB.


        Args:
            relationship (Relationship): the relationship to insert
            source (Source): information about the source from which the relationship has been derived from
        """
        ...

    @abstractmethod
    def insert_graph_documents(self, graph_documents: List[GraphDocument]):
        """
        Insert a list of graph documents into the graph DB.


        Args:
            graph_documents (List[GraphDocuments]): the graph documents to insert
        """
        ...