"""Neo4J DB connector."""
import config

from lambdacommon.common_config import logger

from neo4j import GraphDatabase

from splitio import get_factory
from splitio.exceptions import TimeoutException


split = None
try:
    factory = get_factory(config.SPLITIO_API_KEY)
    factory.block_until_ready(config.SPLIT_BLOCK_UNTIL_READY_TIMEOUT)
    split = factory.client()
except TimeoutException as err:
    logger.error(f'Failed to start split factory due to: {str(err)}')
    raise


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_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


neo4j_driver = None
if config.ENVIRONMENT == config.TEST_ENVIRONMENT:
    neo4j_driver = None
else:
    is_enabled = split and split.get_treatment(
        config.APPLICATION_NAME,
        'neo4j_aura',
        {'service': config.APPLICATION_NAME}
    ) == 'on'
    if is_enabled:
        logger.info('Configuring connection to Aura')
        neo4j_driver = _get_aura_driver()
    if not neo4j_driver:
        # fallback to old neo4j driver.
        logger.info('Configuring connection to Neo4j')
        neo4j_driver = _get_driver()
    # Split client is not required anymore
    split.destroy()