/***
   This file is part of the Tesseract plugin for 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 the Tesseract plugin for fly-ocr. If not, see
   <https://www.gnu.org/licenses/>.
***/

#include <QtDebug>
#include <QFile>
#include <QFileInfo>
#define USE_STD_NAMESPACE
#include <tesseract/baseapi.h>
#if TESSERACT_MAJOR_VERSION < 5
#include <tesseract/genericvector.h>
#endif
#undef USE_STD_NAMESPAC
#include <locale.h>
#include <cmath>
#include "tesseractengine.h"


const QHash<QString, TesseractEngine::TessLang> TesseractEngine::m_tessLangs = {
    { "afr",          { "af", QT_TR_NOOP("Afrikaans") } },
    { "ara",          { "ar", QT_TR_NOOP("Arabic") } },
    { "asm",          { "as", QT_TR_NOOP("Assamese") } },
    { "aze",          { "az", QT_TR_NOOP("Azerbaijani") } },
    { "aze_cyrl",     { "az", QT_TR_NOOP("Azerbaijani (Cyrillic)") } },
    { "bel",          { "be", QT_TR_NOOP("Belarusian") } },
    { "ben",          { "bn", QT_TR_NOOP("Bengali") } },
    { "bod",          { "bo", QT_TR_NOOP("Tibetan (Standard)") } },
    { "bos",          { "bs", QT_TR_NOOP("Bosnian") } },
    { "bre",          { "bs", QT_TR_NOOP("Breton") } },
    { "bul",          { "bg", QT_TR_NOOP("Bulgarian") } },
    { "cat",          { "ca", QT_TR_NOOP("Catalan") } },
    { "ceb",          { "",   QT_TR_NOOP("Cebuano") } },
    { "ces",          { "cs", QT_TR_NOOP("Czech") } },
    { "chi_sim",      { "zh_CN", QT_TR_NOOP("Chinese (Simplified)") } },
    { "chi_sim_vert", { "zh_CN", QT_TR_NOOP("Chinese (Simplified, Vertical)") } },
    { "chi_tra",      { "zh_TW", QT_TR_NOOP("Chinese (Traditional)") } },
    { "chi_tra_vert", { "zh_TW", QT_TR_NOOP("Chinese (Traditional, Vertical)") } },
    { "chr",          { "",   QT_TR_NOOP("Cherokee") } },
    { "cos",          { "co", QT_TR_NOOP("Corsican") } },
    { "cym",          { "cy", QT_TR_NOOP("Welsh") } },
    { "dan",          { "da", QT_TR_NOOP("Danish") } },
    { "dan_frak",     { "da", QT_TR_NOOP("Danish (Fraktur)") } },
    { "deu",          { "de", QT_TR_NOOP("German") } },
    { "deu_frak",     { "de", QT_TR_NOOP("German (Fraktur)") } },
    { "div",          { "dv", QT_TR_NOOP("Dhivehi") } },
    { "dzo",          { "dz", QT_TR_NOOP("Dzongkha") } },
    { "ell",          { "el", QT_TR_NOOP("Greek") } },
    { "eng",          { "en", QT_TR_NOOP("English") } },
    { "enm",          { "en", QT_TR_NOOP("Middle English (1100-1500)") } },
    { "epo",          { "eo", QT_TR_NOOP("Esperanto") } },
    { "equ",          { "",   QT_TR_NOOP("Math / Equations") } },
    { "est",          { "et", QT_TR_NOOP("Estonian") } },
    { "eus",          { "eu", QT_TR_NOOP("Basque") } },
    { "fas",          { "fa", QT_TR_NOOP("Persian (Farsi)") } },
    { "fao",          { "fo", QT_TR_NOOP("Faroese") } },
    { "fil",          { "",   QT_TR_NOOP("Filipino") } },
    { "fin",          { "fi", QT_TR_NOOP("Finnish") } },
    { "fra",          { "fr", QT_TR_NOOP("French") } },
    { "frk",          { "",   QT_TR_NOOP("Frankish") } },
    { "frm",          { "fr", QT_TR_NOOP("Middle French (ca. 1400-1600)") } },
    { "fry",          { "fy", QT_TR_NOOP("Frisian (West)") } },
    { "gla",          { "gd", QT_TR_NOOP("Scottish Gaelic") } },
    { "gle",          { "ga", QT_TR_NOOP("Irish") } },
    { "glg",          { "gl", QT_TR_NOOP("Galician") } },
    { "grc",          { "el", QT_TR_NOOP("Ancient Greek") } },
    { "guj",          { "gu", QT_TR_NOOP("Gujarati") } },
    { "hat",          { "ht", QT_TR_NOOP("Haitian") } },
    { "heb",          { "he", QT_TR_NOOP("Hebrew") } },
    { "hin",          { "hi", QT_TR_NOOP("Hindi") } },
    { "hrv",          { "hr", QT_TR_NOOP("Croatian") } },
    { "hun",          { "hu", QT_TR_NOOP("Hungarian") } },
    { "hye",          { "hy", QT_TR_NOOP("Armenian") } },
    { "iku",          { "iu", QT_TR_NOOP("Inuktitut") } },
    { "ind",          { "id", QT_TR_NOOP("Indonesian") } },
    { "isl",          { "is", QT_TR_NOOP("Icelandic") } },
    { "ita",          { "it", QT_TR_NOOP("Italian") } },
    { "ita_old",      { "it", QT_TR_NOOP("Italian (Old)") } },
    { "jav",          { "jv", QT_TR_NOOP("Javanese") } },
    { "jpn",          { "ja", QT_TR_NOOP("Japanese") } },
    { "jpn_vert",     { "ja", QT_TR_NOOP("Japanese (Vertical)") } },
    { "kan",          { "kn", QT_TR_NOOP("Kannada") } },
    { "kat",          { "ka", QT_TR_NOOP("Georgian") } },
    { "kat_old",      { "ka", QT_TR_NOOP("Georgian (Old)") } },
    { "kaz",          { "kk", QT_TR_NOOP("Kazakh") } },
    { "khm",          { "km", QT_TR_NOOP("Khmer") } },
    { "kir",          { "ky", QT_TR_NOOP("Kyrgyz") } },
    { "kor",          { "ko", QT_TR_NOOP("Korean") } },
    { "kor_vert",     { "ko", QT_TR_NOOP("Korean (Vertical)") } },
    { "kur",          { "ku", QT_TR_NOOP("Kurdish") } },
    { "kur_ara",      { "ku", QT_TR_NOOP("Kurdish (Arabic)") } },
    { "lao",          { "lo", QT_TR_NOOP("Lao") } },
    { "lat",          { "la", QT_TR_NOOP("Latin") } },
    { "lav",          { "lv", QT_TR_NOOP("Latvian") } },
    { "lit",          { "lt", QT_TR_NOOP("Lithuanian") } },
    { "ltz",          { "lb", QT_TR_NOOP("Luxembourgish") } },
    { "mal",          { "ml", QT_TR_NOOP("Malayalam") } },
    { "mar",          { "mr", QT_TR_NOOP("Marathi") } },
    { "mkd",          { "mk", QT_TR_NOOP("Macedonian") } },
    { "mlt",          { "mt", QT_TR_NOOP("Maltese") } },
    { "mon",          { "mn", QT_TR_NOOP("Mongolian") } },
    { "mri",          { "mi", QT_TR_NOOP("Maori") } },
    { "msa",          { "ms", QT_TR_NOOP("Malay") } },
    { "mya",          { "my", QT_TR_NOOP("Burmese") } },
    { "nep",          { "ne", QT_TR_NOOP("Nepali") } },
    { "nld",          { "nl", QT_TR_NOOP("Dutch") } },
    { "nor",          { "no", QT_TR_NOOP("Norwegian") } },
    { "oci",          { "oc", QT_TR_NOOP("Occitan") } },
    { "ori",          { "or", QT_TR_NOOP("Oriya") } },
    { "pan",          { "pa", QT_TR_NOOP("Punjabi") } },
    { "pol",          { "pl", QT_TR_NOOP("Polish") } },
    { "por",          { "pt", QT_TR_NOOP("Portuguese") } },
    { "pus",          { "ps", QT_TR_NOOP("Pashto") } },
    { "que",          { "qu", QT_TR_NOOP("Quechuan") } },
    { "ron",          { "ro", QT_TR_NOOP("Romanian") } },
    { "rus",          { "ru", QT_TR_NOOP("Russian") } },
    { "san",          { "sa", QT_TR_NOOP("Sanskrit") } },
    { "sin",          { "si", QT_TR_NOOP("Sinhala") } },
    { "slk",          { "sk", QT_TR_NOOP("Slovak") } },
    { "slk_frak",     { "sk", QT_TR_NOOP("Slovak (Fraktur)") } },
    { "slv",          { "sl", QT_TR_NOOP("Slovene") } },
    { "snd",          { "sd", QT_TR_NOOP("Sindhi") } },
    { "spa",          { "es", QT_TR_NOOP("Spanish") } },
    { "spa_old",      { "es", QT_TR_NOOP("Spanish (Old)") } },
    { "sqi",          { "sq", QT_TR_NOOP("Albanian") } },
    { "srp",          { "sr", QT_TR_NOOP("Serbian") } },
    { "srp_latn",     { "sr", QT_TR_NOOP("Serbian (Latin)") } },
    { "sun",          { "su", QT_TR_NOOP("Sundanese") } },
    { "swa",          { "sw", QT_TR_NOOP("Swahili") } },
    { "swe",          { "sv", QT_TR_NOOP("Swedish") } },
    { "syr",          { "",   QT_TR_NOOP("Syriac") } },
    { "tam",          { "ta", QT_TR_NOOP("Tamil") } },
    { "tat",          { "tt", QT_TR_NOOP("Tatar") } },
    { "tel",          { "te", QT_TR_NOOP("Telugu") } },
    { "tgk",          { "tg", QT_TR_NOOP("Tajik") } },
    { "tgl",          { "tl", QT_TR_NOOP("Tagalog") } },
    { "tha",          { "th", QT_TR_NOOP("Thai") } },
    { "tir",          { "ti", QT_TR_NOOP("Tigrinya") } },
    { "ton",          { "to", QT_TR_NOOP("Tongan") } },
    { "tur",          { "tr", QT_TR_NOOP("Turkish") } },
    { "uig",          { "ug", QT_TR_NOOP("Uyghur") } },
    { "ukr",          { "uk", QT_TR_NOOP("Ukrainian") } },
    { "urd",          { "ur", QT_TR_NOOP("Urdu") } },
    { "uzb",          { "uz", QT_TR_NOOP("Uzbek") } },
    { "uzb_cyrl",     { "uz", QT_TR_NOOP("Uzbek (Cyrillic)") } },
    { "vie",          { "vi", QT_TR_NOOP("Vietnamese") } },
    { "yid",          { "yi", QT_TR_NOOP("Yiddish") } },
    { "yor",          { "yo", QT_TR_NOOP("Yoruba") } }
};


TesseractEngine::TesseractEngine(QObject* parent) : OcrEngine(parent)
{}

TesseractEngine::~TesseractEngine()
{
    if (m_api)
        delete m_api;
}

bool TesseractEngine::initialize(const QStringList& langs, bool analyseOnly)
{
    if (m_api)
        delete m_api;

    //std::signal(SIGABRT, MainWindow::tesseractCrash);
    QByteArray current = setlocale(LC_ALL, nullptr);
    setlocale(LC_ALL, "C");
    m_api = new tesseract::TessBaseAPI();
    if (analyseOnly) {
        m_api->InitForAnalysePage();
        setlocale(LC_ALL, current.constData()); //restore locale
        return true;
    }

    int ret = m_api->Init(nullptr, langs.isEmpty() ? nullptr : qPrintable(langs.join("+")));
    setlocale(LC_ALL, current.constData());

    if (ret == -1) {
        qWarning("TesseractEngine: unable to initialize tesseract engine");
        delete m_api;
        m_api = nullptr;
        return false;
    }

    m_monitor.progress = 0;
    m_monitor.cancel_this = this;
    m_monitor.cancel = TesseractEngine::cancelCallback;
    m_canceled = false;
    return true;
}

void TesseractEngine::setPageSegmentationMode(int mode)
{
    if (m_api)
        m_api->SetPageSegMode(static_cast<tesseract::PageSegMode>(mode));
}

void TesseractEngine::setAutomaticPageSegmentation()
{
    if (m_api)
        m_api->SetPageSegMode(tesseract::PSM_AUTO_ONLY);
}

void TesseractEngine::setCharacterWhitelist(const QString& characters)
{
    if (m_api)
        m_api->SetVariable("tessedit_char_whitelist", qPrintable(characters));
}

void TesseractEngine::setCharacterBlacklist(const QString& characters)
{
    if (m_api)
        m_api->SetVariable("tessedit_char_blacklist", qPrintable(characters));
}

int TesseractEngine::pageCount() const
{
    return m_pages.count();
}

bool TesseractEngine::recognize(const QImage& image, int resolution, int page)
{
    qDebug() << Q_FUNC_INFO << resolution << page;

    m_monitor.progress = 0;
    m_canceled = false;

    if (m_api) {
        m_api->SetImage(image.bits(), image.width(), image.height(), 4, image.bytesPerLine());
        m_api->SetSourceResolution(resolution);
        if (m_api->Recognize(&m_monitor) != 0)
            return false;

        char* text = m_api->GetUTF8Text();
        m_pages[page].plainText = text;
        delete [] text;
        text = m_api->GetHOCRText(page);
        m_pages[page].hocrText = text;
        delete [] text;
        m_pages[page].recognized = true;

        return true;
    }
    return false;
}

bool TesseractEngine::isRecognized(int page)
{
    return m_pages[page].recognized;
}

QByteArray TesseractEngine::getPageData(int page, OcrEngineFactory::Format format) const
{
    if (format == OcrEngineFactory::TXT)
        return m_pages[page].plainText;
    else if (format == OcrEngineFactory::HOCR)
        return m_pages[page].hocrText;
    return QByteArray();
}

void TesseractEngine::setPageData(int page, OcrEngineFactory::Format format, const QByteArray& data)
{
    if (format == OcrEngineFactory::TXT)
        m_pages[page].plainText = data;
    else if (format == OcrEngineFactory::HOCR)
        m_pages[page].hocrText = data;
}

void TesseractEngine::cancel()
{
    m_canceled = true;
}

int TesseractEngine::progress() const
{
    return m_monitor.progress;
}

QList<QRectF> TesseractEngine::analyseLayout(const QImage& image, double* avgDeskew)
{
    if (!m_api)
        return QList<QRectF>();

    *avgDeskew = 0.0;
    QList<QRectF> rects;
    int nDeskew = 0;
    m_api->SetPageSegMode(tesseract::PSM_AUTO_ONLY);
    m_api->SetImage(image.bits(), image.width(), image.height(), 4, image.bytesPerLine());

    tesseract::PageIterator* it = m_api->AnalyseLayout();
    if (it && !it->Empty(tesseract::RIL_BLOCK)) {
        do {
            int x1, y1, x2, y2;
            tesseract::Orientation orient;
            tesseract::WritingDirection wdir;
            tesseract::TextlineOrder tlo;
            float deskew;
            it->BoundingBox(tesseract::RIL_BLOCK, &x1, &y1, &x2, &y2);
            it->Orientation(&orient, &wdir, &tlo, &deskew);
            *avgDeskew += deskew;
            ++nDeskew;
            float width = x2 - x1, height = y2 - y1;
            if (width > 10 && height > 10) {
                rects.append(QRectF(x1 - 0.5 * image.width(), y1 - 0.5 * image.height(), width, height));
            }
        } while (it->Next(tesseract::RIL_BLOCK));
    }
    delete it;
    // If a somewhat large deskew angle is detected, automatically rotate image and redetect layout,
    // unless we already attempted to rotate (to prevent endless loops)
    *avgDeskew = qRound(((*avgDeskew / nDeskew) / M_PI * 180.0) * 10.0) / 10.0;
    return rects;
}

bool TesseractEngine::exportAs(OcrEngineFactory::Format format, const QString& path)
{
    QFile file(path);
    if (!file.open(QIODevice::WriteOnly))
        return false;

    if (format == OcrEngineFactory::TXT) {
        for (auto it = m_pages.cbegin(); it != m_pages.cend(); ++it) {
            if (it.key() >= 0)
                file.write(it.value().plainText);
        }
        return true;
    }
    return false;
}

void TesseractEngine::removePage(int page)
{
    if (m_pages.remove(page) > 0) {
        //update page numbers
        QMap<int, Page> tmp;
        for (auto it = m_pages.cbegin(); it != m_pages.cend(); ++it) {
            if (it.key() > page)
                tmp.insert(it.key() - 1, it.value());
            else
                tmp.insert(it.key(), it.value());
        }
        m_pages.swap(tmp);
    }
}

QList<FlyOcr::Lang> TesseractEngine::availableLanguages() const
{
    if (!m_api)
        return QList<FlyOcr::Lang>();

#if TESSERACT_MAJOR_VERSION < 5
    GenericVector<STRING> availLanguages;
#else
    std::vector<std::string> availLanguages;
#endif
    m_api->GetAvailableLanguagesAsVector(&availLanguages);

    QStringList langs;
    for (int i = 0; i < availLanguages.size(); ++i) {
        langs.append(availLanguages[i].c_str());
    }

    std::sort(langs.begin(), langs.end(), [](const QString & s1, const QString & s2) {
        bool s1Script = s1.startsWith("script") || s1[0].isUpper();
        bool s2Script = s2.startsWith("script") || s2[0].isUpper();
        if (s1Script != s2Script) {
            return !s1Script;
        } else {
            return s1 < s2;
        }
    });

    qDebug() << Q_FUNC_INFO << "available languages:" << langs;

    QList<FlyOcr::Lang> result;
    for (const QString& lang : qAsConst(langs)) {

        if (lang == QLatin1String("osd"))
            continue;

        FlyOcr::Lang l = {
            .shortName = lang,
            .code = m_tessLangs.value(lang).code,
            .name = tr(m_tessLangs.value(lang).name)
        };
        if (l.name.isEmpty())
            l.name = l.shortName;

        result << l;
    }

    return result;
}

bool TesseractEngine::hasOsd() const
{
    if (!m_api)
        return false;

#if TESSERACT_MAJOR_VERSION < 5
    GenericVector<STRING> availLanguages;
#else
    std::vector<std::string> availLanguages;
#endif

    m_api->GetAvailableLanguagesAsVector(&availLanguages);
    for (int i = 0; i < availLanguages.size(); ++i) {

        if (availLanguages[i] == "osd")
            return true;
    }

    return false;
}

bool TesseractEngine::cancelCallback(void* instance, int)
{
    TesseractEngine* engine = static_cast<TesseractEngine*>(instance);
    return engine->m_canceled;
}
