#include "platform.h"
#include <algorithm>
#include <fstream>
#include <imm.h>
#include <map>
#include <mutex>
#include <set>
#include <shlobj.h>
#include <tlhelp32.h>
#include <wincrypt.h>

namespace ks
{
namespace
{
const UINT scans[] = {0x29, 2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 16, 17, 18,
                      19,   20, 21, 22, 23, 24, 25, 26, 27, 30, 31, 32, 33, 34, 35, 36,
                      37,   38, 39, 40, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 43};
struct ClipboardLock
{
    ClipboardLock(HWND w)
    {
        for (int i = 0; i < 30; ++i)
        {
            if (OpenClipboard(w))
                return;
            Sleep(10);
        }
        throw std::runtime_error("Clipboard is busy");
    }
    ~ClipboardLock()
    {
        CloseClipboard();
    }
};
bool globalFormat(UINT f)
{
    return f != CF_BITMAP && f != CF_PALETTE && f != CF_METAFILEPICT && f != CF_ENHMETAFILE &&
           !(f >= CF_PRIVATEFIRST && f <= CF_PRIVATELAST) && !(f >= CF_GDIOBJFIRST && f <= CF_GDIOBJLAST) &&
           f != CF_OWNERDISPLAY;
}
std::optional<ULONGLONG> started(DWORD id)
{
    HANDLE h = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, id);
    if (!h)
        return {};
    FILETIME c, e, k, u;
    BOOL ok = GetProcessTimes(h, &c, &e, &k, &u);
    CloseHandle(h);
    return ok ? std::optional<ULONGLONG>((ULONGLONG(c.dwHighDateTime) << 32) | c.dwLowDateTime)
              : std::nullopt;
}
} // namespace
std::vector<Layout> readLayouts()
{
    int n = GetKeyboardLayoutList(0, nullptr);
    std::vector<HKL> handles(n);
    n = GetKeyboardLayoutList(n, handles.data());
    std::vector<Layout> out;
    for (int i = 0; i < n; ++i)
    {
        Layout l;
        l.handle = handles[i];
        wchar_t id[32];
        swprintf_s(id, L"win:%08x", (unsigned)(uintptr_t)l.handle);
        l.id = id;
        wchar_t name[128]{};
        GetLocaleInfoW(MAKELCID(LOWORD((uintptr_t)l.handle), SORT_DEFAULT), LOCALE_SNATIVELANGNAME, name,
                       128);
        l.name = name;
        if (l.name.empty())
            l.name = l.id;
        l.name += L" (" + l.id.substr(4) + L")";
        l.convertible =
            ImmGetIMEFileNameW(l.handle, nullptr, 0) == 0 && (((uintptr_t)l.handle >> 16) & 0xf000) != 0xe000;
        for (bool shift : {false, true})
            for (UINT scan : scans)
            {
                BYTE state[256]{};
                if (shift)
                    state[VK_SHIFT] = 0x80;
                wchar_t buf[8]{};
                UINT vk = MapVirtualKeyExW(scan, MAPVK_VSC_TO_VK_EX, l.handle);
                int count = ToUnicodeEx(vk, scan, state, buf, 8, 4, l.handle);
                wchar_t c =
                    count == 1 && buf[0] >= 32 && !(buf[0] >= 0xd800 && buf[0] <= 0xdfff) ? buf[0] : 0;
                (shift ? l.shifted : l.normal) += c;
            }
        out.push_back(std::move(l));
    }
    return out;
}
const Layout &physicalLayout()
{
    static const Layout l{nullptr,
                          L"physical",
                          L"Physical US",
                          L"`1234567890-=qwertyuiop[]asdfghjkl;'zxcvbnm,./\\",
                          L"~!@#$%^&*()_+QWERTYUIOP{}ASDFGHJKL:\"ZXCVBNM<>?|",
                          true};
    return l;
}
wchar_t physicalKey(DWORD scan, bool shift)
{
    for (size_t i = 0; i < std::size(scans); ++i)
        if (scans[i] == scan)
            return (shift ? physicalLayout().shifted : physicalLayout().normal)[i];
    return 0;
}
static HWND inputWindow(HWND w)
{
    if (!w || !IsWindow(w))
        return nullptr;
    GUITHREADINFO info{sizeof(info)};
    DWORD thread = w == GetForegroundWindow() ? 0 : GetWindowThreadProcessId(w, nullptr);
    if (GetGUIThreadInfo(thread, &info) && info.hwndFocus &&
        (info.hwndFocus == w || IsChild(w, info.hwndFocus) || GetAncestor(info.hwndFocus, GA_ROOT) == w))
        return info.hwndFocus;
    return w;
}
HKL windowLayout(HWND w)
{
    HWND input = inputWindow(w);
    return input ? GetKeyboardLayout(GetWindowThreadProcessId(input, nullptr)) : nullptr;
}
bool switchWindowLayout(HWND target, HKL layout, const std::function<bool()> &valid)
{
    HWND input = inputWindow(target);
    auto unchanged = [&]
    { return input && valid() && GetForegroundWindow() == target && inputWindow(target) == input; };
    if (!layout || !unchanged())
        return false;
    if (windowLayout(target) == layout)
        return true;
    // Windows directs layout requests to the focused control, which can belong
    // to a different UI thread/process than a Teams/WebView top-level window.
    if (!PostMessageW(input, WM_INPUTLANGCHANGEREQUEST, 0, (LPARAM)layout))
        return false;
    for (int i = 0; i < 30; ++i)
    {
        if (!unchanged())
            return false;
        if (windowLayout(target) == layout)
            return true;
        // Some hosts handle this at their top level instead of the focused child.
        if (i == 14 && input != target)
            PostMessageW(target, WM_INPUTLANGCHANGEREQUEST, 0, (LPARAM)layout);
        Sleep(10);
    }
    return unchanged() && windowLayout(target) == layout;
}
void showRequestedWindow(HWND window)
{
    ShowWindow(window, SW_SHOW);
    // STARTF_USESHOWWINDOW can override the process's first ShowWindow with
    // SW_HIDE. An explicit user request must still make the window visible.
    if (!IsWindowVisible(window))
        ShowWindow(window, SW_SHOW);
}
bool activateWindow(HWND window)
{
    if (!IsWindow(window))
        return false;
    if (SetForegroundWindow(window))
        return true;
    DWORD foreground = GetWindowThreadProcessId(GetForegroundWindow(), nullptr);
    DWORD current = GetCurrentThreadId();
    bool attached = foreground && foreground != current && AttachThreadInput(current, foreground, TRUE);
    BringWindowToTop(window);
    SetForegroundWindow(window);
    if (attached)
        AttachThreadInput(current, foreground, FALSE);
    return GetForegroundWindow() == window;
}
bool webViewDescendant(DWORD owner, DWORD host,
                       const std::function<std::optional<std::pair<DWORD, std::wstring>>(DWORD)> &read)
{
    if (!owner || !host)
        return false;
    std::set<DWORD> seen;
    while (owner != host)
    {
        if (!seen.insert(owner).second)
            return false;
        auto p = read(owner);
        if (!p || lower(p->second) != L"msedgewebview2.exe" || !p->first)
            return false;
        owner = p->first;
    }
    return true;
}
bool ownsClipboard(HWND window)
{
    HWND owner = GetClipboardOwner();
    if (!owner || !window)
        return false;
    DWORD a{}, b{};
    GetWindowThreadProcessId(owner, &a);
    GetWindowThreadProcessId(window, &b);
    if (!a || !b)
        return false;
    if (a == b)
        return true;
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snapshot == INVALID_HANDLE_VALUE)
        return false;
    std::map<DWORD, std::pair<DWORD, std::wstring>> entries;
    PROCESSENTRY32W e{};
    e.dwSize = sizeof(e);
    if (Process32FirstW(snapshot, &e))
        do
        {
            entries[e.th32ProcessID] = {e.th32ParentProcessID, e.szExeFile};
        } while (Process32NextW(snapshot, &e));
    CloseHandle(snapshot);
    return webViewDescendant(a, b,
                             [&](DWORD id) -> std::optional<std::pair<DWORD, std::wstring>>
                             {
                                 auto it = entries.find(id);
                                 if (it == entries.end())
                                     return {};
                                 auto child = started(id), parent = started(it->second.first);
                                 if (!child || !parent || *parent > *child)
                                     return {};
                                 return it->second;
                             });
}
std::wstring processName(HWND window)
{
    DWORD pid{};
    GetWindowThreadProcessId(window, &pid);
    HANDLE h = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid);
    if (!h)
        return {};
    wchar_t path[32768];
    DWORD n = 32768;
    BOOL ok = QueryFullProcessImageNameW(h, 0, path, &n);
    CloseHandle(h);
    return ok ? std::filesystem::path(std::wstring(path, n)).stem().wstring() : L"";
}
Clip readClipboard(HWND window, size_t limit)
{
    ClipboardLock lock(window);
    Clip c;
    size_t total = 0;
    for (UINT f = EnumClipboardFormats(0); f; f = EnumClipboardFormats(f))
    {
        if (!globalFormat(f))
            continue;
        HANDLE handle = GetClipboardData(f);
        if (!handle)
            continue;
        SIZE_T size = GlobalSize(handle);
        if (!size)
            continue;
        if (size > limit - total)
            throw std::runtime_error("Clipboard exceeds the configured size limit");
        const auto *p = static_cast<const unsigned char *>(GlobalLock(handle));
        if (!p)
            throw std::runtime_error("Cannot read clipboard format");
        Format format;
        format.id = f;
        format.data.assign(p, p + size);
        GlobalUnlock(handle);
        if (f >= 0xc000)
        {
            wchar_t name[256]{};
            GetClipboardFormatNameW(f, name, 256);
            format.name = name;
        }
        total += size;
        c.formats.push_back(std::move(format));
    }
    return c;
}
DWORD writeClipboard(HWND window, const Clip &c)
{
    // Allocate all blocks before replacing any existing data.
    std::vector<std::pair<UINT, HGLOBAL>> blocks;
    try
    {
        for (const auto &f : c.formats)
        {
            if (f.data.empty())
                continue;
            UINT id = f.name.empty() ? f.id : RegisterClipboardFormatW(f.name.c_str());
            if (!id)
                throw std::runtime_error("Cannot register clipboard format");
            HGLOBAL h = GlobalAlloc(GMEM_MOVEABLE, f.data.size());
            if (!h)
                throw std::bad_alloc();
            void *p = GlobalLock(h);
            if (!p)
            {
                GlobalFree(h);
                throw std::bad_alloc();
            }
            memcpy(p, f.data.data(), f.data.size());
            GlobalUnlock(h);
            blocks.emplace_back(id, h);
        }
        ClipboardLock lock(window);
        if (!EmptyClipboard())
            throw std::runtime_error("Cannot replace clipboard");
        for (auto &[id, h] : blocks)
        {
            if (!SetClipboardData(id, h))
                throw std::runtime_error("Cannot write clipboard format");
            h = nullptr;
        }
    }
    catch (...)
    {
        for (auto [id, h] : blocks)
            if (h)
                GlobalFree(h);
        throw;
    }
    return GetClipboardSequenceNumber();
}
Clip textClip(const std::wstring &text)
{
    Format f;
    f.id = CF_UNICODETEXT;
    f.data.resize((text.size() + 1) * sizeof(wchar_t));
    memcpy(f.data.data(), text.c_str(), f.data.size());
    return Clip{{std::move(f)}};
}
std::wstring Clip::text() const
{
    for (const auto &f : formats)
        if (f.id == CF_UNICODETEXT && f.data.size() >= 2)
        {
            std::wstring text(f.data.size() / 2, L'\0');
            memcpy(text.data(), f.data.data(), text.size() * 2);
            auto pos = text.find(L'\0');
            if (pos != text.npos)
                text.resize(pos);
            return text;
        }
    return {};
}
size_t Clip::size() const
{
    size_t n = 0;
    for (const auto &f : formats)
        n += f.data.size();
    return n;
}
ULONG_PTR inputMarker()
{
    static const char marker{};
    return reinterpret_cast<ULONG_PTR>(&marker);
}
bool ownKeyboardInput(const KBDLLHOOKSTRUCT &key)
{
    return (key.flags & LLKHF_INJECTED) && key.dwExtraInfo == inputMarker();
}
INPUT keyboardInput(WORD key, DWORD flags)
{
    INPUT input{};
    input.type = INPUT_KEYBOARD;
    input.ki.wVk = key;
    input.ki.dwFlags = flags;
    input.ki.dwExtraInfo = inputMarker();
    return input;
}
bool chord(WORD key, UINT mods)
{
    std::vector<INPUT> inputs;
    auto add = [&](WORD k, bool up)
    {
        INPUT i = keyboardInput(k, up ? KEYEVENTF_KEYUP : 0);
        if (k == VK_LEFT || k == VK_RIGHT || k == VK_INSERT)
            i.ki.dwFlags |= KEYEVENTF_EXTENDEDKEY;
        inputs.push_back(i);
    };
    for (auto [m, k] : {std::pair<UINT, WORD>{MOD_CONTROL, VK_CONTROL},
                        {MOD_ALT, VK_MENU},
                        {MOD_SHIFT, VK_SHIFT},
                        {MOD_WIN, VK_LWIN}})
        if (mods & m)
            add(k, false);
    add(key, false);
    add(key, true);
    for (auto [m, k] : {std::pair<UINT, WORD>{MOD_WIN, VK_LWIN},
                        {MOD_SHIFT, VK_SHIFT},
                        {MOD_ALT, VK_MENU},
                        {MOD_CONTROL, VK_CONTROL}})
        if (mods & m)
            add(k, true);
    return SendInput((UINT)inputs.size(), inputs.data(), sizeof(INPUT)) == inputs.size();
}
bool needsSelectionDelay(HWND target)
{
    const auto name = lower(processName(target));
    return name == L"ms-teams" || name == L"msedgewebview2";
}
bool selectPrevious(size_t count, const std::function<bool()> &valid)
{
    if (count > 100000)
        return false;
    // Teams' WebView editor applies selection changes asynchronously. A burst of
    // Shift+Left events can all use the same caret position and select only one
    // character. Give each change time to settle before sending the next one.
    const auto target = GetForegroundWindow();
    if (needsSelectionDelay(target))
    {
        for (size_t i = 0; i < count; ++i)
        {
            if (!valid() || GetForegroundWindow() != target || !chord(VK_LEFT, MOD_SHIFT))
                return false;
            Sleep(25);
        }
        return valid() && GetForegroundWindow() == target;
    }
    for (size_t offset = 0; offset < count; offset += 256)
    {
        if (!valid())
            return false;
        std::vector<INPUT> inputs;
        auto add = [&](WORD key, DWORD flags)
        {
            INPUT i = keyboardInput(key, flags);
            inputs.push_back(i);
        };
        add(VK_SHIFT, 0);
        for (size_t i = 0; i < std::min<size_t>(256, count - offset); ++i)
        {
            add(VK_LEFT, KEYEVENTF_EXTENDEDKEY);
            add(VK_LEFT, KEYEVENTF_EXTENDEDKEY | KEYEVENTF_KEYUP);
        }
        add(VK_SHIFT, KEYEVENTF_KEYUP);
        if (SendInput((UINT)inputs.size(), inputs.data(), sizeof(INPUT)) != inputs.size())
            return false;
    }
    return valid();
}
bool modifiersReleased()
{
    for (int k : {VK_CONTROL, VK_MENU, VK_SHIFT, VK_LWIN, VK_RWIN})
        if (GetAsyncKeyState(k) & 0x8000)
            return false;
    return true;
}
void releaseModifiers()
{
    std::vector<INPUT> inputs;
    for (int k : {VK_LCONTROL, VK_RCONTROL, VK_LMENU, VK_RMENU, VK_LSHIFT, VK_RSHIFT, VK_LWIN, VK_RWIN})
        if (GetAsyncKeyState(k) & 0x8000)
        {
            INPUT i = keyboardInput((WORD)k, KEYEVENTF_KEYUP);
            inputs.push_back(i);
        }
    if (!inputs.empty())
        SendInput((UINT)inputs.size(), inputs.data(), sizeof(INPUT));
}
std::wstring executable()
{
    wchar_t path[32768];
    DWORD n = GetModuleFileNameW(nullptr, path, 32768);
    if (!n || n == 32768)
        throw std::runtime_error("Cannot locate executable");
    return std::wstring(path, n);
}
std::filesystem::path dataDirectory()
{
    auto folder = std::filesystem::path(executable()).parent_path() / L"data";
    std::error_code ec;
    std::filesystem::create_directories(folder, ec);
    HANDLE probe = CreateFileW((folder / L".probe").c_str(), GENERIC_WRITE, 0, nullptr, CREATE_ALWAYS,
                               FILE_ATTRIBUTE_TEMPORARY | FILE_FLAG_DELETE_ON_CLOSE, nullptr);
    if (probe != INVALID_HANDLE_VALUE)
    {
        CloseHandle(probe);
        return folder;
    }
    PWSTR base{};
    if (FAILED(SHGetKnownFolderPath(FOLDERID_LocalAppData, 0, nullptr, &base)))
        throw std::runtime_error("Cannot locate user data folder");
    folder = std::filesystem::path(base) / L"KSwitcher";
    CoTaskMemFree(base);
    std::filesystem::create_directories(folder);
    return folder;
}
void atomicWrite(const std::filesystem::path &path, const Bytes &bytes)
{
    auto tmp = path;
    tmp += L".tmp";
    HANDLE h =
        CreateFileW(tmp.c_str(), GENERIC_WRITE, 0, nullptr, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (h == INVALID_HANDLE_VALUE)
        throw std::runtime_error("Cannot open output file");
    DWORD written{};
    BOOL ok = WriteFile(h, bytes.data(), (DWORD)bytes.size(), &written, nullptr);
    BOOL flushed = FlushFileBuffers(h);
    CloseHandle(h);
    if (!ok || written != bytes.size() || !flushed ||
        !MoveFileExW(tmp.c_str(), path.c_str(), MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH))
    {
        DeleteFileW(tmp.c_str());
        throw std::runtime_error("Cannot save file atomically");
    }
}
Bytes readFile(const std::filesystem::path &path, size_t limit)
{
    auto n = std::filesystem::file_size(path);
    if (n > limit)
        throw std::runtime_error("File exceeds size limit");
    Bytes b((size_t)n);
    std::ifstream f(path, std::ios::binary);
    if (!f.read((char *)b.data(), (std::streamsize)b.size()))
        throw std::runtime_error("Cannot read file");
    return b;
}
Bytes protect(const Bytes &bytes)
{
    DATA_BLOB in{(DWORD)bytes.size(), const_cast<BYTE *>(bytes.data())}, out{};
    if (!CryptProtectData(&in, L"KSwitcher clipboard", nullptr, nullptr, nullptr, CRYPTPROTECT_UI_FORBIDDEN,
                          &out))
        throw std::runtime_error("Cannot encrypt history");
    Bytes result(out.pbData, out.pbData + out.cbData);
    LocalFree(out.pbData);
    return result;
}
Bytes unprotect(const Bytes &bytes)
{
    DATA_BLOB in{(DWORD)bytes.size(), const_cast<BYTE *>(bytes.data())}, out{};
    if (!CryptUnprotectData(&in, nullptr, nullptr, nullptr, nullptr, CRYPTPROTECT_UI_FORBIDDEN, &out))
        throw std::runtime_error("Cannot decrypt history for this Windows account");
    Bytes result(out.pbData, out.pbData + out.cbData);
    LocalFree(out.pbData);
    return result;
}
std::wstring lastError(const std::wstring &context)
{
    return context + L" (Windows error " + std::to_wstring(GetLastError()) + L")";
}
std::wstring diagnosticText(const std::wstring &text)
{
    // Escape UTF-16 code units so invisible characters and malformed surrogate pairs
    // can be reproduced without introducing fake log lines. Bound large selections.
    std::wstring out = L"\"";
    constexpr wchar_t hex[] = L"0123456789abcdef";
    const auto count = std::min<size_t>(text.size(), 512);
    for (size_t i = 0; i < count; ++i)
    {
        auto c = static_cast<unsigned short>(text[i]);
        if (c == L'"' || c == L'\\')
        {
            out += L'\\';
            out += c;
        }
        else if (c < 32 || c >= 127)
        {
            out += L"\\u";
            for (int shift = 12; shift >= 0; shift -= 4)
                out += hex[(c >> shift) & 15];
        }
        else
            out += c;
    }
    return out + L"\" units=" + std::to_wstring(text.size()) +
           L" truncated=" + std::to_wstring(count != text.size());
}
namespace
{
std::recursive_mutex logGate;
bool loggingEnabled = false;
DiagnosticBuffer diagnosticBuffer;
}
std::wstring DiagnosticBuffer::record(const std::wstring &message, bool problem, ULONGLONG now)
{
    while (!entries.empty() && (entries.size() >= 40 || now - entries.front().first > 30000))
        entries.pop_front();
    entries.emplace_back(now, message);
    if (!problem)
        return {};
    std::wstring out = L"Diagnostic incident (preceding context, up to 30 seconds):";
    for (const auto &[at, entry] : entries)
        out += L"\n  ageMs=" + std::to_wstring(now - at) + L" " + entry;
    entries.clear();
    return out;
}
void setLoggingEnabled(bool enabled)
{
    std::lock_guard lock(logGate);
    loggingEnabled = enabled;
    if (!enabled)
        diagnosticBuffer.clear();
}
std::wstring readDiagnosticLog(bool previous)
{
    std::lock_guard lock(logGate);
    auto path = dataDirectory() / (previous ? L"kswitcher.log.1" : L"kswitcher.log");
    if (!std::filesystem::exists(path))
        return previous ? L"No previous log is available yet." : L"No log yet. Entries will appear when a problem is detected.";
    auto bytes = readFile(path, 2 * 1024 * 1024);
    auto text = wide(std::string(bytes.begin(), bytes.end()));
    std::wstring display;
    for (size_t i = 0; i < text.size(); ++i)
    {
        if (text[i] == L'\n' && (i == 0 || text[i - 1] != L'\r'))
            display += L'\r';
        display += text[i];
    }
    return display.empty() ? L"The log is empty." : display;
}
void clearDiagnosticLog(bool previous)
{
    // Serialize clearing with writes and rotation. Never delete the other log.
    std::lock_guard lock(logGate);
    auto path = dataDirectory() / (previous ? L"kswitcher.log.1" : L"kswitcher.log");
    std::filesystem::remove(path);
    if (!previous)
        diagnosticBuffer.clear();
}
void diagnosticEvent(const std::wstring &message, bool problem)
{
    std::lock_guard lock(logGate);
    if (!loggingEnabled)
        return;
    auto incident = diagnosticBuffer.record(message, problem, GetTickCount64());
    if (!incident.empty())
        log(incident);
}
void log(const std::wstring &message)
{
    std::lock_guard lock(logGate);
    if (!loggingEnabled)
        return;
    try
    {
        auto path = dataDirectory() / L"kswitcher.log";
        std::error_code ec;
        auto size = std::filesystem::file_size(path, ec);
        if (!ec && size > 256 * 1024)
        {
            auto previous = path;
            previous += L".1";
            std::filesystem::remove(previous, ec);
            ec.clear();
            std::filesystem::rename(path, previous, ec);
        }
        std::ofstream f(path, std::ios::app);
        SYSTEMTIME t;
        GetLocalTime(&t);
        f << t.wYear << '-' << t.wMonth << '-' << t.wDay << ' ' << t.wHour << ':' << t.wMinute << ':'
          << t.wSecond << ' ' << utf8(message) << '\n';
    }
    catch (...)
    {
    }
}
} // namespace ks