#!/usr/bin/env python3

import argparse
import os
import cv2
import logging
import persistqueue
import datetime
import pymongo
import pymongo.errors
import femida_detect


parser = argparse.ArgumentParser()
ANSWERS = 'answers'
PDFS = 'pdfs'
DB = 'femida'

parser.add_argument('--host', type=str, default='localhost')
parser.add_argument('--port', type=int, default=27017)
parser.add_argument('--log-to', type=str,
                    help='path to logfile')
parser.add_argument('--root', type=str, required=True,
                    help='path to local queue')
parser.add_argument('--model-path', type=str, required=True)
parser.add_argument('--max-errors', type=int, default=1)
parser.add_argument('--critical-errors', type=int, default=2)
parser.add_argument('--save-path', type=str, default=os.environ.get('RESULTS_FOLDER', '/media/ocr_results/'))
parser.add_argument('--cuda', action='store_true')
parser.add_argument('--debug', action='store_true')
logger = logging.getLogger('consumer :: answers')


def handle_error(task, error, args):
    logger.error(f"Task _id={task['_id']} :: {task['i']} :: status :: error :: %s", error)
    task['errors'] += 1
    if task['errors'] > args.max_errors:
        task['result'] = dict(
            page=task['i'],
            personal=[],
            UUID=task['UUID'],
            requested_manual=[],
            manual_checks=[],
            test_results=[],
            test_updates=[],
            status='error',
            flipped=False
        )
    else:
        logger.debug(f"Task _id={task['_id']} :: retry {task['errors']}/{args.max_errors}")


def main(args):
    logger.info(f'storing ocr results here: {args.save_path}')
    os.makedirs(args.save_path, exist_ok=True)
    queue = persistqueue.SQLiteQueue(
        os.path.join(args.root, 'answers'))
    logger.info(f'Opened Answers Queue {queue.path}')
    finalize_queue = persistqueue.SQLiteQueue(
        os.path.join(args.root, 'pdf.finalize')
    )
    logger.info(f'Opened Finalize Queue {finalize_queue.path}')
    logger.info(f'Connecting to {args.host}:{args.port} with provided credentials')
    conn = pymongo.MongoClient(
        host=args.host,
        port=args.port,
        username=os.environ.get('MONGODB_USERNAME'),
        password=os.environ.get('MONGODB_PASSWORD')
    )
    femida_detect.utils.setup_rs(conn, logger)
    answers = conn[DB][ANSWERS]
    pdfs = conn[DB][PDFS]
    logger.info(f'Loading model from {args.model_path}')
    predict = femida_detect.detect.eval.load(args.model_path, ('cuda' if args.cuda else 'cpu'))
    logger.info(f'Started listening')
    critical_fails_left = args.critical_errors

    def final_put(final):
        if 'result' in final:
            logger.debug(f"Got result _id={task['_id']} :: {task['i']}")
            answers.insert(final['result'])
            finalize_queue.put(final)
            os.unlink(final['imagef'])
        else:
            logger.debug(f"Got fail _id={task['_id']}")
            queue.put(final)

    while True:
        task = queue.get()
        try:
            logger.info(f"Got new task _id={task['_id']} :: {task['i']} :: {task['imagef']}")
            if 'result' in task:
                final_put(task)
                continue
            image = cv2.imread(task['imagef'])
            test_results = dict.fromkeys(list(map(str, range(1, 41))), '')
            test_updates = [{
                'session_id': 'ocr', 'updates': {}, 'date': datetime.datetime.now()}
            ]

            logger.info(f"UUID={task['UUID']}")
            pdf = pdfs.find({'UUID': task['UUID']})[0]
            pdf_comment = pdf['pdf_comment']
            _, class_, type_, variant = pdf_comment.split('_')

            if type_ == '\u041e\u0422':
                type_ = 'ot'  
            else:
                type_ = 'mat'

            result = dict(
                page=task['i'],
                personal=[{'class': class_,
                          'type': type_}],
                UUID=task['UUID'],
                requested_manual=[],
                manual_checks=[],
                test_results=test_results,
                test_updates=test_updates
            )
            logger.debug(f"Task _id={task['_id']} :: {task['i']} :: CroppedAnswers.from_raw")
            cropped = femida_detect.imgparse.CroppedAnswers.from_raw(image)
            result['flipped'] = cropped.flipped
            img_fio = os.path.join(args.save_path, f"{task['UUID']}__{task['i']}_fio.jpg")
            orig = cropped.personal
            h, w, _ = orig.shape
            cv2.imwrite(
                img_fio,
                cv2.resize(orig, (1000, 1000 * h // w))
            )
            result['img_fio'] = img_fio
            logger.debug(f"Task _id={task['_id']} :: {task['i']} :: predict(cropped)")
            predictions = predict(cropped)
            logger.debug(f"Task _id={task['_id']} :: {task['i']} :: cropped.plot_predicted(predictions)")
            inpainted = cropped.plot_predicted(predictions, only_answers=True)
            img_test_form = os.path.join(args.save_path, f"{task['UUID']}__{task['i']}_test_form.jpg")
            h, w, _ = inpainted.shape
            cv2.imwrite(
                img_test_form,
                cv2.resize(inpainted, (1000, 1000 * h // w))
            )
            result['img_test_form'] = img_test_form
            for (j, letter), pred in zip(
                    cropped.get_labels(),
                    predictions
            ):
                if pred:
                    test_results[str(j)] += letter
            logger.info(f"Task _id={task['_id']} :: {task['i']} :: status :: normal")
            result.update(status='normal')
            task['result'] = result
            final_put(task)
        except pymongo.errors.PyMongoError:
            logger.error(f"Task _id={task['_id']} :: recoverable (restart) :: %s", e)
            break
        except cv2.error as e:
            handle_error(task, e, args)
            final_put(task)
            continue
        except Exception as e:
            logger.critical(f"Task _id={task['_id']} :: {task['i']} :: {e}")
            if critical_fails_left == 0:
                final_put(task)
                raise e
            critical_fails_left -= 1
            handle_error(task, e, args)
            final_put(task)
            continue


if __name__ == '__main__':
    pargs = parser.parse_args()
    formatter = logging.Formatter('%(asctime)s :: %(name)s :: %(levelname)s - %(message)s')
    if pargs.log_to is not None:
        handle = logging.FileHandler(pargs.log_to)
    else:
        handle = logging.StreamHandler()
    handle.setFormatter(formatter)
    logger.addHandler(handle)
    if pargs.debug or os.environ.get('FEMIDA_DEBUG', False):
        logger.setLevel(logging.DEBUG)
    main(pargs)
