/*
 * SPDX-FileCopyrightText: 2017-2017 CSSlayer <wengxt@gmail.com>
 *
 * SPDX-License-Identifier: LGPL-2.1-or-later
 */

#include <array>
#include <cstddef>
#include <iostream>
#include <memory>
#include <ostream>
#include <sstream>
#include <string>
#include <string_view>
#include <unordered_set>
#include <vector>
#include <fcitx-utils/log.h>
#include "libime/core/historybigram.h"
#include "libime/core/lattice.h"
#include "libime/core/userlanguagemodel.h"
#include "libime/pinyin/pinyincontext.h"
#include "libime/pinyin/pinyindecoder.h"
#include "libime/pinyin/pinyindictionary.h"
#include "libime/pinyin/pinyinencoder.h"
#include "libime/pinyin/pinyinime.h"
#include "testdir.h"

using namespace libime;

namespace {

void checkCandidateSet(const PinyinContext &context) {
    std::unordered_set<std::string> candidates;
    for (const auto &candidate : context.candidates()) {
        candidates.insert(candidate.toString());
    }

    FCITX_ASSERT(candidates == context.candidateSet());
}

void checkCandidatesToCursorSet(const PinyinContext &context) {
    std::unordered_set<std::string> candidates;
    for (const auto &candidate : context.candidatesToCursor()) {
        candidates.insert(candidate.toString());
    }

    FCITX_ASSERT(candidates == context.candidatesToCursorSet());
}

} // namespace

int main() {
    PinyinIME ime(
        std::make_unique<PinyinDictionary>(),
        std::make_unique<UserLanguageModel>(LIBIME_BINARY_DIR "/data/sc.lm"));
    ime.model()->history().add({"字迹", "格子"});
    // add a manual dict
    std::stringstream ss;
    ss << "献世 xian'shi 0.0\n";
    ime.dict()->load(PinyinDictionary::SystemDict,
                     LIBIME_BINARY_DIR "/data/sc.dict",
                     PinyinDictFormat::Binary);
    ime.dict()->load(PinyinDictionary::UserDict, ss, PinyinDictFormat::Text);
    ime.dict()->addWord(1, "zi'ji'ge'zi", "自机各自");
    ime.setFuzzyFlags(PinyinFuzzyFlag::Inner);
    PinyinContext c(&ime);

    {
        c.type("xian");
        std::unordered_set<size_t> xiEndIndexes;
        for (const auto &candidate : c.candidates()) {
            if (candidate.toString() == "洗") {
                xiEndIndexes.insert(candidate.sentence().back()->to()->index());
            }
        }
        FCITX_ASSERT(xiEndIndexes.size() == 2) << xiEndIndexes.size();
        FCITX_ASSERT(c.candidateSet().count("洗") == 1);
        c.clear();
    }

    c.type("xianshi");

    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }
    std::cout << "--------------------------------" << std::endl;
    c.select(40);
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }
    std::cout << "--------------------------------" << std::endl;
    c.cancel();
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }
    std::cout << "--------------------------------" << std::endl;
    c.type("shi'");
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }
    std::cout << "--------------------------------" << std::endl;
    int i = 0;
    for (const auto &candidate : c.candidates()) {
        if (candidate.toString() == "西安市") {
            break;
        }
        i++;
    }
    c.select(i);
    FCITX_ASSERT(!c.selected());
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }
    std::cout << "--------------------------------" << std::endl;
    c.select(0);
    FCITX_ASSERT(c.selected()) << c.sentence() << " " << c.preedit();
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }
    std::cout << "--------------------------------" << std::endl;
    c.clear();
    FCITX_ASSERT(!c.selected());
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }
    std::cout << "--------------------------------" << std::endl;
    c.type("zi'ji'ge'zi'");
    FCITX_ASSERT(!c.selected());
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }
    std::cout << "--------------------------------" << std::endl;
    i = 0;
    for (const auto &candidate : c.candidates()) {
        if (candidate.toString() == "子集") {
            break;
        }
        i++;
    }
    c.select(i);
    FCITX_ASSERT(!c.selected());
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    for (const auto &candidate : c.candidates()) {
        std::cout << candidate.toString() << std::endl;
    }

    std::cout << "--------------------------------" << std::endl;
    c.select(0);
    FCITX_ASSERT(c.selected()) << c.sentence() << " " << c.preedit();
    std::cout << c.sentence() << std::endl;
    std::cout << c.preedit() << std::endl;
    c.clear();
    c.type("n");
    for (const auto &candidate : c.candidates()) {
        for (const auto *node : candidate.sentence()) {
            const auto &pinyin = node->as<PinyinLatticeNode>().encodedPinyin();
            std::cout << node->word();
            if (!pinyin.empty()) {
                std::cout << " " << PinyinEncoder::decodeFullPinyin(pinyin);
            }
        }
        std::cout << std::endl;
    }

    std::cout << "--------------------------------" << std::endl;
    {
        c.clear();
        c.type("nvedai");
        std::array<size_t, 7> expectedCursor = {0, 1, 3, 4, 6, 7, 8};
        for (size_t i = 0; i <= c.size(); i++) {
            c.setCursor(i);
            size_t actualCursor =
                c.preeditWithCursor(PinyinPreeditMode::Pinyin).second;
            FCITX_ASSERT(actualCursor == expectedCursor[i])
                << "Raw Cursor: " << i << " Cursor: " << actualCursor
                << " Expected: " << expectedCursor[i];
        }
        FCITX_ASSERT(c.preedit(PinyinPreeditMode::Pinyin) == "nüe dai");
    }

    {
        c.clear();
        c.type("xi'an");
        std::array<size_t, 6> expectedCursor = {0, 1, 2, 4, 6, 7};
        for (size_t i = 0; i <= c.size(); i++) {
            c.setCursor(i);
            size_t actualCursor =
                c.preeditWithCursor(PinyinPreeditMode::RawText).second;
            FCITX_ASSERT(actualCursor == expectedCursor[i])
                << "Raw Cursor: " << i << " Cursor: " << actualCursor
                << " Expected: " << expectedCursor[i];
        }
        FCITX_ASSERT(c.preedit(PinyinPreeditMode::RawText) == "xi ' an");
    }

    {
        c.clear();
        c.type("nianglanghang");
        checkCandidatesToCursorSet(c);
        size_t i = 0;
        for (const auto &candidate : c.candidatesToCursor()) {
            if (candidate.toString() == "娘") {
                break;
            }
            i++;
        }
        FCITX_ASSERT(i < c.candidatesToCursor().size());
        c.selectCandidatesToCursor(i);
        i = c.size();
        while (i > 0) {
            --i;
            c.setCursor(i);
            checkCandidatesToCursorSet(c);
        }
    }

    {
        c.clear();
        c.type("hellonihao");
        c.selectCustom(5, "Hello");
        size_t i = 0;
        for (const auto &candidate : c.candidatesToCursor()) {
            if (candidate.toString() == "你") {
                break;
            }
            i++;
        }
        FCITX_ASSERT(i < c.candidatesToCursor().size());
        c.selectCandidatesToCursor(i);

        FCITX_ASSERT(!c.selected());
        c.selectCustom(3, "What");

        FCITX_ASSERT(c.selected());
        FCITX_ASSERT(c.selectedSentence() == "Hello你What");
    }

    {
        c.clear();
        c.type("shounihao");
        auto shouPinyin = PinyinEncoder::encodeFullPinyin("shou");
        c.selectCustom(4, "✋",
                       std::string_view(shouPinyin.data(), shouPinyin.size()));
        size_t i = 0;
        for (const auto &candidate : c.candidatesToCursor()) {
            if (candidate.toString() == "你好") {
                break;
            }
            i++;
        }
        FCITX_ASSERT(i < c.candidatesToCursor().size());
        c.selectCandidatesToCursor(i);

        FCITX_ASSERT(c.selected());
        c.learn();

        c.clear();
        c.type("shounihao");
        FCITX_ASSERT(c.candidatesToCursorSet().count("✋你好") > 0);
    }
    {
        c.clear();
        c.type("er");
        c.selectCustom(2, "");
        FCITX_ASSERT(c.selected());
        FCITX_ASSERT(c.selectedSentence() == "");
        FCITX_ASSERT(c.selectedWords().size() == 1);
        FCITX_ASSERT(c.selectedWords().front().empty());
        FCITX_ASSERT(c.selectedWordsWithPinyin().size() == 1);
    }

    // Check that context words can change prediction.
    {
        {
            c.clear();
            c.type("ta");
            size_t i = 0;
            for (const auto &candidate : c.candidatesToCursor()) {
                if (candidate.toString() == "她") {
                    break;
                }
                i++;
            }
            FCITX_ASSERT(i > 0) << i;
        }
        {
            c.clear();
            c.setContextWords({"我", "不", "知道"});
            FCITX_ASSERT(c.contextWords() ==
                         std::vector<std::string>{"不", "知道"});
            c.setContextWords({"谁", "他"});
            FCITX_ASSERT(c.contextWords() ==
                         std::vector<std::string>{"谁", "他"});
            c.appendContextWords({"爱"});
            FCITX_ASSERT(c.contextWords() ==
                         std::vector<std::string>{"他", "爱"});
            c.type("ta");
            size_t i = 0;
            for (const auto &candidate : c.candidatesToCursor()) {
                if (candidate.toString() == "她") {
                    break;
                }
                i++;
            }
            FCITX_ASSERT(i == 0) << i;
        }
    }

    {
        c.clear();
        c.clearContextWords();
        auto wordCandidateLimit = ime.wordCandidateLimit();
        ime.setWordCandidateLimit(1);
        c.type("ziran");
        checkCandidateSet(c);
        ime.setWordCandidateLimit(wordCandidateLimit);
    }

    {
        c.clear();
        c.clearContextWords();
        c.type("ziran");
        checkCandidateSet(c);
        c.clear();
        FCITX_ASSERT(!ime.model()->history().containsBigram("他", "爱"));
        c.type("taai");
        checkCandidateSet(c);
        checkCandidatesToCursorSet(c);
        size_t i = 0;
        for (const auto &candidate : c.candidatesToCursor()) {
            if (candidate.toString() == "他爱") {
                break;
            }
            i++;
        }
        FCITX_ASSERT(i < c.candidatesToCursor().size());
        c.selectCandidatesToCursor(i);

        FCITX_ASSERT(c.selected());
        FCITX_ASSERT(c.selectedSentence() == "他爱");
        c.learn();
        c.clear();
        FCITX_ASSERT(ime.model()->history().containsBigram("他", "爱"));
    }

    return 0;
}
