#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <windows.h>
#include <objbase.h>
#include <cstdio>
#include <cstring>
#include <stdexcept>
#include <string>
#include <vector>

// These names deliberately remain visible in the private PDB for WinDbg.
HANDLE gOutputWriter = INVALID_HANDLE_VALUE;
HANDLE gReadyEvent = nullptr;
HANDLE gCloseWriterEvent = nullptr;
HANDLE gWriterClosedEvent = nullptr;
HANDLE gStopEvent = nullptr;
DWORD gHelperPid = 0;
bool gWriterWasInherited = false;

namespace {
constexpr DWORD kReadyTimeoutMs = 15000;
constexpr ULONGLONG kHelperLifetimeMs = 15ull * 60ull * 1000ull;

class Handle {
public:
    explicit Handle(HANDLE value = nullptr) noexcept : value_(value) {}
    ~Handle() { if (valid()) CloseHandle(value_); }
    Handle(const Handle&) = delete;
    Handle& operator=(const Handle&) = delete;
    bool valid() const noexcept { return value_ && value_ != INVALID_HANDLE_VALUE; }
    HANDLE get() const noexcept { return value_; }
private:
    HANDLE value_;
};

class StopHelperOnFailure {
public:
    StopHelperOnFailure(HANDLE stopEvent, HANDLE process) noexcept : stopEvent_(stopEvent), process_(process) {}
    ~StopHelperOnFailure() {
        if (!completed_) {
            // Preserve the original exception. Cleanup is a bounded, graceful
            // request to the helper created by this call, never a forced kill.
            SetEvent(stopEvent_);
            WaitForSingleObject(process_, kReadyTimeoutMs);
        }
    }
    void Complete() noexcept { completed_ = true; }
    StopHelperOnFailure(const StopHelperOnFailure&) = delete;
    StopHelperOnFailure& operator=(const StopHelperOnFailure&) = delete;
private:
    HANDLE stopEvent_;
    HANDLE process_;
    bool completed_ = false;
};

[[noreturn]] void Fail(const char* operation, DWORD error = GetLastError()) {
    throw std::runtime_error(std::string(operation) + " failed; Win32 error " + std::to_string(error));
}

std::string Utf8(const std::wstring& value) {
    if (value.empty()) return {};
    const int size = WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS,
        value.c_str(), static_cast<int>(value.size()), nullptr, 0, nullptr, nullptr);
    if (!size) Fail("WideCharToMultiByte");
    std::string result(static_cast<size_t>(size), '\0');
    if (!WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, value.c_str(),
        static_cast<int>(value.size()), result.data(), size, nullptr, nullptr)) Fail("WideCharToMultiByte");
    return result;
}

std::string Json(const std::wstring& value) {
    const std::string bytes = Utf8(value);
    std::string result = "\"";
    for (const unsigned char c : bytes) {
        if (c == '"' || c == '\\') { result += '\\'; result += static_cast<char>(c); }
        else if (c < 0x20) {
            char escaped[7]{};
            sprintf_s(escaped, "\\u%04x", static_cast<unsigned int>(c));
            result += escaped;
        } else result += static_cast<char>(c);
    }
    return result + "\"";
}

std::wstring Hex(HANDLE value) {
    wchar_t buffer[32]{};
    swprintf_s(buffer, L"0x%llx", static_cast<unsigned long long>(reinterpret_cast<ULONG_PTR>(value)));
    return buffer;
}

std::wstring ExecutablePath() {
    std::vector<wchar_t> buffer(32768);
    const DWORD length = GetModuleFileNameW(nullptr, buffer.data(), static_cast<DWORD>(buffer.size()));
    if (!length || length >= buffer.size()) Fail("GetModuleFileNameW");
    return std::wstring(buffer.data(), length);
}

// Quote an argument according to the MSVC command-line parsing rules.
std::wstring Quote(const std::wstring& value) {
    std::wstring result = L"\"";
    size_t slashes = 0;
    for (const wchar_t c : value) {
        if (c == L'\\') { ++slashes; continue; }
        if (c == L'"') {
            result.append(slashes * 2 + 1, L'\\');
            result += c;
        } else {
            result.append(slashes, L'\\');
            result += c;
        }
        slashes = 0;
    }
    result.append(slashes * 2, L'\\');
    return result + L"\"";
}

void ValidateRunId(const std::wstring& runId) {
    GUID value{};
    const std::wstring braced = L"{" + runId + L"}";
    if (runId.size() != 36 || FAILED(CLSIDFromString(braced.c_str(), &value)))
        throw std::runtime_error("The run ID must be a GUID without braces.");
}

std::wstring EventName(const std::wstring& runId, const wchar_t* suffix) {
    return L"Local\\ScriptHangLab-" + runId + L"-" + suffix;
}

HANDLE NewEvent(const std::wstring& runId, const wchar_t* suffix) {
    const HANDLE event = CreateEventW(nullptr, TRUE, FALSE, EventName(runId, suffix).c_str());
    if (!event) Fail("CreateEventW");
    const DWORD error = GetLastError();
    if (error == ERROR_ALREADY_EXISTS) {
        CloseHandle(event);
        throw std::runtime_error("A lab event already exists for this run ID; choose a new run directory.");
    }
    return event;
}

HANDLE ExistingEvent(const std::wstring& runId, const wchar_t* suffix) {
    const HANDLE event = OpenEventW(SYNCHRONIZE | EVENT_MODIFY_STATE, FALSE,
        EventName(runId, suffix).c_str());
    if (!event) Fail("OpenEventW");
    return event;
}

void Signal(HANDLE event) {
    if (!SetEvent(event)) Fail("SetEvent");
}

void WriteText(const std::wstring& path, const std::string& text) {
    const Handle file(CreateFileW(path.c_str(), GENERIC_WRITE, FILE_SHARE_READ, nullptr,
        CREATE_NEW, FILE_ATTRIBUTE_NORMAL, nullptr));
    if (!file.valid()) Fail("CreateFileW(metadata)");
    if (text.size() > MAXDWORD) throw std::runtime_error("Metadata is too large.");
    DWORD written = 0;
    if (!WriteFile(file.get(), text.data(), static_cast<DWORD>(text.size()), &written, nullptr)) Fail("WriteFile(metadata)");
    if (written != text.size()) throw std::runtime_error("Metadata write was incomplete.");
}

void WriteOutput(const char* text) {
    DWORD written = 0;
    const DWORD size = static_cast<DWORD>(std::strlen(text));
    if (!WriteFile(GetStdHandle(STD_OUTPUT_HANDLE), text, size, &written, nullptr)) Fail("WriteFile(stdout)");
    if (written != size) throw std::runtime_error("Standard output write was incomplete.");
}

DWORD HandleFlags(HANDLE value) {
    DWORD flags = 0;
    if (!GetHandleInformation(value, &flags)) Fail("GetHandleInformation");
    return flags;
}
} // namespace

__declspec(noinline) void CloseOutputWriter() {
    if (gOutputWriter != nullptr && gOutputWriter != INVALID_HANDLE_VALUE) {
        const HANDLE writer = gOutputWriter;
        // SetStdHandle changes the standard-handle slot; it does not close writer.
        if (!SetStdHandle(STD_OUTPUT_HANDLE, INVALID_HANDLE_VALUE)) Fail("SetStdHandle(stdout)");
        if (GetStdHandle(STD_ERROR_HANDLE) == writer &&
            !SetStdHandle(STD_ERROR_HANDLE, INVALID_HANDLE_VALUE)) Fail("SetStdHandle(stderr)");
        if (GetStdHandle(STD_INPUT_HANDLE) == writer &&
            !SetStdHandle(STD_INPUT_HANDLE, INVALID_HANDLE_VALUE)) Fail("SetStdHandle(stdin)");
        if (!CloseHandle(writer)) Fail("CloseHandle(output writer)");
        gOutputWriter = INVALID_HANDLE_VALUE;
    }
    Signal(gWriterClosedEvent);
}

__declspec(noinline) void WaitForHelperCommands() {
    const ULONGLONG started = GetTickCount64();
    bool writerClosed = !gWriterWasInherited;
    for (;;) {
        const ULONGLONG elapsed = GetTickCount64() - started;
        if (elapsed >= kHelperLifetimeMs) return;
        const DWORD remaining = static_cast<DWORD>(kHelperLifetimeMs - elapsed);
        const HANDLE commands[] = { gStopEvent, gCloseWriterEvent };
        const DWORD count = writerClosed ? 1u : 2u;
        const DWORD wait = WaitForMultipleObjects(count, commands, FALSE, remaining);
        if (wait == WAIT_OBJECT_0 || wait == WAIT_TIMEOUT) return;
        if (!writerClosed && wait == WAIT_OBJECT_0 + 1) {
            CloseOutputWriter();
            writerClosed = true;
            continue;
        }
        if (wait == WAIT_FAILED) Fail("WaitForMultipleObjects(helper)");
        throw std::runtime_error("Unexpected helper wait result.");
    }
}

__declspec(noinline) PROCESS_INFORMATION LaunchHelper(const std::wstring& mode,
    const std::wstring& runId, const std::wstring& runDirectory) {
    const bool inheritWriter = mode == L"inherited-writer";
    const std::wstring executable = ExecutablePath();
    const std::wstring command = Quote(executable) + L" helper " + mode + L" " + runId + L" " + Quote(runDirectory);
    std::vector<wchar_t> commandLine(command.begin(), command.end());
    commandLine.push_back(L'\0');

    SECURITY_ATTRIBUTES attributes{ sizeof(SECURITY_ATTRIBUTES), nullptr, TRUE };
    // The broken helper inherits stdout and harmless NUL input/error handles.
    // Stdout may be a pipe (the hang) or a log file (the wrapper workaround).
    // The fixed helper is detached, needs no standard I/O, and inherits no handles.
    const Handle nullHandle(inheritWriter ? CreateFileW(L"NUL", GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_WRITE, &attributes, OPEN_EXISTING, 0, nullptr) : nullptr);
    if (inheritWriter && !nullHandle.valid()) Fail("CreateFileW(NUL)");
    STARTUPINFOW startup{};
    startup.cb = sizeof(startup);
    if (inheritWriter) {
        gOutputWriter = GetStdHandle(STD_OUTPUT_HANDLE);
        const DWORD outputType = GetFileType(gOutputWriter);
        if (outputType != FILE_TYPE_PIPE && outputType != FILE_TYPE_DISK)
            throw std::runtime_error("The child requires stdout redirected to a pipe or file.");
        if (!SetHandleInformation(gOutputWriter, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT))
            Fail("SetHandleInformation(stdout)");
        startup.dwFlags = STARTF_USESTDHANDLES;
        startup.hStdInput = nullHandle.get();
        startup.hStdOutput = gOutputWriter;
        startup.hStdError = nullHandle.get();
    }
    PROCESS_INFORMATION process{};
    if (!CreateProcessW(executable.c_str(), commandLine.data(), nullptr, nullptr,
        inheritWriter ? TRUE : FALSE, DETACHED_PROCESS, nullptr, runDirectory.c_str(), &startup, &process))
        Fail("CreateProcessW(helper)");
    gHelperPid = process.dwProcessId;
    return process;
}

int Child(const std::wstring& mode, const std::wstring& runId, const std::wstring& runDirectory) {
    const Handle ready(NewEvent(runId, L"Ready"));
    const Handle closeWriter(NewEvent(runId, L"CloseWriter"));
    const Handle writerClosed(NewEvent(runId, L"WriterClosed"));
    const Handle stop(NewEvent(runId, L"Stop"));
    const PROCESS_INFORMATION information = LaunchHelper(mode, runId, runDirectory);
    const Handle process(information.hProcess);
    const Handle thread(information.hThread);
    StopHelperOnFailure cleanup(stop.get(), process.get());
    const HANDLE waits[] = { ready.get(), process.get() };
    const DWORD wait = WaitForMultipleObjects(2, waits, FALSE, kReadyTimeoutMs);
    if (wait != WAIT_OBJECT_0) {
        if (wait == WAIT_FAILED) Fail("WaitForMultipleObjects(ready)");
        throw std::runtime_error("Helper did not become ready. Inspect helper-error.txt if present.");
    }
    WriteText(runDirectory + L"\\child.json", "{\n  \"childPid\": " + std::to_string(GetCurrentProcessId()) +
        ",\n  \"helperPid\": " + std::to_string(information.dwProcessId) + ",\n  \"runId\": " + Json(runId) +
        ",\n  \"mode\": " + Json(mode) + ",\n  \"executablePath\": " + Json(ExecutablePath()) + "\n}\n");
    WriteOutput("The native command wrote one line and exited.\r\n");
    cleanup.Complete();
    return 0;
}

int Helper(const std::wstring& mode, const std::wstring& runId, const std::wstring& runDirectory) {
    const Handle ready(ExistingEvent(runId, L"Ready"));
    const Handle closeWriter(ExistingEvent(runId, L"CloseWriter"));
    const Handle writerClosed(ExistingEvent(runId, L"WriterClosed"));
    const Handle stop(ExistingEvent(runId, L"Stop"));
    gReadyEvent = ready.get();
    gCloseWriterEvent = closeWriter.get();
    gWriterClosedEvent = writerClosed.get();
    gStopEvent = stop.get();
    gHelperPid = GetCurrentProcessId();
    gWriterWasInherited = mode == L"inherited-writer";
    const HANDLE standardOutput = GetStdHandle(STD_OUTPUT_HANDLE);
    const DWORD outputType = GetFileType(standardOutput);
    DWORD flags = 0;
    if (gWriterWasInherited) {
        if (outputType != FILE_TYPE_PIPE && outputType != FILE_TYPE_DISK)
            throw std::runtime_error("The broken helper did not receive a pipe or file writer.");
        gOutputWriter = standardOutput;
        flags = HandleFlags(gOutputWriter);
    } else {
        if (outputType == FILE_TYPE_PIPE) throw std::runtime_error("The fixed helper unexpectedly received a pipe.");
        Signal(gWriterClosedEvent);
    }
    WriteText(runDirectory + L"\\helper.json", "{\n  \"helperPid\": " + std::to_string(GetCurrentProcessId()) +
        ",\n  \"runId\": " + Json(runId) + ",\n  \"mode\": " + Json(mode) +
        ",\n  \"executablePath\": " + Json(ExecutablePath()) +
        ",\n  \"inheritedOutputWriter\": " + (gWriterWasInherited ? "true" : "false") +
        ",\n  \"outputWriterHandle\": " + Json(Hex(gOutputWriter)) +
        ",\n  \"standardOutputHandle\": " + Json(Hex(standardOutput)) +
        ",\n  \"standardOutputFileType\": " + std::to_string(outputType) +
        ",\n  \"outputWriterHandleFlags\": " + std::to_string(flags) +
        ",\n  \"safetyTimeoutSeconds\": 900\n}\n");
    Signal(gReadyEvent);
    WaitForHelperCommands();
    CloseOutputWriter();
    return 0;
}

int Control(const std::wstring& runId, const std::wstring& action) {
    const Handle ready(ExistingEvent(runId, L"Ready"));
    const Handle closeWriter(ExistingEvent(runId, L"CloseWriter"));
    const Handle writerClosed(ExistingEvent(runId, L"WriterClosed"));
    const Handle stop(ExistingEvent(runId, L"Stop"));
    if (action == L"close-writer") {
        Signal(closeWriter.get());
        const DWORD wait = WaitForSingleObject(writerClosed.get(), kReadyTimeoutMs);
        if (wait == WAIT_FAILED) Fail("WaitForSingleObject(writer closed)");
        if (wait != WAIT_OBJECT_0) throw std::runtime_error("Helper did not acknowledge closing its writer.");
    } else if (action == L"stop") Signal(stop.get());
    else if (action != L"status") throw std::runtime_error("Unknown control action.");
    const auto signalled = [](HANDLE event) { return WaitForSingleObject(event, 0) == WAIT_OBJECT_0; };
    std::printf("Ready=%s\nWriterClosed=%s\nStopRequested=%s\n",
        signalled(ready.get()) ? "true" : "false",
        signalled(writerClosed.get()) ? "true" : "false",
        signalled(stop.get()) ? "true" : "false");
    return 0;
}

int wmain(int argc, wchar_t* argv[]) {
    try {
        if (argc == 4 && std::wstring(argv[1]) == L"control") {
            ValidateRunId(argv[2]);
            return Control(argv[2], argv[3]);
        }
        if (argc != 5) throw std::runtime_error("Use Run-Lab.ps1 or Control-Lab.ps1; native syntax: child|helper mode run-id run-directory.");
        const std::wstring role = argv[1];
        const std::wstring mode = argv[2];
        const std::wstring runId = argv[3];
        if (mode != L"inherited-writer" && mode != L"fixed") throw std::runtime_error("Unknown lab mode.");
        ValidateRunId(runId);
        if (role == L"child") return Child(mode, runId, argv[4]);
        if (role == L"helper") return Helper(mode, runId, argv[4]);
        throw std::runtime_error("Unknown lab role.");
    } catch (const std::exception& error) {
        if (argc == 5 && std::wstring(argv[1]) == L"helper") {
            try { WriteText(std::wstring(argv[4]) + L"\\helper-error.txt", error.what()); } catch (...) {}
        }
        std::fprintf(stderr, "ScriptHangLab: %s\n", error.what());
        return 1;
    }
}