/***
   This file is part of fly-ocr.

   Copyright (C) 2016-2020 Sandro Mani <manisandro@gmail.com>
   Copyright (C) 2021 LLC "RusBITech-Astra" (https://astralinux.ru)

    Based on gImageReader source code by Sandro Mani under the GNU
    General Public License (GPL) version 3 or later

   Foobar is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   Foobar is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with fly-ocr. If not, see <https://www.gnu.org/licenses/>.
***/

#include <QClipboard>
#include <QDir>
#include <QFileInfo>
#include <QtSpell.hpp>
#include <algorithm>
#include <QMouseEvent>
#include <QPointer>
#include <QtDebug>
#include <fly-ocr/ocrengine.h>
#include "configsettings.h"
#include "displayer.h"
#include "mainwindow.h"
#include "outputeditor.h"
#include "recognitionmenu.h"
#include "recognizer.h"
#include "utils.h"
#include "sourcemanager.h"
#include "mainwindow_ui.h"

class Recognizer::ProgressMonitor : public MainWindow::ProgressMonitor
{
public:
    ProgressMonitor(int nPages) : MainWindow::ProgressMonitor(nPages) {}

    QPointer<OcrEngine> engine;

    void cancel() override
    {
        if (engine)
            engine->cancel();
        MainWindow::ProgressMonitor::cancel();
    }

    int getProgress() const override
    {
        QMutexLocker locker(&mMutex);
        if (engine)
            return 100.0 * ((mProgress + (engine ? engine->progress() : 0) / 100.0) / mTotal);
        else
            return 0;
    }
};

Recognizer::Recognizer(const MainWindowUi* _ui) :
    m_ui(_ui)
{
    QAction* selectedPagesAction = new QAction(tr("Selected Pages"), this);
    QAction* multiplePagesAction = new QAction(tr("Multiple Pages..."), this);
    QAction* batchModeAction = new QAction(tr("Batch mode..."), this);

    m_menuPages = new QMenu(m_ui->toolButtonRecognize);
    m_menuPages->addAction(selectedPagesAction);
    m_menuPages->addAction(multiplePagesAction);
    //m_menuPages->addAction(m_actionBatchMode); //TODO check

    QMenu* m_ocrModeMenu = m_menuPages->addMenu(tr("Output Mode"));
    QAction* plainTextAction = m_ocrModeMenu->addAction(tr("Plain Text"));
    plainTextAction->setCheckable(true);
    plainTextAction->setData(PlainText);

    QAction* formattedTextAction = m_ocrModeMenu->addAction(tr("Formatted Text"));
    formattedTextAction->setCheckable(true);
    formattedTextAction->setData(FormattedText);

    m_ocrModeActionGroup = new QActionGroup(this);
    m_ocrModeActionGroup->addAction(plainTextAction);
    m_ocrModeActionGroup->addAction(formattedTextAction);
    m_ocrModeActionGroup->setExclusive(true);

    m_pagesDialog = new QDialog(MAIN);
    m_pagesDialogUi.setupUi(m_pagesDialog);

    m_batchDialog = new QDialog(MAIN);
    m_batchDialogUi.setupUi(m_batchDialog);
    m_batchDialogUi.comboBoxExisting->addItem(tr("Overwrite existing output"), BatchOverwriteOutput);
    m_batchDialogUi.comboBoxExisting->addItem(tr("Skip processing source"), BatchSkipSource);

    m_ui->toolButtonLanguages->setMenu(MAIN->recognitionMenu());
    m_ui->toolButtonRecognize->setMenu(m_menuPages);

    connect(m_ui->toolButtonLanguages, &QToolButton::pressed,  MAIN->recognitionMenu(), &RecognitionMenu::selectLanguages);
    connect(m_ui->toolButtonRecognize, &QToolButton::clicked, this, &Recognizer::recognizeButtonClicked);
    connect(selectedPagesAction, &QAction::triggered, this, &Recognizer::recognizeSelectedPages);
    connect(multiplePagesAction, &QAction::triggered, this, &Recognizer::recognizeMultiplePages);
    connect(batchModeAction, &QAction::triggered, this, &Recognizer::recognizeBatch);
    connect(m_pagesDialogUi.lineEditPageRange, &QLineEdit::textChanged, this, &Recognizer::clearLineEditPageRangeStyle);
    connect(MAIN->recognitionMenu(), &RecognitionMenu::selectedLanguagesChanged, this, &Recognizer::destroyEngine);
    connect(m_ocrModeMenu, &QMenu::triggered, [ = ](QAction * action) {
        ConfigSettings::get<VarSetting<int>>("outputeditor")->setValue(action->data().toInt());
        emit outputModeChanged();
    });

    ADD_SETTING(ComboSetting("ocrregionstrategy", m_pagesDialogUi.comboBoxRecognitionArea, 0));
    ADD_SETTING(SwitchSetting("ocraddsourcefilename", m_pagesDialogUi.checkBoxPrependFilename));
    ADD_SETTING(SwitchSetting("ocraddsourcepage", m_pagesDialogUi.checkBoxPrependPage));
    ADD_SETTING(VarSetting<int>("outputeditor", PlainText));

    int outputMode = ConfigSettings::get<VarSetting<int>>("outputeditor")->getValue();
    for (QAction* action : m_ocrModeActionGroup->actions()) {
        if (action->data().toInt() == outputMode) {
            action->setChecked(true);
            break;
        }
    }
}

bool Recognizer::exportAs(OcrEngineFactory::Format format, const QString& path)
{
    if (m_engine)
        return m_engine->exportAs(format, path);

    return false;
}

OcrEngine* Recognizer::engine() const
{
    return m_engine;
}

void Recognizer::removePage(int page)
{
    if (m_engine)
        m_engine->removePage(page);
}

void Recognizer::enqueueFiles(const QStringList& filePaths)
{
    QList<int> pages;
    for (const QString& path : qAsConst(filePaths))
        pages << MAIN->sourceManager()->findPages(path);

    if (!pages.isEmpty()) {
        if (!MAIN->sourceManager()->selectedPage())
            MAIN->sourceManager()->selectPage(pages.first());
        if (!m_isRecognizing) {
            recognize(pages);
        } else {
            m_queuedPages << pages;
        }
    }
}

Recognizer::OutputMode Recognizer::outputMode() const
{
    for (QAction* action : m_ocrModeActionGroup->actions()) {
        if (action->isChecked())
            return static_cast<OutputMode>(action->data().toInt());
    }
    return Recognizer::PlainText;
}

void Recognizer::destroyEngine()
{
    if (m_engine) {
        delete m_engine;
        m_engine = nullptr;
        MAIN->sourceManager()->clearReconized();
    }
}

void Recognizer::setRecognizeMode(const QString& mode)
{
    m_modeLabel = mode;
    //TODO display mode
}

void Recognizer::clearLineEditPageRangeStyle()
{
    qobject_cast<QLineEdit*>(QObject::sender())->setStyleSheet(QString());
}

QList<int> Recognizer::selectPages(bool& autodetectLayout)
{
    int nPages = MAIN->displayer()->pageCount();

    m_pagesDialogUi.lineEditPageRange->setText(QString("1-%1").arg(nPages));
    m_pagesDialogUi.lineEditPageRange->setFocus();
    m_pagesDialogUi.labelRecognitionArea->setVisible(MAIN->displayer()->allowAutodetectOCRAreas());
    m_pagesDialogUi.comboBoxRecognitionArea->setVisible(MAIN->displayer()->allowAutodetectOCRAreas());
    m_pagesDialogUi.groupBoxPrepend->setVisible(MAIN->displayer()->allowAutodetectOCRAreas());

    m_pagesDialogUi.comboBoxRecognitionArea->setItemText(0, MAIN->displayer()->hasMultipleOCRAreas() ? tr("Current selection") : tr("Entire page"));

    QRegExp validateRegEx("^[\\d,\\-\\s]+$");
    QList<int> pages;
    while (m_pagesDialog->exec() == QDialog::Accepted) {
        pages.clear();
        QString text = m_pagesDialogUi.lineEditPageRange->text();
        if (validateRegEx.indexIn(text) != -1) {
            text.replace(QRegExp("\\s+"), QString());
#if (QT_VERSION >= QT_VERSION_CHECK(5, 14, 0))
            for (const QString& block : text.split(',', Qt::SkipEmptyParts)) {
                QStringList ranges = block.split('-', Qt::SkipEmptyParts);
#else
            for (const QString& block : text.split(',', QString::SkipEmptyParts)) {
                QStringList ranges = block.split('-', QString::SkipEmptyParts);
#endif
                if (ranges.size() == 1) {
                    int page = ranges[0].toInt();
                    if (page > 0 && page <= nPages) {
                        pages.append(page - 1);
                    }
                } else if (ranges.size() == 2) {
                    int start = std::max(1, ranges[0].toInt());
                    int end = std::min(nPages, ranges[1].toInt());
                    for (int page = start; page <= end; ++page) {
                        pages.append(page - 1);
                    }
                } else {
                    pages.clear();
                    break;
                }
            }
        }
        if (pages.empty()) {
            m_pagesDialogUi.lineEditPageRange->setStyleSheet("background: #FF7777; color: #FFFFFF;");
        } else {
            break;
        }
    }
    std::sort(pages.begin(), pages.end());
    autodetectLayout = m_pagesDialogUi.comboBoxRecognitionArea->isVisible() ? m_pagesDialogUi.comboBoxRecognitionArea->currentIndex() == 1 : false;
    return pages;
}

void Recognizer::recognizeButtonClicked()
{
    m_queuedPages.clear();
    if (MAIN->displayer()->pageCount() > 0)
        recognize( { MAIN->displayer()->currentPage() } );
}

void Recognizer::recognizeSelectedPages()
{
    m_queuedPages.clear();
    QList<int> pages = MAIN->sourceManager()->selectedPages();
    if (pages.isEmpty() && MAIN->displayer()->pageCount() > 0)
        pages << MAIN->displayer()->currentPage();
    recognize(pages);
}

void Recognizer::recognizeMultiplePages()
{
    m_queuedPages.clear();
    bool autodetectLayout = false;
    QList<int> pages = selectPages(autodetectLayout);
    recognize(pages, autodetectLayout);
}

bool Recognizer::setupEngine()
{
    if (!m_engine) {
        QList<FlyOcr::Lang> langs = MAIN->recognitionMenu()->selectedLanguages();
        m_engine = OcrEngine::create(this, langs);
    }
    if (m_engine) {
        m_engine->setPageSegmentationMode(MAIN->recognitionMenu()->getPageSegmentationMode());
        m_engine->setCharacterWhitelist(MAIN->recognitionMenu()->getCharacterWhitelist());
        m_engine->setCharacterBlacklist(MAIN->recognitionMenu()->getCharacterBlacklist());
    } else {
        QMessageBox::critical(MAIN, tr("Recognition errors occurred"), tr("Failed to initialize OCR backend"));
    }
    return m_engine != nullptr;
}

void Recognizer::recognize(const QList<int>& pages, bool autodetectLayout)
{
    if (!setupEngine() || pages.isEmpty())
        return;

    m_isRecognizing = true;
    bool contains = false;
    for (int i : qAsConst(pages)) {

        if (m_engine->isRecognized(i)) {
            contains = true;
            break;
        }
    }
    if (contains) {
        if (QMessageBox::No == QMessageBox::question(MAIN, tr("Source already recognized"), tr("One or more selected sources were aready recognized. Proceed anyway?"))) {
            m_isRecognizing = false;
            m_queuedPages.clear();
            return;
        }
    }

    QStringList errors;
    ProgressMonitor monitor(pages.size());
    monitor.engine = m_engine;
    MAIN->showProgress(&monitor);
    MAIN->displayer()->setBlockAutoscale(true);
    int currentPageNumber = MAIN->displayer()->currentPage();

    if (pages.contains(currentPageNumber)) {
        MAIN->outputEditor()->initialize(m_engine);
    }

    Utils::busyTask([&] {
        int npages = pages.size();
        int idx = 0;
        for (int pageNumber : pages)
        {
            //monitor.desc.progress = 0;
            ++idx;
            QMetaObject::invokeMethod(MAIN, "pushState", Qt::QueuedConnection, Q_ARG(MainWindow::State, MainWindow::State::Busy), Q_ARG(QString, tr("Recognizing page %1 (%2 of %3)").arg(pageNumber + 1).arg(idx).arg(npages)));

            PageData pageData;
            pageData.success = false;
            QMetaObject::invokeMethod(this, "setPage", Qt::BlockingQueuedConnection, Q_RETURN_ARG(PageData, pageData), Q_ARG(int, pageNumber), Q_ARG(bool, autodetectLayout));
            if (!pageData.success) {
                errors.append(tr("- Page %1: failed to render page").arg(pageNumber));
                //MAIN->outputEditor()->readError(tr("\n[Failed to recognize page %1]\n"), readSessionData);
                continue;
            }

            if (pageData.ocrAreas.count() == 1) {
                const QImage& image = pageData.ocrAreas.first();
                m_engine->recognize(image, MAIN->displayer()->currentResolution(), pageData.pageInfo.page);

            } else if (pageData.ocrAreas.count() > 1) {
                QByteArray recognizedText;
                for (const QImage& image : qAsConst(pageData.ocrAreas)) {
                    m_engine->recognize(image, MAIN->displayer()->currentResolution(), pageData.pageInfo.page);
                    recognizedText.append(m_engine->getPageData(pageData.pageInfo.page, OcrEngineFactory::TXT));
                }
                m_engine->setPageData(pageData.pageInfo.page, OcrEngineFactory::TXT, recognizedText);
            }

            if (!monitor.cancelled()) {
                QMetaObject::invokeMethod(this, [ = ]() {
                    MAIN->sourceManager()->setRecognized(pageNumber, true);
                    MAIN->outputEditor()->read(pageNumber, MAIN->sourceManager()->selectedPage().get());
                }, Qt::BlockingQueuedConnection);
            }

            QMetaObject::invokeMethod(MAIN, "popState", Qt::QueuedConnection);
            monitor.increaseProgress();
            if (monitor.cancelled()) {
                break;
            }
        }
        return true;
    }, tr("Recognizing..."));
    m_isRecognizing = false;

    MAIN->displayer()->setBlockAutoscale(false);
    MAIN->hideProgress();
    MAIN->setOutputPaneVisible(true);
    //MAIN->outputEditor()->finalizeRead(readSessionData);
    if (!errors.isEmpty()) {
        showRecognitionErrorsDialog(errors);
    }

    if (monitor.cancelled()) {
        m_queuedPages.clear();
    } else if (!m_queuedPages.isEmpty()) {
        QList<int> tmp = m_queuedPages;
        m_queuedPages.clear();
        QMetaObject::invokeMethod(this, [ = ]() { recognize(tmp); }, Qt::QueuedConnection);
    }
}

void Recognizer::recognizeImage(const QImage& image, OutputDestination dest)
{
    if (!setupEngine())
        return;

    ProgressMonitor monitor(1);
    monitor.engine = m_engine;
    MAIN->showProgress(&monitor);
    if (dest == OutputDestination::Buffer) {
        auto page = MAIN->sourceManager()->selectedPage();
        int currentPageNumber = MAIN->displayer()->currentPage();
        MAIN->outputEditor()->initialize(m_engine);

        Utils::busyTask([&] {
            m_engine->recognize(image, MAIN->displayer()->currentResolution(), currentPageNumber);
            return true;
        }, tr("Recognizing..."));

        if (!monitor.cancelled()) {
            MAIN->outputEditor()->read(currentPageNumber, page.get());
            MAIN->sourceManager()->setRecognized(currentPageNumber, true);
            MAIN->setOutputPaneVisible(true);
        }

    } else if (dest == OutputDestination::Clipboard) {
        QByteArray output;
        if (Utils::busyTask([&] {
        m_engine->recognize(image, MAIN->displayer()->currentResolution(), -1);
            if (!monitor.cancelled())
            {
                output = m_engine->getPageData(-1, OcrEngineFactory::TXT);
                return true;
            }
            return false;
        }, tr("Recognizing..."))) {
            QApplication::clipboard()->setText(QString::fromUtf8(output));
        }
    }
    MAIN->hideProgress();
}

void Recognizer::recognizeBatch()
{
    m_batchDialogUi.checkBoxPrependPage->setVisible(MAIN->displayer()->allowAutodetectOCRAreas());
    m_batchDialogUi.checkBoxAutolayout->setVisible(MAIN->displayer()->allowAutodetectOCRAreas());
    if (m_batchDialog->exec() != QDialog::Accepted) {
        return;
    }
    BatchExistingBehaviour existingBehaviour = static_cast<BatchExistingBehaviour>(m_batchDialogUi.comboBoxExisting->currentData().toInt());
    bool prependPage = MAIN->displayer()->allowAutodetectOCRAreas() && m_batchDialogUi.checkBoxPrependPage->isChecked();
    bool autolayout = MAIN->displayer()->allowAutodetectOCRAreas() && m_batchDialogUi.checkBoxAutolayout->isChecked();
    int nPages = MAIN->displayer()->pageCount();

    if (!setupEngine())
        return;

    QMap<QString, QVariant> batchOptions;
    batchOptions["prependPage"] = prependPage;
    OutputEditor::BatchProcessor* batchProcessor = MAIN->outputEditor()->createBatchProcessor(batchOptions);

    QStringList errors;
    ProgressMonitor monitor(nPages);
    monitor.engine = m_engine;
    MAIN->showProgress(&monitor);
    MAIN->displayer()->setBlockAutoscale(true);
    Utils::busyTask([&] {
        int idx = 0;
        QString currFilename;
        QFile outputFile;
        for (int page = 1; page <= nPages; ++page)
        {
            //monitor.desc.progress = 0;
            ++idx;
            QMetaObject::invokeMethod(MAIN, "pushState", Qt::QueuedConnection, Q_ARG(MainWindow::State, MainWindow::State::Busy), Q_ARG(QString, tr("Recognizing page %1 (%2 of %3)").arg(page + 1).arg(idx).arg(nPages)));

            PageData pageData;
            pageData.success = false;
            QMetaObject::invokeMethod(this, "setPage", Qt::BlockingQueuedConnection, Q_RETURN_ARG(PageData, pageData), Q_ARG(int, page), Q_ARG(bool, autolayout));
            if (!pageData.success) {
                errors.append(tr("- %1:%2: failed to render page").arg(QFileInfo(pageData.pageInfo.filename).fileName()).arg(page));
                continue;
            }
            if (pageData.pageInfo.filename != currFilename) {
                if (outputFile.isOpen()) {
                    batchProcessor->writeFooter(&outputFile);
                    outputFile.close();
                }
                currFilename = pageData.pageInfo.filename;
                QFileInfo finfo(pageData.pageInfo.filename);
                QString fileName = QDir(finfo.absolutePath()).absoluteFilePath(finfo.baseName() + batchProcessor->fileSuffix());
                bool exists = QFileInfo(fileName).exists();
                if (exists && existingBehaviour == BatchSkipSource) {
                    errors.append(tr("- %1: output already exists, skipping").arg(finfo.fileName()));
                } else {
                    outputFile.setFileName(fileName);
                    if (!outputFile.open(QIODevice::WriteOnly)) {
                        errors.append(tr("- %1: failed to create output file").arg(finfo.fileName()).arg(page));
                    } else {
                        batchProcessor->writeHeader(&outputFile, pageData.pageInfo);
                    }
                }
            }
            if (outputFile.isOpen()) {
                bool firstChunk = true;
                for (const QImage& image : pageData.ocrAreas) {
                    m_engine->recognize(image, MAIN->displayer()->currentResolution(), -1);
                    if (!monitor.cancelled()) {
                        batchProcessor->appendOutput(&outputFile, m_engine, pageData.pageInfo, firstChunk);
                    }
                    firstChunk = false;
                }
            }
            QMetaObject::invokeMethod(MAIN, "popState", Qt::QueuedConnection);
            monitor.increaseProgress();
            if (monitor.cancelled()) {
                break;
            }
        }
        if (outputFile.isOpen())
        {
            batchProcessor->writeFooter(&outputFile);
            outputFile.close();
        }
        return true;
    }, tr("Recognizing..."));
    MAIN->displayer()->setBlockAutoscale(false);
    MAIN->hideProgress();
    if (!errors.isEmpty()) {
        showRecognitionErrorsDialog(errors);
    }
    delete batchProcessor;
}

Recognizer::PageData Recognizer::setPage(int page, bool autodetectLayout)
{
    PageData pageData;
    pageData.success = MAIN->displayer()->setup(&page);
    MAIN->displayer()->setPage(MAIN->sourceManager()->page(page));
    if (pageData.success) {
        if (autodetectLayout) {
            MAIN->displayer()->autodetectOCRAreas();
        }
        pageData.pageInfo.filename = MAIN->displayer()->currentImage();
        pageData.pageInfo.angle = MAIN->displayer()->currentAngle();
        pageData.pageInfo.resolution = MAIN->displayer()->currentResolution();
        pageData.pageInfo.page = page;
        pageData.ocrAreas = MAIN->displayer()->getOCRAreas();
    }
    return pageData;
}

void Recognizer::showRecognitionErrorsDialog(const QStringList& errors)
{
    Utils::messageBox(MAIN, tr("Recognition errors occurred"), tr("The following errors occurred:"), errors.join("\n"), QMessageBox::Warning, QDialogButtonBox::Close);
}
