#!/usr/bin/env python
from __future__ import print_function

import os
import sys
import subprocess
import json
import threading
import timeit

from logger import Logger, Msg
import psycopg2

def movie_in_db(oclc_id):
    '''
    Checks whether or not a movie exists in the media_metadata
    table of the Database

    :param: oclc_id The OCLC ID of the movie to check for

    :returns: A boolean representing whether or  not the movie was found
    '''
    READ_CUR_T1.execute('''
        SELECT oclc_id FROM media_metadata
        WHERE oclc_id = %(oclc_id)s LIMIT 1;
    ''', {'oclc_id': oclc_id})
    res = None
    try:
        res = READ_CUR_T1.fetchone()
    except psycopg2.ProgrammingError as err:
        print(err.message)
    if res is None:
        return False
    return True


def line_in_db(oclc_id, line_no):
    '''
    Checks whether or not a line exists in the media_recognized_objects
    table of the Database

    :param: oclc_id The OCLC ID of the movie to check for
    :param: line_no The line number in the movie
    
    :returns: A boolean representing whether or not the line was found
    '''
    data = {'oclc_id': int(oclc_id), 'line_no': int(line_no)}
    READ_CUR_T1.execute('''
        SELECT db_line_id FROM media_recognized_objects
        WHERE db_line_id in (
            SELECT db_line_id FROM media_text
            WHERE oclc_id = %(oclc_id)s
            AND line_number = %(line_no)s
            LIMIT 1
        ) LIMIT 1;
    ''', data)
    res = None
    try:
        res = READ_CUR_T1.fetchone()
    except psycopg2.ProgrammingError as err:
        print(err.message, 'for', line_no, '@', oclc_id)
    if res is None:
        return False
    return True


def process(img):
    '''
    Pipes a filename into darknet for recognition

    :param: img The path to the image
    '''
    LOGGER.log(
        Msg(
            'traverse',
            'Adding: {}'.format(img),
            prefix_color='green'
        ),
        tag_id='T1_DISP'
    )
    DARKNET_PROC.stdin.write(img + '\n')
    DARKNET_PROC.stdin.flush()


def listfiles(media_dir):
    '''
    Recursively generates all of the files in the given directory
    that have .png and .jpg extensions

    :param: media_dir The directory to list files from
    '''
    for root, _, files in os.walk(media_dir):
        for filename in files:
            if filename.endswith('.jpg') or filename.endswith('.png'):
                yield os.path.join(root, filename)


def traverse(media_dir):
    '''
    Traverses all of the files in media_dir, and checks if a given
    one has been processed already. If it has not, adds the file
    for processing. Needs to be run in a thread

    :param: media_dir The directory to recursively check
    '''
    for img in listfiles(media_dir):
        path = os.path.dirname(img)
        fname = os.path.basename(img)
        oclc_id = os.path.basename(path)
        line_no = os.path.splitext(fname)[0]
        if movie_in_db(oclc_id):
            if not line_in_db(oclc_id, line_no):
                process(img)
    process('')


def upload():
    '''
    Uploads data that is generated by the DARKNET_PROC to the
    database. Needs to be run in a thread
    '''
    step = 0

    movies = {}

    while DARKNET_PROC.poll() is None:
        recognized_object = DARKNET_PROC.stdout.readline().rstrip()
        # print(recognized_object)
        # Done - exit
        if (recognized_object == ''):
            break

        # JSON data
        data = json.loads(recognized_object)

        oclc_id = data['oclc_id']
        line_no = data['db_line_id']

        # Populate cache
        if oclc_id not in movies.keys():
            movies[oclc_id] = {}
        if line_no not in movies[oclc_id].keys():
            # Get DB Line ID
            READ_CUR_T2.execute('''
                SELECT db_line_id FROM media_text WHERE oclc_id = %(oclc_id)s AND line_number = %(db_line_id)s LIMIT 1;
            ''', data)
            res = READ_CUR_T2.fetchone()
            if res is None:
                LOGGER.log(
                    Msg(
                        'error',
                        'Postgres entry not found for: {line_no} @ {oclc_id}'.format(**locals()),
                        prefix_color='red',
                        msg_color='red'
                    )
                )
                continue
            else:
                movies[oclc_id][line_no] = res[0]

        data['db_line_id'] = movies[oclc_id][line_no]

        # Insert new data
        WRITE_CUR_T2.execute('''
            INSERT INTO media_recognized_objects
              (db_line_id, text_label, confidence, bounding_left, bounding_right, bounding_top, bounding_bottom)
            VALUES
              (%(db_line_id)s, %(label)s, %(confidence)s, %(l)s, %(r)s, %(t)s, %(b)s)
            ON CONFLICT DO NOTHING;
            ''', data)

        LOGGER.log(
            Msg(
                'upload',
                '{step} entered into DB'.format(**locals()),
                prefix_color='blue'
            ),
            tag_id='T2_COUNT'
        )
        step += 1

    # commit all remaining lines
    LOGGER.log(
        Msg(
            'upload',
            '{step} entered into DB'.format(**locals()),
            prefix_color='blue'
        ),
        tag_id='T2_COUNT'
    )
    CONN.commit()
    CONN.close()


if __name__ == '__main__':

    #==============================
    #      LOGGER/DB SETUP
    #==============================

    LOGGER = Logger()

    LOGGER.log(Msg('populate.py', 'Populate Database'))
    LOGGER.log(Msg('populate.py', '-----------------'))

    THIS_FILE = os.path.dirname(os.path.realpath(__file__))
    CFG_FILE = os.path.join(THIS_FILE, '..', 'src', 'dbConfig.json')

    # Load the config file
    with open(CFG_FILE) as DB_CFG:
        PG_CFG = json.loads(DB_CFG.read())

    DB_CFG = {
        'dbname': 'filmtvse',
        'user': PG_CFG["DATABASE"]["username"],
        'password': PG_CFG["DATABASE"]["password"],
        'host': PG_CFG["DATABASE"]["host"],
        'port': PG_CFG["DATABASE"]["port"]
    }

    LOGGER.tag('db_login').log(Msg('DB Login', 'Logging into DB...', prefix_color='yellow'))
    try:
        CONN = psycopg2.connect(**DB_CFG)
        READ_CUR_T1 = CONN.cursor()
        READ_CUR_T2 = CONN.cursor()
        WRITE_CUR_T2 = CONN.cursor()
        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, READ_CUR_T1)
        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, READ_CUR_T2)
        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, WRITE_CUR_T2)
    except Exception as err:
        LOGGER.log(
            Msg(
                'DB Login',
                'Login Failed: {err.message}'.format(**locals()),
                prefix_color='yellow',
                msg_color='red'
            ),
            tag_id='db_login'
        )
        sys.exit(1)
    LOGGER.log(
        Msg(
            'DB Login',
            'Login Success!'.format(**locals()),
            prefix_color='yellow',
            msg_color='green'
        ),
        tag_id='db_login'
    )

    #==============================
    #       DARKNET SETUP
    #==============================

    # Start darknet
    DARKNET_DIR = os.path.join(THIS_FILE, '..', 'submodules', 'darknet')
    DARKNET = os.path.join('.', 'darknet')
    CMD = './darknet detector test cfg/coco.data cfg/yolo.cfg yolo.weights'
    DARKNET_PROC = subprocess.Popen(
        CMD,
        cwd=DARKNET_DIR,
        shell=True,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

    #==============================
    #         DATA ENTRY
    #==============================

    START = timeit.default_timer()

    # Start thread for adding all necessary files
    T1 = threading.Thread(target=traverse, args=[sys.argv[1]])
    T1.start()
    LOGGER.tag('T1_WAIT').log(Msg('traverse', 'Traversing directories', prefix_color='green'))
    LOGGER.wait('T1_WAIT', Msg('traverse', 'Traversing directories', prefix_color='green'), T1)
    LOGGER.tag('T1_DISP').log(Msg('traverse', '', prefix_color='green'))

    # Start thread for uploading to DB
    T2 = threading.Thread(target=upload)
    T2.start()
    LOGGER.tag('T2_WAIT').log(Msg('upload', 'Entering into DB', prefix_color='blue'))
    LOGGER.wait('T2_WAIT', Msg('upload', 'Entering into DB', prefix_color='blue'), T2)
    LOGGER.tag('T2_COUNT').log(Msg('upload', '0 Entered into DB', prefix_color='blue'))

    # Wait for threads to finish
    T1.join()
    LOGGER.log(
        Msg(
            'traverse',
            'Done in {}s'.format(timeit.default_timer() - START),
            prefix_color='green',
            msg_color='green'
        ),
        tag_id='T1_WAIT'
    )

    DARKNET_PROC.wait()

    T2.join()
    LOGGER.log(
        Msg(
            'upload',
            'Done in {}s'.format(timeit.default_timer() - START),
            prefix_color='blue',
            msg_color='green'
        ),
        tag_id='T2_WAIT'
    )

    LOGGER.log(Msg('populate.py', '-----------------'))
    LOGGER.log(Msg('populate.py', 'Done.'))
