/*
 * This file is part of the Flowee project
 * Copyright (C) 2021 Tom Zander <tom@flowee.org>
 *
 * This program 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.
 *
 * This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "Mnemonic.h"

#include <sha256.h>
#include <Logger.h>
#include <QFile>

namespace {
int findOn(const QStringList &haystack, const QString &needle)
{
    // TODO replace with binary search
    return haystack.indexOf(needle);
}

inline uint8_t bip39Shift(uint32_t bit)
{
    return (1 << (8 - (bit % 8) - 1));
}
}

int Mnemonic::findWord(const QString &word)
{
    if (!m_words.isEmpty())
        return findOn(m_words, word);

    QMapIterator<QString, QString> iter(m_wordLists);
    while (iter.hasNext()) {
        iter.next();
        QFile file(iter.value());
        if (file.open(QIODevice::ReadOnly)) {
            auto txt = QString::fromUtf8(file.readAll());
            auto words = txt.split('\n', QString::SkipEmptyParts);
            assert(words.size() == 2048); // deployment error
            if (words.size() != 2048)
                throw std::runtime_error("Invalid mnemonics lexicon");

            int rc = findOn(words, word);
            if (rc >= 0) {
                m_words = words;
                m_languageId = iter.key();
                return rc;
            }
        }
        else {
            logWarning(Log::Bitcoin) << "Mnemonic: No lexicon found at" << iter.value() << "id:" << iter.key();
        }
    }
    return -1;
}

void Mnemonic::clearSelectedLanguage()
{
    m_languageId.clear();
    m_words.clear();
}

void Mnemonic::registerWordList(const QString &id, const QString &filename)
{
    m_wordLists.insert(id, filename);
}

static constexpr int BitsPerWord = 11;
static constexpr int EntropyBitDivisor = 32;
static constexpr int MnemonicWordMultiple = 3;

Mnemonic::Validity Mnemonic::validateMnemonic(const QString &mnemonic, int &errorIndex)
{
    // start with some sanity checks
    if (mnemonic.isEmpty())
        return IncorrectWordCount;
    if (mnemonic.at(0).isSpace() || mnemonic.back().isSpace()) // caller should chomp
        return WhitespaceError;
    // can't verify until all words are present.
    auto words = mnemonic.split(' ', QString::SkipEmptyParts);
    if (words.length() == 1)
        words = mnemonic.split(QChar(0x3000), QString::SkipEmptyParts); // IDEOGRAPHIC SPACE
    if (words.length() < 12 || words.length() > 24 || words.length() % MnemonicWordMultiple != 0)
        return IncorrectWordCount;

    const uint32_t totalBits = BitsPerWord * words.length();
    const uint32_t checkBits = totalBits / (EntropyBitDivisor + 1);
    const uint32_t entropyBits = totalBits - checkBits;
    assert(entropyBits % 8 == 0);

    std::vector<uint8_t> data(entropyBits / 8 + 1); // the +1 fits the checksum
    uint32_t bit = 0;
    errorIndex = 0;
    for (const QString &word : qAsConst(words)) {
        const auto index = findWord(word);
        if (index == -1) {
            if (m_languageId.isEmpty())
                return UnknownLanguage;
            return UnknownWord;
        }
        errorIndex += word.length() + 1;
        assert(index >= 0);
        uint32_t wordPos = static_cast<uint32_t>(index);

        for (size_t loop = 0; loop < BitsPerWord; ++loop, ++bit) {
            if (wordPos & (1 << (BitsPerWord - loop - 1))) {
                assert(bit <= totalBits); // don't extend beyond range of vector.
                data[bit / 8] |= bip39Shift(bit);
            }
        }
    }
    errorIndex = -1;

    // calc checksum
    const uint8_t checksum = data.back(); // remember checksum
    data.resize(data.size() - 1); // cut off the checksum bits (between 3 and 8 bits)
    CSHA256 hasher;
    hasher.Write(data.data(), data.size());
    unsigned char buf[CSHA256::OUTPUT_SIZE];
    hasher.Finalize(buf);
    uint8_t mask = 0;
    for (uint32_t i = 0; i < checkBits; ++i) { // build mask by adding one bit at a time, high bit first
        mask = mask >> 1;
        mask += 0x80;
    }
    assert(checksum == (checksum & mask));
    if ((buf[0] & mask) != checksum)
        return ChecksumFailure;
    return Valid;
}

QString Mnemonic::generateMnemonic(const std::vector<uint8_t> &entropy, const QString &langId) const
{
    if ((entropy.size() % 4) != 0)
        throw std::runtime_error("Entropy sizing invalid; multiple of 4 bytes expected");
    QStringList lexicon;
    if (m_wordLists.contains(langId)) {
        QFile file(m_wordLists[langId]);
        if (file.open(QIODevice::ReadOnly)) {
            auto txt = QString::fromUtf8(file.readAll());
            lexicon= txt.split('\n', QString::SkipEmptyParts);
            assert(lexicon.size() == 2048); // deployment error
            if (lexicon.size() != 2048)
                throw std::runtime_error("Invalid mnemonics lexicon");
        }
        else {
            logWarning(Log::Bitcoin) << "Mnemonic: No lexicon found at" << m_wordLists.value(langId) << "id:" << langId;
            throw std::runtime_error("Missing mnemonics lexicon");
        }
    }
    else {
        throw std::runtime_error("Unknown language Id");
    }

    const int entropyBits = (entropy.size() * 8);
    const int checkBits = (entropyBits / EntropyBitDivisor);
    const int totalBits = (entropyBits + checkBits);
    const int wordCount = (totalBits / BitsPerWord);

    assert((totalBits % BitsPerWord) == 0);
    assert((wordCount % MnemonicWordMultiple) == 0);

    CSHA256 hasher;
    hasher.Write(entropy.data(), entropy.size());
    std::vector<uint8_t> data(entropy);
    data.resize(entropy.size() + CSHA256::OUTPUT_SIZE); // make space
    hasher.Finalize(&data[entropy.size()]); // append

    size_t bit = 0;
    QStringList words;
    for (int word = 0; word < wordCount; word++) {
        size_t position = 0;
        for (size_t loop = 0; loop < BitsPerWord; loop++) {
            bit = (word * BitsPerWord + loop);
            position <<= 1;

            const auto byte = bit / 8;

            if ((data[byte] & bip39Shift(bit)) > 0)
                position++;
        }

        assert(position < (size_t) lexicon.size());
        words.push_back(lexicon[position]);
    }

    assert(words.size() == ((bit + 1) / BitsPerWord));
    return words.join(" ");
}