# Copyright (c) 2018 gps.az LLC
#
# @author:      Tural Mahmudov <prodoccenter@gmail.com>
#               https://github.com/NimaBavari
#
# @date:        11.01.2018
#
# @description: Creates the desktop GUI application as one of our front-end
#               points
import json
import os
import sys

import cv2
import face_recognition
from PIL import Image
from PyQt5.QtCore import pyqtSlot, Qt
from PyQt5.QtGui import QIcon, QStandardItemModel
from PyQt5.QtWidgets import (QApplication, QDialog, QDialogButtonBox,
                             QFileDialog, QFormLayout, QGridLayout, QGroupBox,
                             QHBoxLayout, QLabel, QLineEdit, QMainWindow,
                             QMessageBox, QPushButton, QTreeView, QVBoxLayout,
                             QWidget)

import config
from helpers import (add_to_json, add_user_to_db, add_user_to_model,
                     edit_user_in_db, edit_user_in_model, get_stranger_faces,
                     populate_model, remove_from_json, remove_user_in_db,
                     remove_user_in_model)


class AddDialog(QDialog):

    def __init__(self, title, parent, **dial_attr):
        super(AddDialog, self).__init__()
        self.title = title
        self.parent = parent
        self.dial_attr = dial_attr
        self.img_to_upload = None
        self.create_form_group_box()

        self.buttonBox = QDialogButtonBox(
            QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        self.buttonBox.button(QDialogButtonBox.Ok).setText('Add')
        self.buttonBox.accepted.connect(self.submit_add)
        self.buttonBox.rejected.connect(self.reject)

        main_layout = QVBoxLayout()
        main_layout.addWidget(self.formGroupBox)
        main_layout.addWidget(self.buttonBox)
        self.setLayout(main_layout)

        self.setWindowTitle(self.title)
        self.setWindowIcon(self.dial_attr['favicon'])

    def create_form_group_box(self):
        self.formGroupBox = QGroupBox('Add New Member')
        layout = QFormLayout()
        self.f_name_box = QLineEdit()
        layout.addRow(QLabel('First name:'), self.f_name_box)
        self.l_name_box = QLineEdit()
        layout.addRow(QLabel('Last name:'), self.l_name_box)
        self.img_url_btn = QPushButton('Upload image')
        layout.addRow(QLabel('Upload image:'), self.img_url_btn)
        self.img_url_btn.clicked.connect(self.img_upload_dialog)
        self.formGroupBox.setLayout(layout)

    @pyqtSlot()
    def img_upload_dialog(self):
        options = QFileDialog.Options()
        options |= QFileDialog.DontUseNativeDialog
        self.img_to_upload, _ = QFileDialog.getOpenFileName(
            self,
            'Upload an image',
            os.path.expanduser('~'),
            'All Files (*);;Image Files (*.jpg *.jpeg *.png *.jfif)',
            options=options
        )

    @pyqtSlot()
    def submit_add(self):
        first_name = self.f_name_box.text()
        last_name = self.l_name_box.text()
        if first_name and last_name and self.img_to_upload:
            full_name = '%s %s' % (first_name, last_name)
            temp_img = face_recognition.load_image_file(self.img_to_upload)
            if face_recognition.face_encodings(temp_img):
                _, ext = os.path.splitext(self.img_to_upload)
                img_name = '%s_%s%s' % (
                    first_name.lower(), last_name.lower(), ext)
                img_path = os.path.join(config.member_dir, img_name)
                img = Image.open(self.img_to_upload)
                img.save(img_path)
                add_user_to_db(first_name, last_name, img_path)
                QMessageBox.about(
                    self,
                    'Success',
                    '%s successfully added as a new customer.' % full_name
                )
                face = face_recognition.face_encodings(temp_img)[0]
                capture_times_in_db, face_encs_in_db = get_stranger_faces()
                res = face_recognition.compare_faces(face_encs_in_db, face)
                capture_times = [capture_times_in_db[idx].isoformat()
                                 for idx, truth in enumerate(res) if truth]
                seen_at = None
                if capture_times:
                    face_data = {
                        'full_name': full_name,
                        'capture_times': capture_times
                    }
                    add_to_json(config.strangers_json_data, face_data)
                    seen_at = ', '.join(capture_times)
                    QMessageBox.about(
                        self,
                        'Info',
                        '%s were captured %d times before at %s.' % (
                            full_name, len(capture_times), seen_at)
                    )
                self.close()
                add_user_to_model(
                    self.parent.model,
                    first_name,
                    last_name,
                    img_path,
                    seen_at
                )
            else:
                QMessageBox.about(
                    self,
                    'Error',
                    'Bad image: face not shown. Please, upload another image.'
                )
        else:
            QMessageBox.about(self, 'Forbidden!',
                              'Please, fill in all the blanks.')


class EditDialog(QDialog):

    def __init__(self, title, parent, **dial_attr):
        super(EditDialog, self).__init__()
        self.title = title
        self.parent = parent
        self.dial_attr = dial_attr

        self.buttonBox = QDialogButtonBox(
            QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        self.buttonBox.button(QDialogButtonBox.Ok).setText('Edit')
        self.buttonBox.accepted.connect(self.submit_edit)
        self.buttonBox.rejected.connect(self.reject)

        self.setWindowTitle(self.title)
        self.setWindowIcon(self.dial_attr['favicon'])

    def show(self):
        self.selected_row = [itm.data() for itm in
                             self.parent.data_view.selectedIndexes()]
        if self.selected_row:
            super(EditDialog, self).show()
            self.old_first_name = self.selected_row[0]
            self.old_last_name = self.selected_row[1]
            self.old_img_path = self.selected_row[2]
            self.old_seen_at = self.selected_row[3]
            self.create_form_group_box()
            main_layout = QVBoxLayout()
            main_layout.addWidget(self.formGroupBox)
            main_layout.addWidget(self.buttonBox)
            self.setLayout(main_layout)
        else:
            QMessageBox.about(self, 'Forbidden!',
                              'Please, choose a row to edit.')

    def create_form_group_box(self):
        self.formGroupBox = QGroupBox('Edit Member')
        layout = QFormLayout()
        self.f_name_box = QLineEdit(self.old_first_name)
        layout.addRow(QLabel('First name:'), self.f_name_box)
        self.l_name_box = QLineEdit(self.old_last_name)
        layout.addRow(QLabel('Last name:'), self.l_name_box)
        self.img_url_btn = QPushButton('Upload image')
        layout.addRow(QLabel('Upload image:'), self.img_url_btn)
        self.img_url_btn.clicked.connect(self.img_upload_dialog)
        self.formGroupBox.setLayout(layout)

    @pyqtSlot()
    def img_upload_dialog(self):
        options = QFileDialog.Options()
        options |= QFileDialog.DontUseNativeDialog
        self.img_to_upload, _ = QFileDialog.getOpenFileName(
            self,
            'Upload an image',
            os.path.expanduser('~'),
            'All Files (*);;Image Files (*.jpg *.jpeg *.png *.jfif)',
            options=options
        )

    @pyqtSlot()
    def submit_edit(self):
        first_name = self.f_name_box.text()
        last_name = self.l_name_box.text()
        if first_name and last_name and self.img_to_upload:
            full_name = '%s %s' % (first_name, last_name)
            temp_img = face_recognition.load_image_file(self.img_to_upload)
            if face_recognition.face_encodings(temp_img):
                os.remove(self.old_img_path)
                _, ext = os.path.splitext(self.img_to_upload)
                img_name = '%s_%s%s' % (
                    first_name.lower(), last_name.lower(), ext)
                img_path = os.path.join(config.member_dir, img_name)
                img = Image.open(self.img_to_upload)
                img.save(img_path)
                edit_user_in_db(
                    self.old_first_name,
                    self.old_last_name,
                    self.old_img_path,
                    first_name,
                    last_name,
                    img_path
                )
                old_capture_times = []
                if self.old_seen_at:
                    old_capture_times = self.old_seen_at.split(', ')
                old_face_data = {
                    'full_name': '%s %s' % (self.old_first_name,
                                            self.old_last_name),
                    'capture_times': old_capture_times
                }
                remove_from_json(config.strangers_json_data, old_face_data)
                QMessageBox.about(
                    self,
                    'Success',
                    '%s successfully edited as a customer.' % full_name
                )
                face = face_recognition.face_encodings(temp_img)[0]
                capture_times_in_db, face_encs_in_db = get_stranger_faces()
                res = face_recognition.compare_faces(face_encs_in_db, face)
                capture_times = [capture_times_in_db[idx].isoformat()
                                 for idx, truth in enumerate(res) if truth]
                seen_at = None
                if capture_times:
                    face_data = {
                        'full_name': full_name,
                        'capture_times': capture_times
                    }
                    add_to_json(config.strangers_json_data, face_data)
                    seen_at = ', '.join(capture_times)
                    QMessageBox.about(
                        self,
                        'Info',
                        '%s were captured %d times before at %s.' % (
                            full_name, len(capture_times), seen_at)
                    )
                self.close()
                edit_user_in_model(
                    self.parent.model,
                    self.parent.data_view.currentIndex().row(),
                    first_name,
                    last_name,
                    img_path,
                    seen_at
                )
            else:
                QMessageBox.about(self, 'Error!',
                                  'Bad image: face not shown. Please,'
                                  'upload another image.')
        else:
            QMessageBox.about(self, 'Forbidden!',
                              'Please, fill in all the blanks.')


class EditChannelDialog(QDialog):

    def __init__(self, title, parent, **dial_attr):
        super(EditChannelDialog, self).__init__()
        self.title = title
        self.parent = parent
        self.dial_attr = dial_attr
        self.create_form_group_box()

        self.buttonBox = QDialogButtonBox(
            QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        self.buttonBox.button(QDialogButtonBox.Ok).setText('Edit Channel')
        self.buttonBox.accepted.connect(self.submit_channel_edit)
        self.buttonBox.rejected.connect(self.reject)

        main_layout = QVBoxLayout()
        main_layout.addWidget(self.formGroupBox)
        main_layout.addWidget(self.buttonBox)
        self.setLayout(main_layout)

        self.setWindowTitle(self.title)
        self.setWindowIcon(self.dial_attr['favicon'])

    def create_form_group_box(self):
        with open('config.json') as conf:
            user_conf = json.load(conf)
        old_channel = user_conf['channel']
        old_uname = user_conf['username']
        old_pwd = user_conf['password']
        self.formGroupBox = QGroupBox('Edit Channel')
        layout = QFormLayout()
        self.channel_box = QLineEdit(old_channel)
        layout.addRow(QLabel('Channel:'), self.channel_box)
        self.uname_box = QLineEdit(old_uname)
        layout.addRow(QLabel('Username:'), self.uname_box)
        self.pwd_box = QLineEdit(old_pwd)
        layout.addRow(QLabel('Password:'), self.pwd_box)
        self.formGroupBox.setLayout(layout)

    @pyqtSlot()
    def submit_channel_edit(self):
        new_channel = self.channel_box.text()
        new_uname = self.uname_box.text()
        new_pwd = self.pwd_box.text()
        if new_channel:
            new_stream_url = new_channel
            if new_uname:
                sep_pos = new_stream_url.find('//') + 2
                protocol = new_stream_url[:sep_pos]
                new_stream_url = '%s%s:%s@%s' % (
                    protocol, new_uname, new_pwd, new_stream_url[sep_pos:])
            ret, frame = cv2.VideoCapture(new_stream_url)
            if ret and frame:
                new_conf = {
                    'channel': new_channel,
                    'username': new_uname,
                    'password': new_uname
                }
                with open('config.json', 'w') as conf:
                    json.dump(new_conf, conf)
                QMessageBox.about(
                    self,
                    'Success',
                    'Channel is successfully edited to %s' % new_stream_url
                )
                self.close()
            else:
                QMessageBox.about(
                    self,
                    'Error!',
                    'Could not read from the source. Please, make sure the '
                    'channel, username and password are all correct.'
                )
        else:
            QMessageBox.about(self, 'Forbidden!', 'Please, specify a channel.')


class App(QMainWindow):

    def __init__(self, title, **win_attr):
        super().__init__()
        self.title = title
        self.win_attr = win_attr
        self.initUI()

    def initUI(self):
        # Set the grid layout
        self.central_widget = QWidget(self)
        self.central_widget.resize(800, 640)
        grid = QGridLayout(self.central_widget)

        # Treeview
        self.members_data_group_box = QGroupBox('Members')
        self.data_view = QTreeView()
        self.data_view.setRootIsDecorated(False)
        self.data_view.setAlternatingRowColors(True)

        data_layout = QHBoxLayout()
        data_layout.addWidget(self.data_view)
        self.members_data_group_box.setLayout(data_layout)

        self.model = self.create_members_model(self)
        self.data_view.setModel(self.model)
        populate_model(self.model)
        grid.addWidget(self.members_data_group_box, 1, 0, 12, 8)

        # Define the children dialogs
        self.add_new_member_dialog = AddDialog(
            'Add New Member', self, favicon=QIcon('favicon.png'))
        self.edit_member_dialog = EditDialog(
            'Edit Member', self, favicon=QIcon('favicon.png'))
        self.edit_channel_dialog = EditChannelDialog(
            'Edit Channel', self, favicon=QIcon('favicon.png')
        )

        # Buttons
        add_new_member_btn = QPushButton('Add New', self.central_widget)
        add_new_member_btn.setToolTip('Add New Customer')
        grid.addWidget(add_new_member_btn, 13, 5, 1, 1)
        add_new_member_btn.clicked.connect(self.add_new_member_dialog.show)

        edit_member_btn = QPushButton('Edit', self.central_widget)
        edit_member_btn.setToolTip('Edit Customer Info')
        grid.addWidget(edit_member_btn, 13, 6, 1, 1)
        edit_member_btn.clicked.connect(self.edit_member_dialog.show)

        remove_member_btn = QPushButton('Remove', self.central_widget)
        remove_member_btn.setToolTip('Remove Customer')
        grid.addWidget(remove_member_btn, 13, 7, 1, 1)
        remove_member_btn.clicked.connect(self.remove_member)

        watch_live_stream_btn = QPushButton('Get Live Stream',
                                            self.central_widget)
        watch_live_stream_btn.setToolTip('Get Live Stream from Camera')
        grid.addWidget(watch_live_stream_btn, 15, 0, 1, 4)
        watch_live_stream_btn.clicked.connect(watch_live)

        change_ch_btn = QPushButton('Edit Channel', self.central_widget)
        change_ch_btn.setToolTip('Edit Streaming Channel')
        grid.addWidget(change_ch_btn, 16, 0, 1, 4)
        change_ch_btn.clicked.connect(self.edit_channel_dialog.show)

        # Set the window components
        self.setWindowTitle(self.title)
        self.setGeometry(
            self.win_attr['left'],
            self.win_attr['top'],
            self.win_attr['width'],
            self.win_attr['height']
        )
        self.setWindowIcon(self.win_attr['favicon'])
        # self.statusBar().showMessage('Status Bar Message.')

        # Show the window
        self.show()

    @staticmethod
    def create_members_model(parent):
        model = QStandardItemModel(0, 4, parent)
        model.setHeaderData(0, Qt.Horizontal, 'First name')
        model.setHeaderData(1, Qt.Horizontal, 'Last name')
        model.setHeaderData(2, Qt.Horizontal, 'Image URL')
        model.setHeaderData(3, Qt.Horizontal, 'Seen at')
        return model

    def remove_member(self):
        selected_row = [itm.data() for itm in self.data_view.selectedIndexes()]
        if selected_row:
            fname_rem, lname_rem, img_path_rem, seen_at_rem = selected_row
            ret = QMessageBox.question(
                self,
                'Remove',
                'Are you sure you want to remove %s %s?' % (
                    fname_rem, lname_rem),
                QMessageBox.Yes | QMessageBox.No
            )
            if ret == QMessageBox.Yes:
                capture_times = []
                if seen_at_rem:
                    capture_times = seen_at_rem.split(', ')
                face_data = {
                    'full_name': '%s %s' % (lname_rem, fname_rem),
                    'capture_times': capture_times
                }
                os.remove(img_path_rem)
                remove_user_in_db(fname_rem, lname_rem, img_path_rem)
                remove_from_json(config.strangers_json_data, face_data)
                remove_user_in_model(
                    self.model,
                    self.data_view.currentIndex().row()
                )
        else:
            QMessageBox.about(self, 'Forbidden!',
                              'Please, choose a row to remove.')


def watch_live(): pass


if __name__ == '__main__':
    app = QApplication(sys.argv)
    app.setStyle('Fusion')
    ex = App(
        'ITS FaceDetect Demo',
        left=40,
        top=40,
        width=800,
        height=640,
        favicon=QIcon('favicon.png')
    )
    sys.exit(app.exec_())