"""Neo4j monitor."""

import csv
import datetime
import io
import json
import os
import textwrap
import zipfile

import boto3
import botocore.exceptions
import neo4j
import prettytable


AWS_REGION = os.environ.get('AWS_DEFAULT_REGION', 'us-east-1')
MAX_TABLE_COLUMN_WIDTH = int(os.environ.get('MAX_TABLE_COLUMN_WIDTH', 70))

boto_session = boto3.Session()


queries = {
    'cluster-overview': 'SHOW SERVERS YIELD * RETURN *;',
    'list-transactions': 'SHOW TRANSACTIONS YIELD * RETURN *;',
    'list-connections': (
        'CALL dbms.listConnections() '
        'YIELD username, connectTime, serverAddress '
        'RETURN username, min(connectTime) AS minConnectTime, max(connectTime) AS maxConnectTime, ' 
        'count(*) as connectionCount, collect(DISTINCT serverAddress) AS serverList '
        'ORDER BY connectionCount DESC;'
    ),
}


def get_neo_credentials_from_secretsmanager(secret_arn):
    """Fetch Neo4j credentials from SecretsManager."""
    client = boto_session.client(
        service_name='secretsmanager',
        region_name=AWS_REGION,
    )

    try:
        secret = client.get_secret_value(
            SecretId=secret_arn
        )
        data = json.loads(secret['SecretString'])

        url = data['URL'] if 'URL' in data else None
        username = data['USERNAME'] if 'USERNAME' in data else None
        password = data['PASSWORD'] if 'PASSWORD' in data else None

        if None in (url, username, password):
            raise SystemExit('Error getting Neo4j credentials from '
                             'SecretsManager. At least one of the URL, '
                             'USERNAME, or PASSWORD key-pair does not exist.')

        return (url, username, password)
    except botocore.exceptions.ClientError as e:
        raise SystemExit(f'Error fetching secret {secret_arn} from '
                         f'SecretsManager: {e}')


def get_neo_credentials():
    """Fetch Neo4j credentials."""
    credentials_arn = os.environ.get('NEO4J_CREDENTIALS_SECRET_ARN')
    if credentials_arn:
        return get_neo_credentials_from_secretsmanager(credentials_arn)

    url = os.environ.get('NEO4J_URL')
    username = os.environ.get('NEO4J_USERNAME')
    password = os.environ.get('NEO4J_PASSWORD')

    if None in (url, username, password):
        raise SystemExit('Error getting Neo4j credentials. They should be '
                         'configured as one of:\n'
                         '- environment variables: NEO4J_URL, NEO4J_USERNAME '
                         'and NEO4J_PASSWORD\n'
                         '- SecretsManager key-value pairs URL, USERNAME and '
                         'PASSWORD with ARN specified in '
                         'NEO4J_CREDENTIALS_SECRET_ARN environment variable')

    return (url, username, password)


def convert_to_printable(results):
    """Convert results to a human-readable form."""
    printable_results = {}
    for id, result in results.items():
        records = []
        max_width = 0
        for row in result['data']:
            record = []
            for k, v in row.items():
                text = f'{v}'
                width = len(text)
                if width > MAX_TABLE_COLUMN_WIDTH:
                    width = MAX_TABLE_COLUMN_WIDTH
                    text = textwrap.fill(text, width=MAX_TABLE_COLUMN_WIDTH)
                if width > max_width:
                    max_width = width
                record.append([k, text])
            records.append(record)

        tables = []
        for record in records:
            value_text = 'Value' + ' ' * (max_width - 5)
            table = prettytable.PrettyTable(border=True)
            table.field_names = ['Key', value_text]
            table.add_rows(record)
            table.align = 'l'
            tables.append(table.get_string())

        printable_results[id] = '\n\n'.join(tables)

    return printable_results


def print_to_stdout(printable_results):
    """Print results as a table to stdout."""
    for id, content in printable_results.items():
        print(f'Result for {id} query:\n{content}\n')


def convert_to_csv(results):
    """Convert results to a csv form."""
    csv_results = {}
    for id, result in results.items():
        csv_buffer = io.StringIO()
        writer = csv.DictWriter(csv_buffer, fieldnames=result['keys'])
        writer.writeheader()
        writer.writerows(result['data'])

        csv_results[id] = csv_buffer.getvalue()

    return csv_results


def export_to_zip(results, file_type='txt'):
    """Export results to a bytes buffer which represents zip file."""
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, mode='w',
                         compression=zipfile.ZIP_DEFLATED) as zip_file:
        for id, content in results.items():
            zip_file.writestr(f'{id}.{file_type}', content)
    zip_buffer.seek(0)

    return zip_buffer


def upload_to_s3(bucket_name, file_name, content):
    """Upload file content to S3 bucket_name."""
    try:
        client = boto_session.client(service_name='s3')

        client.upload_fileobj(content, bucket_name, file_name)
    except botocore.exceptions.ClientError as e:
        raise SystemExit(f'Error uploading file to S3 bucket {bucket_name} '
                         f'as {file_name}: {e}')


def export_to_s3(printable_results, csv_results):
    """Store results as files in a ZIP archive and upload it to S3."""
    bucket_name = os.environ.get('S3_BUCKET_NAME')
    bucket_prefix = os.environ.get('S3_BUCKET_PREFIX', '')

    if not bucket_name:
        print('Skipping upload to S3 as there is no S3_BUCKET_NAME '
              'environment variable configured.')
        return

    print('Uploading results to S3:')

    now = datetime.datetime.now()
    name = now.strftime('%Y-%m-%dT%H%M%S')
    file_prefix = name if bucket_prefix == '' else f'{bucket_prefix}/{name}'

    file_name = f'{file_prefix}-txt.zip'
    upload_to_s3(bucket_name, file_name, export_to_zip(printable_results))
    print(f'- txt output has been uploaded to s3://{bucket_name}/{file_name}')

    file_name = f'{file_prefix}-csv.zip'
    upload_to_s3(bucket_name, file_name, export_to_zip(csv_results, 'csv'))
    print(f'- csv output has been uploaded to s3://{bucket_name}/{file_name}')


def lambda_handler(event, context):
    """Entrypoint function."""
    url, username, password = get_neo_credentials()
    driver = neo4j.GraphDatabase.driver(url, auth=(username, password))

    results = {}
    for id, query in queries.items():
        print(f'Running query: {query}')
        session = driver.session()
        result = session.run(query)
        results[id] = {
            'keys': result.keys(),
            'data': result.data()
        }
    driver.close()

    printable_results = convert_to_printable(results)
    print_to_stdout(printable_results)

    csv_results = convert_to_csv(results)
    export_to_s3(printable_results, csv_results)

    return {}