Tools/rembg/native/src/main.cpp

#include <windows.h>
#include <dwmapi.h>
#include <shellapi.h>
#include <memory>
#include <string>
#include <nlohmann/json.hpp>

#include "WebViewHost.h"
#include "SidecarProcess.h"
#include "Dialogs.h"
#include "PathUtil.h"

#pragma comment(lib, "dwmapi.lib")

using json = nlohmann::json;
using namespace rbg;

namespace {

constexpr wchar_t kWindowClass[] = L"RemoveBGStudioWindow";
constexpr wchar_t kWindowTitle[] = L"RemoveBG Studio";
constexpr UINT WM_APP_SIDECAR_EVENT = WM_APP + 1;

#ifndef DWMWA_USE_IMMERSIVE_DARK_MODE
#define DWMWA_USE_IMMERSIVE_DARK_MODE 20
#endif

struct AppState {
    HWND hwnd = nullptr;
    std::unique_ptr<WebViewHost> webView;
    std::unique_ptr<SidecarProcess> sidecar;
};

AppState* g_app = nullptr;

std::vector<FilterSpec> ImageFilters() {
    return { { L"Image Files", L"*.png;*.jpg;*.jpeg;*.webp;*.bmp;*.tiff" }, { L"All Files", L"*.*" } };
}

std::vector<FilterSpec> VideoFilters() {
    return { { L"Video Files", L"*.mp4;*.mov;*.avi;*.mkv;*.webm;*.flv" }, { L"All Files", L"*.*" } };
}

// Handles JS messages addressed to the native host directly (file/folder/
// color dialogs) rather than the Python sidecar. Replies synchronously via
// PostWebMessageAsJson since we're already on the UI thread here.
void HandleNativeCommand(AppState& app, const json& req) {
    int id = req.value("id", 0);
    std::string cmd = req.value("cmd", "");
    json params = req.value("params", json::object());

    json reply = { {"id", id} };

    auto ok = [&](json result) {
        reply["event"] = "done";
        reply["result"] = std::move(result);
    };
    auto fail = [&](const std::string& msg) {
        reply["event"] = "error";
        reply["message"] = msg;
    };

    if (cmd == "dialog-open-image") {
        auto path = ShowOpenFileDialog(app.hwnd, L"Select Image File", ImageFilters());
        ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
    } else if (cmd == "dialog-open-images") {
        auto paths = ShowOpenMultiFileDialog(app.hwnd, L"Select Images", ImageFilters());
        json arr = json::array();
        for (auto& p : paths) arr.push_back(WideToUtf8(p));
        ok(json{ {"paths", arr} });
    } else if (cmd == "dialog-open-folder") {
        std::wstring title = Utf8ToWide(params.value("title", std::string("Select Folder")));
        auto path = ShowOpenFolderDialog(app.hwnd, title);
        ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
    } else if (cmd == "dialog-open-video") {
        auto path = ShowOpenFileDialog(app.hwnd, L"Select Video File", VideoFilters());
        ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
    } else if (cmd == "dialog-save-image") {
        std::wstring suggested = Utf8ToWide(params.value("suggestedName", std::string("output.png")));
        auto path = ShowSaveFileDialog(app.hwnd, L"Save Processed Image As",
            { {L"PNG Image", L"*.png"}, {L"WebP Image", L"*.webp"}, {L"JPEG Image", L"*.jpg"} },
            L"png", suggested);
        ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
    } else if (cmd == "dialog-save-video") {
        std::wstring suggested = Utf8ToWide(params.value("suggestedName", std::string("output.mov")));
        auto path = ShowSaveFileDialog(app.hwnd, L"Save Output Video As",
            { {L"QuickTime MOV", L"*.mov"}, {L"WebM Video", L"*.webm"}, {L"MP4 Video", L"*.mp4"} },
            L"mov", suggested);
        ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
    } else if (cmd == "dialog-pick-color") {
        std::wstring initial = Utf8ToWide(params.value("initial", std::string("#00FF00")));
        auto hex = ShowColorPickerDialog(app.hwnd, initial);
        ok(hex ? json{ {"hex", WideToUtf8(*hex)} } : json{ {"hex", nullptr} });
    } else {
        fail("Unknown native command: " + cmd);
    }

    app.webView->PostMessageToJS(Utf8ToWide(reply.dump()));
}

void OnWebMessage(AppState& app, const std::wstring& jsonStr) {
    json req;
    try {
        req = json::parse(WideToUtf8(jsonStr));
    } catch (const std::exception&) {
        return;
    }

    if (req.value("target", "") == "native") {
        HandleNativeCommand(app, req);
        return;
    }

    // Everything else is a sidecar command; forward the raw JSON line as-is.
    if (app.sidecar) {
        app.sidecar->SendCommand(req.dump());
    }
}

void OnSidecarEvent(AppState& app, const std::string& jsonLine) {
    // Marshal from the SidecarProcess background thread to the UI thread;
    // WebView2's COM objects are single-threaded apartment and must only
    // be touched from the thread that created them.
    auto* payload = new std::wstring(Utf8ToWide(jsonLine));
    PostMessageW(app.hwnd, WM_APP_SIDECAR_EVENT, 0, reinterpret_cast<LPARAM>(payload));
}

LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    AppState* app = g_app;

    switch (msg) {
        case WM_SIZE:
            if (app && app->webView) app->webView->Resize();
            return 0;

        case WM_GETMINMAXINFO: {
            auto* mmi = reinterpret_cast<MINMAXINFO*>(lParam);
            mmi->ptMinTrackSize = { 760, 560 };
            return 0;
        }

        case WM_APP_SIDECAR_EVENT: {
            std::unique_ptr<std::wstring> payload(reinterpret_cast<std::wstring*>(lParam));
            if (app && app->webView && app->webView->IsReady()) {
                app->webView->PostMessageToJS(*payload);
            }
            return 0;
        }

        case WM_DESTROY:
            PostQuitMessage(0);
            return 0;
    }

    return DefWindowProcW(hwnd, msg, wParam, lParam);
}

void ApplyDarkTitleBar(HWND hwnd) {
    BOOL dark = TRUE;
    DwmSetWindowAttribute(hwnd, DWMWA_USE_IMMERSIVE_DARK_MODE, &dark, sizeof(dark));
}

} // namespace

int APIENTRY wWinMain(HINSTANCE hInstance, HINSTANCE, PWSTR, int nCmdShow) {
    CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);

    WNDCLASSEXW wc{};
    wc.cbSize = sizeof(wc);
    wc.style = CS_HREDRAW | CS_VREDRAW;
    wc.lpfnWndProc = WndProc;
    wc.hInstance = hInstance;
    wc.hCursor = LoadCursorW(nullptr, IDC_ARROW);
    wc.hbrBackground = CreateSolidBrush(RGB(0x11, 0x11, 0x14));
    wc.lpszClassName = kWindowClass;
    RegisterClassExW(&wc);

    HWND hwnd = CreateWindowExW(
        0, kWindowClass, kWindowTitle,
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT, 1180, 820,
        nullptr, nullptr, hInstance, nullptr);

    if (!hwnd) return 1;

    ApplyDarkTitleBar(hwnd);

    AppState app;
    app.hwnd = hwnd;
    g_app = &app;

    ShowWindow(hwnd, nCmdShow);
    UpdateWindow(hwnd);

    std::wstring repoRoot = FindRepoRoot();
    std::wstring pythonExe = ResolvePythonInterpreter(repoRoot);
    std::wstring sidecarScript = repoRoot + L"\\core\\sidecar.py";

    app.sidecar = std::make_unique<SidecarProcess>(
        pythonExe, sidecarScript, repoRoot,
        [&app](const std::string& line) { OnSidecarEvent(app, line); });

    if (!app.sidecar->Start()) {
        MessageBoxW(hwnd,
            L"Could not start the RemoveBG sidecar process (Python).\n"
            L"Make sure Python 3.10+ and the project dependencies are installed\n"
            L"(see README.md), then relaunch.",
            L"RemoveBG Studio", MB_ICONWARNING);
    }

    std::wstring userDataFolder = repoRoot + L"\\native\\.webview2";
    CreateDirectoryW(userDataFolder.c_str(), nullptr);

#ifdef RBG_CONSOLE
    bool devTools = true;
#else
    bool devTools = false;
#endif

    app.webView = std::make_unique<WebViewHost>(hwnd, userDataFolder, devTools);
    app.webView->SetMessageHandler([&app](const std::wstring& j) { OnWebMessage(app, j); });

    std::wstring indexPath = GetExeDir() + L"\\web\\index.html";
    app.webView->Initialize([&app, indexPath]() {
        app.webView->Navigate(L"file:///" + indexPath);
    });

    MSG msg;
    while (GetMessageW(&msg, nullptr, 0, 0)) {
        TranslateMessage(&msg);
        DispatchMessageW(&msg);
    }

    app.sidecar.reset();
    CoUninitialize();
    return (int)msg.wParam;
}