import os
from neo4j import GraphDatabase
from contextlib import contextmanager
from typing import Any, List, Dict, Optional


class Neo4jClient:
    def __init__(self):
        self.uri = os.getenv("NEO4J_URI", "bolt://localhost:7687")
        self.username = os.getenv("NEO4J_USERNAME", "neo4j")
        self.password = os.getenv("NEO4J_PASSWORD", "password")
        self.driver = None

    def connect(self):
        try:
            self.driver = GraphDatabase.driver(
                self.uri,
                auth=(self.username, self.password),
            )
            self.driver.verify_connectivity()
        except Exception as e:
            raise ConnectionError(f"Failed to connect to Neo4j: {e}")

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

    @contextmanager
    def session(self):
        if not self.driver:
            self.connect()
        session = self.driver.session()
        try:
            yield session
        finally:
            session.close()

    def query(self, cypher: str, parameters: Optional[Dict[str, Any]] = None) -> List[Dict]:
        """Execute a Cypher query and return results."""
        with self.session() as session:
            result = session.run(cypher, parameters or {})
            return [record.data() for record in result]

    def find_nodes(self, label: str, filters: Optional[Dict[str, Any]] = None) -> List[Dict]:
        """Find nodes by label and optional filters."""
        filters = filters or {}
        filter_str = " AND ".join([f"n.{k} = ${k}" for k in filters.keys()])
        where_clause = f" WHERE {filter_str}" if filter_str else ""
        query = f"MATCH (n:{label}){where_clause} RETURN n"
        with self.session() as session:
            result = session.run(query, filters)
            return [record.data() for record in result]