"""Neo4j connector Client."""

from contextlib import contextmanager

from flask import g
from neo4j import GraphDatabase
from pythonfeatures import pythonfeatures
from pythonfeatures.constants import split as split_constants

from permissions import config
from permissions.api import app
from permissions.connectors.sentry import sentry_client
from permissions.constants import constants

neo4j_driver = GraphDatabase.driver(
    config.NEO4J_URL,
    auth=(config.NEO4J_USERNAME, config.NEO4J_PASSWORD),
    max_transaction_retry_time=config.NEO4J_MAX_RETRY_TIME,
)
aura_driver = None

try:
    aura_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,
    )
except Exception:
    sentry_client.capture_exception()
    print('Failed to connect to Aura on app start.')


@app.teardown_appcontext
def close_db(exc_info):
    """Close session on request teardown.

    So there is no need to do that throughout the code. This gets called
    once per request, so a session lives for each request instead of
    each query scope.

    Args:
        exc_info(exception): exception instance, if it was not handled.
    """
    if hasattr(g, 'neo4j_sessions'):
        for session in g.neo4j_sessions.values():
            session.close()


@contextmanager
def db_session(access_mode=constants.NEO4j_READ_ACCESS):
    """Provide a session for series of operations.

    Session will be closed once the flask request is done.
    Commit and rollbacks has to be managed by caller.
    Maintains separate read and write sessions per flask request.

    READ and WRITE sessions are kept separate to avoid transaction management conflicts.
    READ operations will always use a READ session with auto-commit support.
    WRITE operations will always use a WRITE session requiring explicit transactions.

    Args:
        access_mode: Access mode (read or write) for transactions in this session.
                    Sessions are cached per access mode within each flask request.
    Usage:
        # Read operation (auto-commit)
        with db_session(access_mode=constants.NEO4j_READ_ACCESS) as session:
            session.run("MATCH (i:Identity {email: $email}) RETURN i", email=email)

        # Write operation requires explicit transaction
        with db_session(access_mode=constants.NEO4j_WRITE_ACCESS) as session:
            tx = session.begin_transaction()
            try:
                result = tx.run("CREATE (i:Identity {email: $email})", email=email)
                tx.commit()
            except Exception as e:
                tx.rollback()
                raise

        # Multiple operations in write transaction
        with db_session(access_mode=constants.NEO4j_WRITE_ACCESS) as session:
            tx = session.begin_transaction()
            try:
                tx.run("CREATE (i:Identity {email: 't1@t.com'})")
                tx.run("CREATE (n:Person {email: 't2@t.com'})")
                tx.commit()
            except Exception as e:
                tx.rollback()
                raise
    """
    if not hasattr(g, 'neo4j_sessions'):
        g.neo4j_sessions = {}

    # use separate sessions for READ and WRITE operations
    if access_mode not in g.neo4j_sessions:
        g.neo4j_sessions[access_mode] = _get_neo4j_session(access_mode)

    yield g.neo4j_sessions[access_mode]


def _get_neo4j_session(access_mode):
    """Select appropriate driver based on FFlag."""
    is_enabled = (
        pythonfeatures.get_single_feature_by_attributes(
            'neo4j_aura', {'service': config.SERVICE_NAME}
        )
    ).message == split_constants.FEATURE_ENABLED
    if is_enabled and aura_driver:
        return aura_driver.session(
            default_access_mode=access_mode, database=config.NEO4J_DATABASE_NAME
        )
    # default
    return neo4j_driver.session(
        default_access_mode=access_mode, database=config.NEO4J_DATABASE_NAME
    )