from flask import current_app
from neo4j import GraphDatabase
from neo4j.exceptions import ServiceUnavailable

class Neo4jDriver:
    def __init__(self):
        self._driver = None
        
    def connect(self):
        if not self._driver:
            self._driver = GraphDatabase.driver(
                current_app.config['NEO4J_URI'],
                auth=(current_app.config['NEO4J_USER'], current_app.config['NEO4J_PASSWORD']),
                encrypted=current_app.config['NEO4J_ENCRYPTED']
            )
        return self._driver
    
    def close(self):
        if self._driver:
            self._driver.close()
            self._driver = None
            
    def verify_connectivity(self):
        try:
            self.connect().verify_connectivity()
            return True
        except ServiceUnavailable:
            return False
    
    def get_session(self):
        """获取 Neo4j 会话"""
        return self.connect().session()

# 单例实例
neo4j_driver = Neo4jDriver()