"""Create Missing OSRs."""

import datetime
import json
import os
import sys
import time

import boto3

from neo4j import GraphDatabase

LOG_FILE = os.path.join(os.path.dirname(__file__), 'osr-logs.txt')
LOG_INFO = 'INFO'
LOG_ERROR = 'ERROR'

ENVIRONMENT = os.environ.get('ENVIRONMENT', 'qa')
NEO4J_URL = os.environ.get('NEO4J_URL', '***')
NEO4J_USERNAME = os.environ.get('NEO4J_USERNAME', '***')
NEO4J_PASSWORD = os.environ.get('NEO4J_PASSWORD', '***')
NEO4J_MAX_RETRY_TIME = 30


TRACK_LIMIT = int(os.environ.get('TRACK_LIMIT', 100))

GET_TRACK_IDS = """
    MATCH(tr:Track)-[x:HAS_ASSET]->(oa:OrchardAsset)-[y:FINGERPRINTED_AS]->(acr:ACRID)
    WHERE NOT (acr)<-[:CONTAINS_FINGERPRINT]-(:OrchardSoundRecording)
    AND tr.id > toInteger($last_id)
    AND tr.isrc IS NOT NULL
    AND NOT tr.isrc = ''
    WITH
        DISTINCT(tr.id) AS tuids
        ORDER BY tr.id ASC
        LIMIT $provided_limit
    MATCH(tr:Track)-[x:HAS_ASSET]->(oa:OrchardAsset)
    WHERE tr.id IN [tuids]
    RETURN
        tr.id AS tuid,
        count(oa) AS num_assets
"""

CHECK_TRACK_WITH_OSR = """
    MATCH(tr:Track)-[:HAS_ASSET]->(:OrchardAsset)-[:FINGERPRINTED_AS]->(:ACRID)<-[:CONTAINS_FINGERPRINT]-(:OrchardSoundRecording)
    WHERE tr.id = $provided_id
    RETURN tr
"""


def log(log_type, log_msg, final_log=False):
    """Write log line to file."""
    with open(LOG_FILE, 'a') as log_file:
        ts = datetime.datetime.now().isoformat()
        msg = '[{}] [{}] {}\n'.format(ts, log_type, log_msg)
        log_file.write(msg)

        if final_log:
            log_file.write('----------\n')


def _get_driver():
    """Connect to Neo4J and return Neo4J driver."""
    try:
        driver = GraphDatabase.driver(
            NEO4J_URL,
            auth=(NEO4J_USERNAME, NEO4J_PASSWORD),
            encrypted=True,
            max_transaction_retry_time=NEO4J_MAX_RETRY_TIME)
        return driver
    except Exception as e:
        raise e


def _get_track_ids_to_process(neo4j_driver, last_id):
    with neo4j_driver.session() as session:
        nodes = session.run(
            GET_TRACK_IDS,
            provided_limit=TRACK_LIMIT,
            last_id=last_id
        )
        rows = [x for x in nodes.data()]
        session.close()
        return rows


def _check_track_processed(neo4j_driver, track_id):
    with neo4j_driver.session() as session:
        nodes = session.run(
            CHECK_TRACK_WITH_OSR,
            provided_id=track_id
        )
        result = bool(nodes.data())
        session.close()
        return result


def chunk(lst, n):
    """Break list into equal size chunks."""
    for i in range(0, len(lst), n):
        yield lst[i:i + n]


try:
    log(LOG_INFO, 'Starting osr create script for {}'.format(ENVIRONMENT))
    neo4j_driver = _get_driver()
    last_track_id = 0

    lambda_client = boto3.client('lambda', region_name='us-east-1')

    while True:
        log(LOG_INFO, 'Fetching batch of track ids')
        rows = _get_track_ids_to_process(neo4j_driver, last_track_id)

        if not rows:
            log(LOG_INFO, 'No tracks to process')
            break

        # move cursor up for next batch
        last_track_id = rows[-1]['tuid']

        # filter out problem tracks with > 1 asset
        track_ids = [
            x['tuid']
            for x in rows
            if x['num_assets'] == 1
        ]
        if not track_ids:
            continue

        # fetch new batch to process
        for batch_ids in chunk(track_ids, 10):
            payload = json.dumps({'track_ids': batch_ids}).encode('utf-8')
            lambda_client.invoke(
                FunctionName='{}-lambda-sr-create'.format(ENVIRONMENT),  # noqa:E501
                InvocationType='Event',
                Payload=payload
            )
            for track_id in batch_ids:
                log(LOG_INFO, 'Track with id: {} sent for OSR create'.format(track_id))  # noqa:E501

        # verify previous batch is done
        track_id_to_check = track_ids[-1]
        while True:
            if _check_track_processed(neo4j_driver, track_id_to_check):
                break
            log(LOG_INFO, 'Track with id: {} has no OSR connection, waiting....'.format(track_id_to_check))  # noqa:E501
            time.sleep(15)

    neo4j_driver.close()
    log(LOG_INFO, 'End osr create script', True)
    sys.exit()
except Exception as e:
    log(LOG_ERROR, 'Unexpected error: {}'.format(e), True)
    sys.exit(e)