#!/usr/bin/env python3
"""Apply owner-thread-destroy + pump-gate crash fix to main.cpp (LF file)."""
import sys, io

p = r"C:/Users/locpham/SonicForgeStudio/native_bridge/src/main.cpp"
s = open(p, "rb").read().decode("utf-8")
orig = s

def rep(old, new, tag):
    global s
    n = s.count(old)
    assert n == 1, f"[{tag}] count={n}"
    s = s.replace(old, new, 1)
    print(f"ok: {tag}")

# 1. globals
rep(
"""static std::map<uint32_t, std::unique_ptr<ChannelWorker>>* g_workers = nullptr;""",
"""static std::map<uint32_t, std::unique_ptr<ChannelWorker>>* g_workers = nullptr;
static std::mutex g_tidMutex;
static std::map<DWORD, ChannelWorker*> g_tidToWorker;  // worker tid -> worker (owner-thread destroy)""",
"globals")

# 2. fwd decl
rep(
"""static void post_close_gui(uint32_t ch, HWND hwnd);  // defined after ChannelWorker""",
"""static void post_close_gui(uint32_t ch, HWND hwnd);  // defined after ChannelWorker
static bool is_teardown_window(HWND hwnd);  // fwd — defined after ChannelWorker""",
"fwd")

# 3. constructor register tid
rep(
"""        th_ = std::thread([this] {
#ifdef _WIN32
            OleInitialize(nullptr);""",
"""        th_ = std::thread([this] {
#ifdef _WIN32
            {
                std::lock_guard<std::mutex> lk(g_tidMutex);
                g_tidToWorker[GetCurrentThreadId()] = this;
            }
            OleInitialize(nullptr);""",
"register tid")

# 4. pump gate
rep(
"""                for (int pumped = 0; pumped < 16; ++pumped) {
                    if (!PeekMessageW(&msg, nullptr, 0, 0, PM_REMOVE)) break;
                    TranslateMessage(&msg);""",
"""                for (int pumped = 0; pumped < 16; ++pumped) {
                    if (!PeekMessageW(&msg, nullptr, 0, 0, PM_REMOVE)) break;
                    // CRASH FIX (0xc000041d): JUCE editor child windows of a
                    // plugin DLL are owned by the FIRST worker that ran the
                    // DLL's global JUCE MessageManager - not by the channel
                    // worker doing the teardown. While another worker's LOAD /
                    // close job tears an instance down (reloading / close in
                    // flight), dispatching a message (e.g. WM_PAINT) into the
                    // plugin wndproc reads instance state being destroyed ->
                    // AV in USER32 (observed: crash on the owner worker's
                    // pump, right after the teardown job closed the view).
                    // Drop the message instead (PeekMessageW already removed
                    // it): the editor may glitch, the bridge survives.
                    if (is_teardown_window(msg.hwnd)) continue;
                    TranslateMessage(&msg);""",
"pump gate")

# 5. unregister tid
rep(
"""#ifdef _WIN32
            OleUninitialize();
#endif
        });""",
"""#ifdef _WIN32
            {
                std::lock_guard<std::mutex> lk(g_tidMutex);
                g_tidToWorker.erase(GetCurrentThreadId());
            }
            OleUninitialize();
#endif
        });""",
"unregister tid")

# 6. post_and_wait
rep(
"""    void post(std::function<void()> job) {
        {
            std::lock_guard<std::mutex> lk(mu_);
            jobs_.push_back(std::move(job));
        }
        cv_.notify_all();
    }

private:""",
"""    void post(std::function<void()> job) {
        {
            std::lock_guard<std::mutex> lk(mu_);
            jobs_.push_back(std::move(job));
        }
        cv_.notify_all();
    }
    // Run job on THIS worker and wait (5s cap) until it finished. Used to
    // serialize a window destroy with the owner worker's pump. MUST NOT be
    // called from a job running on this same worker (deadlock) - callers
    // destroy directly when owner tid == current tid.
    bool post_and_wait(std::function<void()> job) {
        std::mutex doneMx;
        std::condition_variable doneCv;
        bool done = false;
        {
            std::lock_guard<std::mutex> lk(mu_);
            jobs_.push_back([&]() {
                job();
                {
                    std::lock_guard<std::mutex> dk(doneMx);
                    done = true;
                }
                doneCv.notify_all();
            });
        }
        cv_.notify_all();
        std::unique_lock<std::mutex> dk(doneMx);
        return doneCv.wait_for(dk, std::chrono::seconds(5), [&] { return done; });
    }

private:""",
"post_and_wait")

# 7. helpers after close_in_flight
rep(
"""static bool close_in_flight(uint32_t ch) {
    std::lock_guard<std::mutex> lk(g_attachMutex);
    return g_closeInFlight[ch];
}""",
"""static bool close_in_flight(uint32_t ch) {
    std::lock_guard<std::mutex> lk(g_attachMutex);
    return g_closeInFlight[ch];
}

// True if hwnd (or any ancestor) is a plugin editor window whose channel is
// mid-teardown (reloading or close job in flight). Editor child windows
// (JUCE) have the native VST window as parent; the native window is in
// g_hwndToCh. JUCE's global message window has no such parent - walk stops
// -> false (never dropped).
static bool is_teardown_window(HWND hwnd) {
    for (HWND h = hwnd; h; h = GetParent(h)) {
        uint32_t ch = UINT32_MAX;
        {
            std::lock_guard<std::mutex> lock(g_guiMutex);
            auto it = g_hwndToCh.find(h);
            if (it != g_hwndToCh.end()) ch = it->second;
        }
        if (ch != UINT32_MAX) {
            if (close_in_flight(ch)) return true;
            return g_engine ? g_engine->isReloading(ch) : false;
        }
    }
    return false;
}

// CRASH FIX (0xc000041d): plugin editor child windows are owned by the
// thread that first ran the plugin's JUCE MessageManager (a single worker),
// NOT by the channel worker doing the teardown. Cross-thread DestroyWindow
// races the owner's message pump. Post the destroy to the OWNER worker
// (serialized with its pump - no wndproc entry can race) and wait. Fall
// back to direct destroy when the owner is the current thread or unknown
// (leak > crash). Non-trivial: destroy jobs left queued on timeout are
// harmless - IsWindow guards them, and the pump gate already stopped
// dispatching to teardown windows.
static void destroy_window_on_owner(HWND hwnd) {
    if (!hwnd || !IsWindow(hwnd)) return;
    DWORD ownerTid = GetWindowThreadProcessId(hwnd, nullptr);
    DWORD curTid = GetCurrentThreadId();
    if (ownerTid == 0 || ownerTid == curTid) {
        if (IsWindow(hwnd)) DestroyWindow(hwnd);
        return;
    }
    ChannelWorker* owner = nullptr;
    {
        std::lock_guard<std::mutex> lk(g_tidMutex);
        auto it = g_tidToWorker.find(ownerTid);
        if (it != g_tidToWorker.end()) owner = it->second;
    }
    if (!owner) {
        // Window belongs to a thread we don't control (e.g. main thread).
        // Cross-thread DestroyWindow is unsafe -> leave it (leak > crash).
        std::cerr << "[NativeBridge] destroy_window_on_owner: owner tid=" << ownerTid
                  << " not a bridge worker - leaving window " << hwnd << std::endl;
        return;
    }
    if (!owner->post_and_wait([hwnd]() {
            if (IsWindow(hwnd)) DestroyWindow(hwnd);
        }))
        std::cerr << "[NativeBridge] destroy_window_on_owner: timeout waiting owner tid="
                  << ownerTid << " to destroy " << hwnd << std::endl;
}""",
"helpers")

# 8. post_close_gui: destroy outside g_loadMutex, via owner
rep(
"""                bool reloadOk = false;
                {
                    // destroy children + reload in ONE lock: destroying the
                    // editor windows runs the plugin's window proc (DLL
                    // entry) on this worker; it must not race another
                    // thread's createInstance/createView of the same DLL.
                    // Same-thread destroy is valid (attachView created them
                    // on this worker); without destroy the editor keeps
                    // pumping (0x47b flood) while the audio loop processes
                    // the instance (2 threads in one DLL) -> Nexus USER32 crash.
                    std::lock_guard<std::mutex> lg(g_loadMutex);
                    // CRASH FIX: reload() = terminate + loadPlugin
                    // (createInstance) — Nexus may throw a C++ exception
                    // here too; catch it, treat as failed reload, and let
                    // the restore below unmute the channel (instance may be
                    // broken; next load/assign rebuilds it).
                    try {
#ifdef _WIN32
                        disable_ime_contexts(hwnd);
                        if (hwnd && IsWindow(hwnd)) {
                            while (HWND c = FindWindowExA(hwnd, nullptr, nullptr, nullptr)) {
                                std::cerr << "[dbg] closeGUI destroy child=" << c
                                          << " owner_tid=" << GetWindowThreadProcessId(c, nullptr)
                                          << " cur_tid=" << GetCurrentThreadId() << std::endl;
                                DestroyWindow(c);
                            }
                        }
#endif
                        reloadOk = i->reload();
                    } catch (...) {
                        std::cerr << "[NativeBridge] closeGUI reload EXCEPTION ch=" << ch
                                  << " — plugin threw (createInstance)" << std::endl;
                        reloadOk = false;
                    }
                }""",
"""                bool reloadOk = false;
#ifdef _WIN32
                // CRASH FIX (0xc000041d): destroy the editor's child windows
                // on the thread that OWNS them (the plugin's JUCE
                // MessageManager thread - a single worker), never
                // cross-thread. Wait for each destroy to finish (owner worker
                // serializes it with its pump) BEFORE reload() - otherwise
                // reload's createInstance enters the DLL while the owner
                // thread is still inside the plugin wndproc (2 threads in one
                // DLL -> Nexus crash). Kept OUTSIDE g_loadMutex: the wait
                // must not block other loads (the owner worker may itself be
                // waiting on that lock).
                disable_ime_contexts(hwnd);
                if (hwnd && IsWindow(hwnd)) {
                    while (HWND c = FindWindowExA(hwnd, nullptr, nullptr, nullptr)) {
                        std::cerr << "[dbg] closeGUI destroy child=" << c
                                  << " owner_tid=" << GetWindowThreadProcessId(c, nullptr)
                                  << " cur_tid=" << GetCurrentThreadId() << std::endl;
                        destroy_window_on_owner(c);
                    }
                }
#endif
                {
                    // reload() = terminate + loadPlugin (createInstance) —
                    // serialized with all other plugin-DLL entry points.
                    // Nexus may throw a C++ exception here too; catch it,
                    // treat as failed reload, and let the restore below
                    // unmute the channel (instance may be broken; next
                    // load/assign rebuilds it).
                    std::lock_guard<std::mutex> lg(g_loadMutex);
                    try {
                        reloadOk = i->reload();
                    } catch (...) {
                        std::cerr << "[NativeBridge] closeGUI reload EXCEPTION ch=" << ch
                                  << " — plugin threw (createInstance)" << std::endl;
                        reloadOk = false;
                    }
                }""",
"post_close_gui destroy")

# 9. LOAD close-editor destroy loop via owner
rep(
"""                            // Destroy the old editor's child windows (worker-owned,
                            // created by attachView here) so they stop pumping;
                            // the parent window is hidden and kept for reuse.
                            // See post_close_gui for the same pattern.
                            while (HWND c = FindWindowExA(hToHide, nullptr, nullptr, nullptr)) {
                                std::cerr << "[dbg] load close-editor destroy child=" << c
                                          << " owner_tid=" << GetWindowThreadProcessId(c, nullptr)
                                          << " cur_tid=" << GetCurrentThreadId() << std::endl;
                                DestroyWindow(c);
                            }""",
"""                            // CRASH FIX (0xc000041d): destroy the old
                            // editor's child windows on the thread that OWNS
                            // them (JUCE MessageManager worker), never
                            // cross-thread - DestroyWindow from this LOAD
                            // worker raced the owner's pump dispatching
                            // WM_PAINT into the plugin wndproc -> AV. Wait
                            // for each destroy before assign() below. The
                            // parent window is hidden and kept for reuse.
                            // See post_close_gui for the same pattern.
                            while (HWND c = FindWindowExA(hToHide, nullptr, nullptr, nullptr)) {
                                std::cerr << "[dbg] load close-editor destroy child=" << c
                                          << " owner_tid=" << GetWindowThreadProcessId(c, nullptr)
                                          << " cur_tid=" << GetCurrentThreadId() << std::endl;
                                destroy_window_on_owner(c);
                            }""",
"load destroy loop")

assert s != orig
open(p, "wb").write(s.encode("utf-8"))
print("written")