"""Neo4J DB connector."""

from lambdacommon.common_config import logger
from neo4j import GraphDatabase
from splitio import get_factory
from splitio.exceptions import TimeoutException

import config

neo4j_driver = None


def _get_driver():
    """Connect to Neo4J and return Neo4J driver."""
    try:
        driver = GraphDatabase.driver(
            config.NEO4J_URL,
            auth=(config.NEO4J_USERNAME, config.NEO4J_PASSWORD),
            max_transaction_retry_time=config.NEO4J_MAX_RETRY_TIME,
        )
        return driver
    except Exception as e:
        logger.exception('Failed to connect to Neo4J.')
        raise e


def _get_aura_driver():
    """Connect to Neo4J Aura and return driver."""
    try:
        driver = GraphDatabase.driver(
            config.NEO4J_AURA_URL,
            auth=(config.NEO4J_AURA_USERNAME, config.NEO4J_AURA_PASSWORD),
            max_transaction_retry_time=config.NEO4J_MAX_RETRY_TIME,
        )
        return driver
    except Exception as e:
        logger.exception(f'Failed to connect to Neo4J Aura due to: {str(e)}')
        # don't raise exception. use fallback.


def get_neo4j_driver():
    """Return a cached Neo4j driver."""
    global neo4j_driver

    if neo4j_driver is not None or config.ENVIRONMENT == config.TEST_ENVIRONMENT:
        return neo4j_driver

    split = None
    aura_enabled = False

    try:
        if config.SPLITIO_API_KEY:
            try:
                factory = get_factory(config.SPLITIO_API_KEY)
                factory.block_until_ready(config.SPLIT_BLOCK_UNTIL_READY_TIMEOUT)
                split = factory.client()
                aura_enabled = (
                    split.get_treatment(
                        config.APPLICATION_NAME,
                        'neo4j_aura',
                        {'service': config.APPLICATION_NAME},
                    )
                    == 'on'
                )
            except TimeoutException as err:
                logger.error(f'Failed to start split factory due to: {str(err)}')
        else:
            logger.info('Split is not configured. Falling back to Neo4j without feature flag evaluation.')

        if aura_enabled:
            logger.info('Configuring connection to Aura')
            neo4j_driver = _get_aura_driver()

        if not neo4j_driver:
            # Fallback to the legacy Neo4j deployment when Aura is unavailable.
            logger.info('Configuring connection to Neo4j')
            neo4j_driver = _get_driver()

        return neo4j_driver
    finally:
        if split:
            split.destroy()