import csv
import json
import os
import sys
import time

import psycopg2 as db

import settings

LOGGER = settings.get_logger('database')


def create():
    """Function to create the table underlying the analysis.

    Sufficient to be called once, but it is safe to be called multiple times.

    Returns
    -------
    int
        The number of rows in the table.
    """
    apks = 0
    try:
        db_connection = db.connect(settings.db_string)
    except db.Error:
        LOGGER.fatal('Could not establish a connection to the database.')
        sys.exit(1)
    cursor = db_connection.cursor()
    try:
        cursor.execute("CREATE TABLE apks (sha256 varchar PRIMARY KEY, dex_date date, apk_size int,"
                       " pkg_name varchar, version_code int, vt_detection int, vt_date date, dex_size int,"
                       " markets varchar[]);")
        db_connection.commit()
        LOGGER.info('Successfully created table "apks"')
    except db.Error:
        db_connection.rollback()
        cursor.execute("SELECT * FROM apks;")
        apks = cursor.rowcount
        LOGGER.info(f'Table "apks" was already present with {apks} rows')
    try:
        cursor.execute("CREATE TABLE results (sha256 varchar PRIMARY KEY, packer varchar, libs varchar[],"
                       " reflection int, reflection_code_loading boolean, class_loading boolean,"
                       " dex_class_loading boolean, encrypted varchar[], permissions varchar[], vt json);")
        db_connection.commit()
        LOGGER.info('Successfully created table "results"')
    except db.Error:
        db_connection.rollback()
        cursor.execute("SELECT * FROM results;")
        LOGGER.info(f'Table "results" was already present with {cursor.rowcount} rows')
    try:
        cursor.execute("CREATE TABLE libs (sha256 varchar PRIMARY KEY, name varchar, entropy double precision,"
                       " size int);")
        db_connection.commit()
        LOGGER.info('Successfully created table "libs"')
    except db.Error:
        db_connection.rollback()
        cursor.execute("SELECT * FROM libs;")
        LOGGER.info(f'Table "libs" was already present with {cursor.rowcount} rows')
    try:
        cursor.execute("CREATE TABLE encrypted (sha256 varchar PRIMARY KEY, name varchar, entropy double precision,"
                       " size int);")
        db_connection.commit()
        LOGGER.info('Successfully created table "encrypted"')
    except db.Error:
        db_connection.rollback()
        cursor.execute("SELECT * FROM encrypted;")
        LOGGER.info(f'Table "encrypted" was already present with {cursor.rowcount} rows')
    cursor.close()
    db_connection.close()
    return apks


def populate(filepath):
    """Populates the database with the contents of a .csv file.

    The intended use is with a description of the AndroZoo dataset (https://androzoo.uni.lu/),
    but any csv file with the correct columns will work.
    Please refer to https://androzoo.uni.lu/lists for a documentation on the format.

    Parameters
    ----------
    filepath : str
        The path to a csv file containing the information to populate the database with.

    Returns
    -------
    int
        Number of rows in the table after population.
    """
    if filepath is None or filepath == '':
        LOGGER.error('No input file found.')
        sys.exit(1)
    if not os.path.isfile(filepath):
        LOGGER.error(f'File "{filepath}" not found.')
        sys.exit(1)
    try:
        db_connection = db.connect(settings.db_string)
    except db.Error:
        LOGGER.error('Could not establish a connection to the database.')
        sys.exit(1)
    cursor = db_connection.cursor()
    start = time.time()
    with open(filepath, 'r') as csv_file:
        for row in csv.DictReader(csv_file, skipinitialspace=True):
            cursor.execute(
                "INSERT INTO apks (sha256, dex_date, apk_size, pkg_name, version_code, vt_detection, vt_date, dex_size,"
                " markets) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s);",
                (row['sha256'],
                 row['dex_date'],
                 int(row['apk_size']) if row['apk_size'] else None,
                 row['pkg_name'],
                 int(row['vercode']) if row['vercode'] else None,
                 int(row['vt_detection']) if row['vt_detection'] else None,
                 row['vt_scan_date'] if row['vt_scan_date'] else None,
                 int(row['dex_size']) if row['dex_size'] else None,
                 [market for market in row['markets'].split('|')]))
    db_connection.commit()
    cursor.execute("SELECT * FROM apks;")
    size = cursor.rowcount
    cursor.close()
    db_connection.close()
    LOGGER.info(f'Successfully populated "apks" with {size} rows. Took {time.time() - start}s')
    return size


def access(query, args=None):
    try:
        db_connection = db.connect(settings.db_string)
    except db.Error:
        LOGGER.fatal('Could not establish a connection to the database.')
        sys.exit(1)
    cursor = db_connection.cursor()
    if args:
        cursor.execute(query, args)
    else:
        cursor.execute(query)
    for row in cursor:
        yield row
    db_connection.commit()


def store_result(sha256, permissions, packer, method_results, libraries, encrypted):
    try:
        db_connection = db.connect(settings.db_string)
    except db.Error:
        LOGGER.fatal('Could not establish a connection to the database.')
        sys.exit(1)
    cursor = db_connection.cursor()
    encrypted_sha256 = [enc[0] for enc in encrypted]
    libs_sha256 = [lib[0] for lib in libraries]
    reflection_acccess, reflection_code_loading, class_loading, dex_class_loading = method_results
    try:
        cursor.execute("INSERT INTO results (sha256, packer, libs, reflection, reflection_code_loading, class_loading,"
                       " dex_class_loading, encrypted, permissions) "
                       "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING;",
                       (sha256, packer, libs_sha256, reflection_acccess, reflection_code_loading, class_loading,
                        dex_class_loading, encrypted_sha256, permissions))
        for lib in libraries:
            cursor.execute("INSERT INTO libs (sha256, name, entropy, size) VALUES (%s, %s, %s, %s)"
                           " ON CONFLICT DO NOTHING;", lib)
        for enc in encrypted:
            cursor.execute("INSERT INTO encrypted (sha256, name, entropy, size) VALUES (%s, %s, %s, %s)"
                           " ON CONFLICT DO NOTHING;", enc)
        db_connection.commit()
    except db.Error:
        db_connection.rollback()
        LOGGER.fatal(f'Could not store results for {sha256}')


def store_vt(sha256, vt_data):
    try:
        db_connection = db.connect(settings.db_string)
    except db.Error:
        LOGGER.error('Could not establish a connection to the database.')
        return
    cursor = db_connection.cursor()
    try:
        cursor.execute("UPDATE results SET vt = %s WHERE sha256 = %s;", (json.dumps(vt_data), sha256))
        db_connection.commit()
    except db.Error:
        db_connection.rollback()
        LOGGER.fatal(f'Could not update results with VirusTotal data for {sha256}')
