import logging
from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator

from neo4j import AsyncDriver, AsyncGraphDatabase

import config
from src.errors import UserProcessError

logger = logging.getLogger('users_cleanup')


_driver: AsyncDriver | None = None


@asynccontextmanager
async def Neo4jContext(concurrency: int = 1) -> AsyncGenerator[None, Any]:
    """
    Context Manager to handle the global driver lifecycle.
    Use this in your main entry point.
    """
    global _driver

    if _driver is None:
        logger.debug('Initializing Neo4j Driver...')
        _driver = AsyncGraphDatabase.driver(
            config.NEO4J_URL,
            auth=(config.NEO4J_USERNAME, config.NEO4J_PASSWORD),
            max_transaction_retry_time=config.NEO4J_MAX_RETRY_TIME,
            max_connection_pool_size=concurrency,
        )
        await _driver.verify_connectivity()
        logger.debug('Neo4j Connected.')

    try:
        yield
    finally:
        if _driver:
            logger.debug('Closing Neo4j Driver...')
            await _driver.close()
            _driver = None
            logger.debug('Neo4j Closed.')


def _get_driver() -> AsyncDriver:
    if _driver is None:
        raise RuntimeError("Neo4j driver is not initialized. Wrap your logic in 'async with Neo4jContext():'")
    return _driver


class Neo4jClient:
    """Singleton async neo4j client."""

    async def get_identity_by_email(self, email: str) -> dict[str, Any]:
        driver = _get_driver()

        query = """
            MATCH (i:Identity)
            WHERE i.email = $email
            RETURN i as identity
        """
        result = await driver.execute_query(
            query,
            {'email': email},
            routing_='r',
        )

        if len(result.records) == 0:
            raise UserProcessError('Identity not found')

        if len(result.records) > 1:
            raise UserProcessError('More than one identity found')

        return dict(result.records[0]['identity'])


db_client = Neo4jClient()