#include "autocorrect.h"
#include <algorithm>
#include <charconv>
#include <uiautomation.h>
namespace ks
{
static bool openingDelimiter(wchar_t c, const Settings &settings)
{
    return settings.autoPrefixChars.find(c) != std::wstring::npos;
}
static bool trailingDelimiter(wchar_t c, const Settings &settings)
{
    return c == L' ' || settings.autoSuffixChars.find(c) != std::wstring::npos;
}
bool autoClosingSelection(const std::wstring &text, const Settings &settings)
{
    return !text.empty() && text.size() <= 6 &&
           std::all_of(text.begin(), text.end(), [&](wchar_t c) { return trailingDelimiter(c, settings); }) &&
           text.find_first_not_of(L' ') != text.npos;
}
std::wstring layoutLanguage(const Layout &layout)
{
    if (layout.id == L"en" || layout.id == L"uk" || layout.id == L"ru")
        return layout.id;
    switch (PRIMARYLANGID(LOWORD((uintptr_t)layout.handle)))
    {
    case LANG_ENGLISH:
        return L"en";
    case LANG_UKRAINIAN:
        return L"uk";
    case LANG_RUSSIAN:
        return L"ru";
    default:
        return {};
    }
}
bool safeAutoWord(const std::wstring &word, int minimumLength)
{
    if (minimumLength < 1 || minimumLength > 32 || word.size() < (size_t)minimumLength || word.size() > 32 ||
        !std::all_of(word.begin(), word.end(), letter))
        return false;
    auto small = lower(word);
    return word == small || (word.substr(1) == small.substr(1) && word != upper(word));
}
std::string_view AutoCorrector::Dictionary::word(size_t index) const
{
    auto start = offsets[index];
    auto end = bytes.find('\t', start);
    return std::string_view(bytes).substr(start, end - start);
}
AutoCorrector::Dictionary::Dictionary(const std::filesystem::path &path)
{
    auto raw = readFile(path, 2 * 1024 * 1024);
    bytes.assign(raw.begin(), raw.end());
    for (size_t start = 0; start < bytes.size();)
    {
        auto end = bytes.find('\n', start);
        if (end == bytes.npos)
            end = bytes.size();
        auto tab = bytes.find('\t', start);
        if (tab == bytes.npos || tab >= end || tab == start || tab - start > 160)
            throw std::runtime_error("Invalid auto-correction dictionary");
        int rank{};
        auto rankEnd = end > tab + 1 && bytes[end - 1] == '\r' ? end - 1 : end;
        auto parse = std::from_chars(bytes.data() + tab + 1, bytes.data() + rankEnd, rank);
        if (parse.ec != std::errc() || parse.ptr != bytes.data() + rankEnd || rank <= 0)
            throw std::runtime_error("Invalid dictionary rank");
        offsets.push_back((uint32_t)start);
        start = end + 1;
    }
    if (offsets.empty())
        throw std::runtime_error("Empty auto-correction dictionary");
    auto entry = [&](uint32_t offset) { return std::string_view(bytes).substr(offset, bytes.find('\t', offset) - offset); };
    auto frequency = [&](uint32_t offset)
    {
        int rank{};
        auto start = offset + entry(offset).size() + 1;
        std::from_chars(bytes.data() + start, bytes.data() + bytes.size(), rank);
        return rank;
    };
    std::sort(offsets.begin(), offsets.end(), [&](uint32_t a, uint32_t b)
              { return entry(a) == entry(b) ? frequency(a) < frequency(b) : entry(a) < entry(b); });
    offsets.erase(std::unique(offsets.begin(), offsets.end(), [&](uint32_t a, uint32_t b)
                              { return entry(a) == entry(b); }), offsets.end());
    offsets.shrink_to_fit();
    bytes.shrink_to_fit();
}
int AutoCorrector::Dictionary::rank(const std::wstring &text) const
{
    auto key = utf8(lower(text));
    size_t begin = 0, end = offsets.size();
    while (begin < end)
    {
        auto mid = begin + (end - begin) / 2;
        if (word(mid) < key)
            begin = mid + 1;
        else
            end = mid;
    }
    if (begin == offsets.size() || word(begin) != key)
        return 0;
    auto start = offsets[begin] + word(begin).size() + 1;
    int rank{};
    std::from_chars(bytes.data() + start, bytes.data() + bytes.size(), rank);
    return rank;
}
void AutoCorrector::configure(const std::filesystem::path &folder, const std::vector<std::wstring> &wanted)
{
    if (languages == wanted)
        return;
    std::map<std::wstring, Dictionary> next;
    for (auto &lang : wanted)
    {
        if (lang != L"en" && lang != L"uk" && lang != L"ru")
            throw std::runtime_error("Unsupported dictionary language");
        next.emplace(lang, Dictionary(folder / (lang + L".txt")));
    }
    dictionaries = std::move(next);
    languages = wanted;
}
void AutoCorrector::reject(std::wstring word, const Settings &settings)
{
    while (!word.empty() && openingDelimiter(word.front(), settings))
        word.erase(word.begin());
    while (!word.empty() && trailingDelimiter(word.back(), settings))
        word.pop_back();
    if (rejected.size() >= 128)
        rejected.clear();
    rejected.insert(lower(word));
}
size_t AutoCorrector::memoryBytes() const
{
    size_t bytes = 0;
    for (auto &[language, d] : dictionaries)
        bytes += d.bytes.capacity() + d.offsets.capacity() * sizeof(uint32_t);
    return bytes;
}
std::optional<AutoSuggestion> AutoCorrector::suggest(const std::wstring &typed, int source,
                                                     const std::vector<Layout> &layouts,
                                                     const Settings &settings) const
{
    if (source < 0 || source >= (int)layouts.size() || !layouts[source].convertible)
        return {};
    std::wstring word = typed, suffix, prefix;
    while (!word.empty() && openingDelimiter(word.front(), settings) && prefix.size() < 3)
    {
        prefix += word.front();
        word.erase(word.begin());
    }
    while (!word.empty() && trailingDelimiter(word.back(), settings))
    {
        suffix.insert(suffix.begin(), word.back());
        word.pop_back();
    }
    // Punctuation keys can be letters in the target layout (for example ',' -> 'б').
    // Validate letters and capitalization after conversion, not before it.
    if (settings.autoMinLength < 1 || settings.autoMinLength > 32 ||
        word.size() < (size_t)settings.autoMinLength || word.size() > 32 ||
        suffix.size() > 6 || rejected.contains(lower(word)))
        return {};
    for (auto &never : settings.autoNever)
        if (lower(never) == lower(word))
            return {};
    auto sourceLanguage = layoutLanguage(layouts[source]);
    if (!dictionaries.contains(sourceLanguage))
        return {};
    auto personal = [&](const std::wstring &language, const std::wstring &candidate)
    {
        auto found = settings.autoWords.find(language);
        return found != settings.autoWords.end() &&
            std::any_of(found->second.begin(), found->second.end(), [&](const std::wstring &known)
                        { return lower(known) == lower(candidate); });
    };
    // Valid words from either selected language are never changed.
    for (auto &[language, d] : dictionaries)
        if (d.rank(word) || personal(language, word))
            return {};
    std::optional<AutoSuggestion> candidate;
    std::optional<std::wstring> matchedText;
    auto priority = [&](int layout)
    {
        auto language = layoutLanguage(layouts[layout]);
        return std::find(settings.autoLanguages.begin(), settings.autoLanguages.end(), language) - settings.autoLanguages.begin();
    };
    for (int i = 0; i < (int)layouts.size(); ++i)
    {
        if (i == source || !layouts[i].convertible ||
            std::find(settings.cycle.begin(), settings.cycle.end(), layouts[i].id) == settings.cycle.end())
            continue;
        auto dictionary = dictionaries.find(layoutLanguage(layouts[i]));
        if (dictionary == dictionaries.end())
            continue;
        auto converted = convert(word, layouts[source], layouts[i], true);
        if (converted == word || !safeAutoWord(converted, settings.autoMinLength))
            continue;
        int rank = personal(dictionary->first, converted) ? 1 : dictionary->second.rank(converted);
        if (!rank)
            continue;
        if (matchedText && *matchedText != converted)
            return {};
        matchedText = converted;
        if ((settings.autoMaxRank == 0 || rank <= settings.autoMaxRank) &&
            (!candidate || priority(i) < priority(candidate->target)))
            candidate = AutoSuggestion{source, i, prefix + converted + suffix};
    }
    return candidate;
}
int matchAutoInput(const Run &run, const std::wstring &text, const std::vector<Layout> &layouts,
                   const Layout &physical, const Settings &settings)
{
    int full = run.match(text, layouts, physical);
    if (full >= 0 || run.overflow || text.empty())
        return full;
    // Ctrl+Shift+Left may omit opening delimiters. Only these delimiters may
    // be omitted; never accept an arbitrary matching suffix of a longer word.
    for (int i = 0; i < (int)layouts.size(); ++i)
    {
        if (!layouts[i].convertible)
            continue;
        std::wstring expected;
        for (wchar_t c : run.keys)
            expected += c == L' ' ? L' ' : layouts[i].at(physical.code(c));
        for (int removed = 0;
             removed < 3 && !expected.empty() && openingDelimiter(expected.front(), settings); ++removed)
        {
            expected.erase(expected.begin());
            if (lower(expected) == lower(text))
                return i;
        }
    }
    return -1;
}
bool autoCorrectionAllowed(HWND target, const Settings &settings)
{
    if (!target || GetForegroundWindow() != target)
        return false;
    auto process = lower(processName(target));
    for (auto &excluded : settings.autoExcluded)
    {
        auto name = lower(excluded);
        if (name.ends_with(L".exe"))
            name.resize(name.size() - 4);
        if (name == process)
            return false;
    }
    // Fail closed: do not copy from password fields or inaccessible controls.
    IUIAutomation2 *automation{};
    if (FAILED(CoCreateInstance(CLSID_CUIAutomation8, nullptr, CLSCTX_INPROC_SERVER, IID_IUIAutomation2,
                                (void **)&automation)))
        return false;
    automation->put_ConnectionTimeout(200);
    automation->put_TransactionTimeout(200);
    IUIAutomationElement *focused{};
    HRESULT hr = automation->GetFocusedElement(&focused);
    automation->Release();
    if (FAILED(hr) || !focused)
        return false;
    BOOL password = TRUE, enabled = FALSE, hasFocus = FALSE;
    CONTROLTYPEID type{};
    bool allowed = SUCCEEDED(focused->get_CurrentIsPassword(&password)) && !password &&
                   SUCCEEDED(focused->get_CurrentIsEnabled(&enabled)) && enabled &&
                   SUCCEEDED(focused->get_CurrentHasKeyboardFocus(&hasFocus)) && hasFocus &&
                   SUCCEEDED(focused->get_CurrentControlType(&type)) &&
                   (type == UIA_EditControlTypeId || type == UIA_DocumentControlTypeId);
    focused->Release();
    return allowed && GetForegroundWindow() == target;
}
} // namespace ks