from decouple import config
from neo4j import GraphDatabase

class Neo4jConnection:
    def __init__(self):
        self.uri = config("NEO4J_URI")
        self.user = config("NEO4J_USER")
        self.password = config("NEO4J_PASSWORD")
        self.driver = GraphDatabase.driver(self.uri, auth=(self.user, self.password))

        # Create the schema
        self.create_schema()

        self.test_retrieval()

        # Close the connection
        self.close()

    def close(self):
        self.driver.close()

    def create_schema(self):
        with self.driver.session() as session:
            session.write_transaction(self._create_constraints_and_indexes)

    def test_retrieval(tx):
        #run MATCH (n) RETURN n LIMIT 25
        result = tx.run("MATCH (n) RETURN n LIMIT 25;")


    @staticmethod
    def _create_constraints_and_indexes(tx):
        # Constraints and indexes for Person
        tx.run("CREATE CONSTRAINT ON (p:Person) ASSERT p.person_id IS UNIQUE;")
        
        # Constraints and indexes for Memory
        tx.run("CREATE CONSTRAINT ON (m:Memory) ASSERT m.memory_id IS UNIQUE;")
        tx.run("CREATE INDEX ON :Memory(content);")
        tx.run("CREATE INDEX ON :Memory(timestamp);")
        
        # Constraints and indexes for Tag
        tx.run("CREATE CONSTRAINT ON (t:Tag) ASSERT t.tag_id IS UNIQUE;")
        tx.run("CREATE INDEX ON :Tag(tag_name);")
        
        # Constraints and indexes for Category
        tx.run("CREATE CONSTRAINT ON (c:Category) ASSERT c.category_id IS UNIQUE;")
        
        # Constraints and indexes for Skill
        tx.run("CREATE CONSTRAINT ON (s:Skill) ASSERT s.name IS UNIQUE;")
        
        # Constraints and indexes for Fact
        tx.run("CREATE CONSTRAINT ON (f:Fact) ASSERT f.certainty IS UNIQUE;")
        
        # Additional schema definitions can be added here, such as relationship constraints or more complex indexing strategies.

if __name__ == "__main__":
    conn = Neo4jConnection()