import os
import time

from neo4j import GraphDatabase

neo4j_config_dev = {
    'url': 'bolt+routing://dev-neo4j-cluster.dev.theorchard.io:7687',
    'user': os.environ.get('DEV_NEO4J_USER'),
    'password': os.environ.get('DEV_NEO4J_PASSWORD')
}

# prod is default one
neo4j_config = {
    'url': 'bolt+routing://prod-neo4j-cluster.theorchard.io:7687',
    'user': os.environ.get('PROD_NEO4J_USER'),
    'password': os.environ.get('PROD_NEO4J_PASSWORD')
}

neo4j_driver = GraphDatabase.driver(
    neo4j_config['url'],
    auth=(neo4j_config['user'], neo4j_config['password']), encrypted=True)


# make sure that there is an index on PublicParticipant(chartmetricId)
query = """
CALL apoc.export.json.query(
"
    MATCH (pp:PublicParticipant)
    WHERE pp.chartmetricId >= __start__ AND pp.chartmetricId <= __end__

    OPTIONAL MATCH (gp)-[:REPRESENTS]->(pp)

    OPTIONAL MATCH (gp)-[merged:MERGED_TO]->(other_gp)
    WHERE merged IS NULL

    OPTIONAL MATCH (gp)-[:REPRESENTS]->(lp:LabelParticipant)

    OPTIONAL MATCH (lp)-[:PARTICIPATED_IN]->(label_track:Track:Orchard)

    RETURN pp, gp, lp, COLLECT(DISTINCT label_track) AS label_tracks

", null, {stream: true})
YIELD file, source, format, nodes, relationships, properties, time, rows, batchSize, batches, done, data
RETURN file, source, format, nodes, relationships, properties, time, rows, batchSize, batches, done, data;
"""


def export_result(tx, query_parametrized, end_id):
    file_num = divmod(end_id, 100000)[0]  # write ~1000000 PublicParticipant to each file
    with open('neo4j_{}.json'.format(str(file_num)), 'at+') as new_file:
        start = time.perf_counter()
        for record in tx.run(query_parametrized):
            data = record['data']
            if data:
                new_file.write(record['data'])
        end = time.perf_counter()
        print('neo4j query time: {}'.format(end - start))


def export_to_json():
    """~6 hours in one thread."""
    start = time.perf_counter()
    start_id = 10  # skip 10 first chartmetricId ("Various artists", etc.), deal with them manually
    end_id = 110

    while end_id <= 3752477:  # MATCH (pp:PublicParticipant) RETURN MAX(pp.chartmetricId)
        with neo4j_driver.session() as session:
            query_parametrized = query.replace(
                '__start__', str(start_id)).replace('__end__', str(end_id))
            session.read_transaction(export_result, query_parametrized, end_id)
            start_id += 100
            end_id += 100

    end = time.perf_counter()
    total_time = end - start
    print(f'\nTotal time: {total_time}')

    neo4j_driver.close()


if __name__ == '__main__':
    export_to_json()