/*
 * This file is part of the Flowee project
 * Copyright (C) 2021-2025 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 <hmac_sha512.h>
#include <sha256.h>
#include <Logger.h>
#include <QFile>

#include <string_view>

namespace {
enum FindOnType {
    FindComplete,
    FindPartial
};

int findOn(const QStringList &haystack, const QString &needle, FindOnType type)
{
    int left = 0;
    int right = haystack.size() - 1;
    while (left <= right) {
        int mid = left + (right - left) / 2;
        const int comparison = haystack.at(mid).compare(needle);
        if (comparison == 0)
            return mid;  // Found the needle
        else if (comparison < 0)
            left = mid + 1;  // Search in right half
        else
            right = mid - 1;  // Search in left half
    }

    if (type == FindPartial)
        return left;
    return -1;
}

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, FindComplete);

    QMapIterator<QString, QString> iter(m_wordLists);
    while (iter.hasNext()) {
        iter.next();
        loadLanguage(iter.key());
        int rc = findOn(m_words, word, FindComplete);
        if (rc >= 0)
            return rc;
    }
    m_words.clear();
    m_languageId.clear();
    return -1;
}

QString Mnemonic::wordForIndex(int index)
{
    assert(index >= 0);
    assert(index < 2048);
    if (m_words.isEmpty())
        throw std::runtime_error("No wordlist selected");
    assert(!m_words.isEmpty());
    assert(m_words.size() > index);
    return m_words.at(index);
}

QStringList Mnemonic::completeWords(const QString &partialWord) const
{
    if (m_words.isEmpty())
        throw std::runtime_error("No wordlist selected");
    int from = findOn(m_words, partialWord, FindPartial);
    QStringList answer;
    while (from < m_words.size()) {
        if (m_words.at(from).startsWith(partialWord))
            answer.append(m_words.at(from++));
        else
            break;
    }
    return answer;
}

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;

std::vector<uint8_t> Mnemonic::toRawData(const QString &mnemonic, Mnemonic::Validity *error, int *errorIndex)
{
    auto words = mnemonic.split(' ', Qt::SkipEmptyParts);
    if (words.length() == 1)
        words = mnemonic.split(QChar(0x3000), Qt::SkipEmptyParts); // IDEOGRAPHIC SPACE
    if (words.length() < 12 || words.length() > 24 || words.length() % MnemonicWordMultiple != 0) {
        if (error)
            *error = IncorrectWordCount;
        return std::vector<uint8_t>();
    }

    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;
    int stringIndex = 0;
    for (const QString &word : std::as_const(words)) {
        const auto index = findWord(word);
        if (index == -1) {
            if (m_languageId.isEmpty()) {
                if (error)
                    *error = UnknownLanguage;
                if (errorIndex)
                    *errorIndex = 0;
            }
            else {
                if (error)
                    *error = UnknownWord;
                if (errorIndex)
                    *errorIndex = stringIndex;
            }
            return std::vector<uint8_t>();
        }
        stringIndex += 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);
            }
        }
    }
    if (errorIndex) // clear it, no error.
        *errorIndex = -1;

    if (error) {
        // calc checksum
        const uint8_t checksum = data.back(); // remember checksum
        CSHA256 hasher;
        hasher.write(data.data(), data.size() - 1);
        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));
        *error = (buf[0] & mask) != checksum ? ChecksumFailure : Valid;
    }

    return data;
}

Mnemonic::Validity Mnemonic::validateMnemonic(const QString &mnemonic, int &errorIndex, bool *maybeElectrum)
{
    if (maybeElectrum)
        *maybeElectrum = false; // clear out param, if any
    // start with some sanity checks
    if (mnemonic.isEmpty())
        return IncorrectWordCount;
    if (mnemonic.at(0).isSpace() || mnemonic.back().isSpace()) // caller should chomp
        return WhitespaceError;

    Mnemonic::Validity error;
    auto data = toRawData(mnemonic, &error, &errorIndex);
    if (data.empty())
        return error;

    if (maybeElectrum)
        *maybeElectrum = maybeValidElectrumMnemonic(mnemonic);
    return error;
}

bool Mnemonic::maybeValidElectrumMnemonic(const QString &mnemonic) const
{
    // Try the phrase versus the Electrum-style checksum check
    // The way it works is:
    // 1. Initialize an hmac sha512 with "Seed version" as key,
    // 2. then write the raw utf-8 encoded bytes of the mnemonic phrase to this hasher,
    // 3. and finally, compare the digest to the byte 0x01 (seed version == 0x01).
    constexpr std::string_view key = "Seed version";
    CHMAC_SHA512 hmac(reinterpret_cast<const unsigned char *>(key.data()), key.size());
    const auto mnemonicBytes = mnemonic.toUtf8();
    hmac.write(reinterpret_cast<const unsigned char *>(mnemonicBytes.data()), mnemonicBytes.size());
    unsigned char digest[CHMAC_SHA512::OUTPUT_SIZE];
    hmac.finalize(digest);
    // The seed version will always be 0x01, no other versions were ever created.
    constexpr unsigned char seedVersion = 0x01u;
    return digest[0] == seedVersion;
}

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', Qt::SkipEmptyParts);
            assert(lexicon.size() == 2048); // deployment error
            if (lexicon.size() != 2048)
                throw std::runtime_error("Invalid mnemonics lexicon");
        }
        else {
            logWarning(Log::Flowee) << "Mnemonic: No lexicon found at" << m_wordLists.value(langId) << "id:" << langId;
            throw std::runtime_error("Missing mnemonics lexicon");
        }
    }
    else {
        logWarning().nospace() << "Unknown language ID requested: '" << langId << "'. Available are: " << m_wordLists.keys();
        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(" ");
}

void Mnemonic::loadLanguage(const QString &id)
{
    auto dict = m_wordLists.find(id);
    if (dict == m_wordLists.end()) {
        m_words.clear();
        m_languageId = QString();
        return;
    }

    QFile file(dict.value());
    if (!file.open(QIODevice::ReadOnly))
        throw std::runtime_error("Missing mnemonics lexicon");
    auto txt = QString::fromUtf8(file.readAll());
    auto words = txt.split('\n', Qt::SkipEmptyParts);
    assert(words.size() == 2048); // deployment error
    if (words.size() != 2048)
        throw std::runtime_error("Invalid mnemonics lexicon");

    m_words = words;
    m_languageId = id;
}


QString Mnemonic::languageId() const
{
    return m_languageId;
}