#define UNICODE
#define _UNICODE
#include <windows.h>
#include <userenv.h>
#include <sddl.h>
#include <aclapi.h>
#include <stdio.h>
#include <string.h>
#include <string>
#include <algorithm>
#include <vector>

#pragma comment(lib, "userenv.lib")
#pragma comment(lib, "ole32.lib")

static bool GrantContainerFolderAccess(const std::wstring& path, PSID sid) {
    PACL oldDacl = nullptr;
    PSECURITY_DESCRIPTOR descriptor = nullptr;
    DWORD error = GetNamedSecurityInfoW(
        const_cast<LPWSTR>(path.c_str()), SE_FILE_OBJECT, DACL_SECURITY_INFORMATION,
        nullptr, nullptr, &oldDacl, nullptr, &descriptor);
    if (error != ERROR_SUCCESS) return false;

    EXPLICIT_ACCESSW access{};
    access.grfAccessPermissions = GENERIC_ALL;
    access.grfAccessMode = GRANT_ACCESS;
    access.grfInheritance = OBJECT_INHERIT_ACE | CONTAINER_INHERIT_ACE;
    BuildTrusteeWithSidW(&access.Trustee, sid);
    PACL updatedDacl = nullptr;
    error = SetEntriesInAclW(1, &access, oldDacl, &updatedDacl);
    if (error == ERROR_SUCCESS) {
        error = SetNamedSecurityInfoW(
            const_cast<LPWSTR>(path.c_str()), SE_FILE_OBJECT, DACL_SECURITY_INFORMATION,
            nullptr, nullptr, updatedDacl, nullptr);
    }
    if (updatedDacl) LocalFree(updatedDacl);
    if (descriptor) LocalFree(descriptor);
    return error == ERROR_SUCCESS;
}

static std::wstring QuoteArg(const std::wstring& value) {
    std::wstring out = L"\"";
    size_t slashes = 0;
    for (wchar_t ch : value) {
        if (ch == L'\\') {
            ++slashes;
        } else if (ch == L'\"') {
            out.append(slashes * 2 + 1, L'\\');
            out += ch;
            slashes = 0;
        } else {
            out.append(slashes, L'\\');
            slashes = 0;
            out += ch;
        }
    }
    out.append(slashes * 2, L'\\');
    out += L'\"';
    return out;
}

static std::vector<wchar_t> SafeEnvironmentBlock() {
    static const wchar_t* const allowed[] = {
        L"ALLUSERSPROFILE", L"APPDATA", L"COMSPEC", L"HOMEDRIVE", L"HOMEPATH",
        L"LOCALAPPDATA", L"NUMBER_OF_PROCESSORS", L"OS", L"PATH", L"PATHEXT",
        L"PROCESSOR_ARCHITECTURE", L"PROCESSOR_IDENTIFIER", L"PROGRAMDATA",
        L"PROGRAMFILES", L"PROGRAMFILES(X86)", L"PUBLIC", L"SYSTEMDRIVE",
        L"SYSTEMROOT", L"TEMP", L"TMP", L"USERDOMAIN", L"USERNAME", L"USERPROFILE",
        L"WINDIR"
    };
    std::vector<std::wstring> entries;
    LPWCH inherited = GetEnvironmentStringsW();
    if (!inherited) return {};
    for (const wchar_t* item = inherited; *item;) {
        std::wstring entry(item);
        item += entry.size() + 1;
        size_t separator = entry.find(L'=');
        if (separator == std::wstring::npos || separator == 0) continue;
        std::wstring key = entry.substr(0, separator);
        bool keep = false;
        for (const wchar_t* candidate : allowed) {
            if (_wcsicmp(key.c_str(), candidate) == 0) {
                keep = true;
                break;
            }
        }
        if (keep) entries.push_back(std::move(entry));
    }
    FreeEnvironmentStringsW(inherited);
    std::sort(entries.begin(), entries.end(), [](const std::wstring& left, const std::wstring& right) {
        return _wcsicmp(left.c_str(), right.c_str()) < 0;
    });
    std::vector<wchar_t> block;
    for (const auto& entry : entries) {
        block.insert(block.end(), entry.begin(), entry.end());
        block.push_back(L'\0');
    }
    block.push_back(L'\0');
    if (entries.empty()) block.push_back(L'\0');
    return block;
}

static void SetEnvironmentValue(
    std::vector<wchar_t>& block, const std::wstring& key, const std::wstring& value) {
    std::vector<std::wstring> entries;
    for (const wchar_t* item = block.data(); item && *item;) {
        std::wstring entry(item);
        item += entry.size() + 1;
        size_t separator = entry.find(L'=');
        if (separator == std::wstring::npos || _wcsicmp(entry.substr(0, separator).c_str(), key.c_str()) != 0) {
            entries.push_back(std::move(entry));
        }
    }
    entries.push_back(key + L"=" + value);
    std::sort(entries.begin(), entries.end(), [](const std::wstring& left, const std::wstring& right) {
        return _wcsicmp(left.c_str(), right.c_str()) < 0;
    });
    block.clear();
    for (const auto& entry : entries) {
        block.insert(block.end(), entry.begin(), entry.end());
        block.push_back(L'\0');
    }
    block.push_back(L'\0');
}

static bool FilesMatch(const std::wstring& leftPath, const std::wstring& rightPath) {
    HANDLE left = CreateFileW(leftPath.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr,
        OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    HANDLE right = CreateFileW(rightPath.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr,
        OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (left == INVALID_HANDLE_VALUE || right == INVALID_HANDLE_VALUE) {
        if (left != INVALID_HANDLE_VALUE) CloseHandle(left);
        if (right != INVALID_HANDLE_VALUE) CloseHandle(right);
        return false;
    }
    LARGE_INTEGER leftSize{}, rightSize{};
    bool matches = GetFileSizeEx(left, &leftSize) && GetFileSizeEx(right, &rightSize) &&
        leftSize.QuadPart == rightSize.QuadPart && leftSize.QuadPart >= 0 &&
        leftSize.QuadPart <= 1024 * 1024;
    std::vector<char> leftBytes(matches ? static_cast<size_t>(leftSize.QuadPart) : 0);
    std::vector<char> rightBytes(matches ? static_cast<size_t>(rightSize.QuadPart) : 0);
    DWORD leftRead = 0, rightRead = 0;
    if (matches && !leftBytes.empty()) {
        matches = ReadFile(left, leftBytes.data(), static_cast<DWORD>(leftBytes.size()),
                &leftRead, nullptr) &&
            ReadFile(right, rightBytes.data(), static_cast<DWORD>(rightBytes.size()),
                &rightRead, nullptr) &&
            leftRead == static_cast<DWORD>(leftBytes.size()) &&
            rightRead == static_cast<DWORD>(rightBytes.size()) &&
            leftBytes == rightBytes;
    }
    CloseHandle(left);
    CloseHandle(right);
    return matches;
}

static DWORD RunContained(
    PSID appContainerSid,
    const std::wstring& command,
    const wchar_t* environment,
    const std::wstring& cwd,
    HANDLE job
) {
    SIZE_T attributeBytes = 0;
    InitializeProcThreadAttributeList(nullptr, 1, 0, &attributeBytes);
    if (GetLastError() != ERROR_INSUFFICIENT_BUFFER) return GetLastError();

    auto* attributes = static_cast<LPPROC_THREAD_ATTRIBUTE_LIST>(
        HeapAlloc(GetProcessHeap(), 0, attributeBytes)
    );
    if (!attributes) return ERROR_OUTOFMEMORY;
    if (!InitializeProcThreadAttributeList(attributes, 1, 0, &attributeBytes)) {
        DWORD error = GetLastError();
        HeapFree(GetProcessHeap(), 0, attributes);
        return error;
    }

    SECURITY_CAPABILITIES security{};
    security.AppContainerSid = appContainerSid;
    security.Capabilities = nullptr;
    security.CapabilityCount = 0;

    if (!UpdateProcThreadAttribute(
            attributes,
            0,
            PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES,
            &security,
            sizeof(security),
            nullptr,
            nullptr)) {
        DWORD error = GetLastError();
        DeleteProcThreadAttributeList(attributes);
        HeapFree(GetProcessHeap(), 0, attributes);
        return error;
    }

    STARTUPINFOEXW startup{};
    startup.StartupInfo.cb = sizeof(startup);
    startup.StartupInfo.dwFlags = STARTF_USESHOWWINDOW;
    startup.StartupInfo.wShowWindow = SW_HIDE;
    startup.lpAttributeList = attributes;
    PROCESS_INFORMATION process{};
    std::wstring commandLine = QuoteArg(L"C:\\Windows\\System32\\cmd.exe") +
        L" /d /s /c \"" + command + L"\"";
    std::wstring mutableLine = commandLine;
    constexpr DWORD flags = EXTENDED_STARTUPINFO_PRESENT | CREATE_UNICODE_ENVIRONMENT |
        CREATE_NO_WINDOW | CREATE_SUSPENDED;

    BOOL created = CreateProcessW(
        L"C:\\Windows\\System32\\cmd.exe",
        mutableLine.empty() ? nullptr : &mutableLine[0],
        nullptr,
        nullptr,
        FALSE,
        flags,
        const_cast<wchar_t*>(environment),
        cwd.c_str(),
        &startup.StartupInfo,
        &process
    );
    DWORD error = created ? ERROR_SUCCESS : GetLastError();
    DeleteProcThreadAttributeList(attributes);
    HeapFree(GetProcessHeap(), 0, attributes);
    if (!created) return error;

    if (!AssignProcessToJobObject(job, process.hProcess)) {
        error = GetLastError();
        TerminateProcess(process.hProcess, error);
        CloseHandle(process.hThread);
        CloseHandle(process.hProcess);
        return error;
    }
    ResumeThread(process.hThread);
    DWORD waitResult = WaitForSingleObject(process.hProcess, 30000);
    DWORD exitCode = ERROR_TIMEOUT;
    if (waitResult == WAIT_OBJECT_0) {
        GetExitCodeProcess(process.hProcess, &exitCode);
    } else {
        TerminateJobObject(job, ERROR_TIMEOUT);
    }
    CloseHandle(process.hThread);
    CloseHandle(process.hProcess);
    return exitCode;
}

int wmain() {
    wchar_t tempPath[MAX_PATH]{};
    if (!GetTempPathW(MAX_PATH, tempPath)) return 10;
    GUID id{};
    if (FAILED(CoCreateGuid(&id))) return 11;
    wchar_t idText[40]{};
    if (StringFromGUID2(id, idText, 40) == 0) return 12;
    std::wstring nonce = idText;
    if (!nonce.empty() && nonce.front() == L'{') nonce.erase(nonce.begin());
    if (!nonce.empty() && nonce.back() == L'}') nonce.pop_back();
    std::wstring profileName = L"SovereignAIProbe" + nonce;
    std::wstring secretPath = std::wstring(tempPath) + L"SovereignAIProbe" + nonce + L".txt";
    std::wstring writePath = std::wstring(tempPath) + L"SovereignAIProbe" + nonce + L"-write.txt";
    const char marker[] = "host-secret-must-stay-private";
    HANDLE secret = CreateFileW(secretPath.c_str(), GENERIC_WRITE, 0, nullptr, CREATE_NEW,
        FILE_ATTRIBUTE_TEMPORARY, nullptr);
    if (secret == INVALID_HANDLE_VALUE) return 13;
    DWORD bytesWritten = 0;
    if (!WriteFile(secret, marker, sizeof(marker) - 1, &bytesWritten, nullptr) ||
        bytesWritten != sizeof(marker) - 1) {
        CloseHandle(secret);
        return 14;
    }
    CloseHandle(secret);

    PSID appContainerSid = nullptr;
    HRESULT profileResult = CreateAppContainerProfile(
        profileName.c_str(), L"SovereignAI isolated command probe",
        L"Temporary test profile for command isolation", nullptr, 0, &appContainerSid);
    if (FAILED(profileResult)) {
        fwprintf(stderr, L"CreateAppContainerProfile failed: 0x%08lx (Win32 %lu)\n",
            static_cast<unsigned long>(profileResult),
            static_cast<unsigned long>(HRESULT_CODE(profileResult)));
        DeleteFileW(secretPath.c_str());
        return static_cast<int>(HRESULT_CODE(profileResult));
    }
    LPWSTR sidText = nullptr;
    PWSTR appContainerFolder = nullptr;
    bool appContainerPathReady = ConvertSidToStringSidW(appContainerSid, &sidText) &&
        SUCCEEDED(GetAppContainerFolderPath(sidText, &appContainerFolder));
    if (!appContainerPathReady) {
        FreeSid(appContainerSid);
        DeleteAppContainerProfile(profileName.c_str());
        DeleteFileW(secretPath.c_str());
        return 16;
    }
    if (sidText) LocalFree(sidText);

    HANDLE job = CreateJobObjectW(nullptr, nullptr);
    JOBOBJECT_EXTENDED_LIMIT_INFORMATION limits{};
    limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE |
        JOB_OBJECT_LIMIT_ACTIVE_PROCESS | JOB_OBJECT_LIMIT_JOB_MEMORY;
    limits.BasicLimitInformation.ActiveProcessLimit = 8;
    limits.JobMemoryLimit = 512ull * 1024 * 1024;
    bool jobReady = job && SetInformationJobObject(
        job, JobObjectExtendedLimitInformation, &limits, sizeof(limits));

    auto environment = SafeEnvironmentBlock();
    if (environment.empty()) return 15;
    std::wstring appDataPath(appContainerFolder);
    std::wstring appTempPath = appDataPath + L"\\Temp";
    std::wstring sandboxWorkspacePath = appDataPath + L"\\agent-workspace";
    std::wstring stagedInputPath = sandboxWorkspacePath + L"\\README.md";
    std::wstring stagedCopyPath = sandboxWorkspacePath + L"\\staged-copy.md";
    CreateDirectoryW(appTempPath.c_str(), nullptr);
    bool workspaceFolderReady = CreateDirectoryW(sandboxWorkspacePath.c_str(), nullptr) != FALSE ||
        GetLastError() == ERROR_ALREADY_EXISTS;
    workspaceFolderReady = workspaceFolderReady &&
        GrantContainerFolderAccess(sandboxWorkspacePath, appContainerSid);
    bool stagedInputCopied = workspaceFolderReady &&
        CopyFileW(L"README.md", stagedInputPath.c_str(), FALSE);
    std::wstring curlPath = sandboxWorkspacePath + L"\\curl.exe";
    bool curlCopied = workspaceFolderReady &&
        CopyFileW(L"C:\\Windows\\System32\\curl.exe", curlPath.c_str(), FALSE);
    SetEnvironmentValue(environment, L"PATH", L"C:\\Windows\\System32");
    SetEnvironmentValue(environment, L"APPDATA", appDataPath);
    SetEnvironmentValue(environment, L"LOCALAPPDATA", appDataPath);
    SetEnvironmentValue(environment, L"TEMP", appTempPath);
    SetEnvironmentValue(environment, L"TMP", appTempPath);
    std::wstring cwd = L"C:\\Windows\\System32";
    DWORD shellResult = jobReady
        ? RunContained(appContainerSid, L"exit 0", environment.data(), cwd, job)
        : GetLastError();
    std::wstring readAttempt = L"type " + secretPath + L" >nul 2>nul";
    DWORD readResult = jobReady
        ? RunContained(appContainerSid, readAttempt, environment.data(), cwd, job)
        : ERROR_INVALID_HANDLE;
    std::wstring writeAttempt = L"echo modified>" + writePath + L" 2>nul";
    DWORD writeResult = jobReady
        ? RunContained(appContainerSid, writeAttempt, environment.data(), cwd, job)
        : ERROR_INVALID_HANDLE;
    std::wstring allowedWritePath = sandboxWorkspacePath + L"\\probe-output.txt";
    DWORD allowedWriteResult = jobReady && workspaceFolderReady
        ? RunContained(appContainerSid, L"echo contained>" + allowedWritePath,
            environment.data(), cwd, job)
        : ERROR_INVALID_HANDLE;
    std::wstring stagedFileUrl = L"file:///";
    for (wchar_t ch : stagedInputPath) {
        stagedFileUrl += ch == L'\\' ? L'/' : ch;
    }
    DWORD stagedReadResult = jobReady && stagedInputCopied && curlCopied
        ? RunContained(appContainerSid, QuoteArg(curlPath) + L" --fail --silent " +
            stagedFileUrl + L" -o " + stagedCopyPath, environment.data(), cwd, job)
        : ERROR_INVALID_HANDLE;
    DWORD curlVersionResult = jobReady && curlCopied
        ? RunContained(appContainerSid, QuoteArg(curlPath) + L" --version",
            environment.data(), cwd, job)
        : ERROR_INVALID_HANDLE;
    DWORD curlNetworkResult = jobReady && curlCopied
        ? RunContained(appContainerSid, QuoteArg(curlPath) +
            L" --noproxy \"*\" --max-time 4 http://127.0.0.1:8100/health -o NUL",
            environment.data(), cwd, job)
        : ERROR_INVALID_HANDLE;

    DWORD attrs = GetFileAttributesW(writePath.c_str());
    DWORD attrsError = attrs == INVALID_FILE_ATTRIBUTES ? GetLastError() : ERROR_SUCCESS;
    bool hostSecretPreserved = false;
    secret = CreateFileW(secretPath.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr,
        OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (secret != INVALID_HANDLE_VALUE) {
        char check[sizeof(marker)]{};
        DWORD bytesRead = 0;
        hostSecretPreserved = ReadFile(secret, check, sizeof(marker) - 1, &bytesRead, nullptr) &&
            bytesRead == sizeof(marker) - 1 && memcmp(check, marker, sizeof(marker) - 1) == 0;
        CloseHandle(secret);
    }
    bool allowedWorkspaceWriteVisible = false;
    HANDLE allowedOutput = CreateFileW(allowedWritePath.c_str(), GENERIC_READ, FILE_SHARE_READ,
        nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (allowedOutput != INVALID_HANDLE_VALUE) {
        char contents[32]{};
        DWORD count = 0;
        allowedWorkspaceWriteVisible = ReadFile(allowedOutput, contents, sizeof(contents) - 1,
            &count, nullptr) && count >= 9 && memcmp(contents, "contained", 9) == 0;
        CloseHandle(allowedOutput);
    }
    bool stagedInputRoundTripMatches = stagedInputCopied && FilesMatch(
        L"README.md", stagedCopyPath);

    if (job) CloseHandle(job);
    if (appContainerFolder) CoTaskMemFree(appContainerFolder);
    FreeSid(appContainerSid);
    DeleteFileW(curlPath.c_str());
    DeleteFileW(stagedInputPath.c_str());
    DeleteFileW(stagedCopyPath.c_str());
    DeleteAppContainerProfile(profileName.c_str());
    DeleteFileW(secretPath.c_str());
    DeleteFileW(writePath.c_str());
    DeleteFileW(allowedWritePath.c_str());

    bool writeWasBlocked = attrs == INVALID_FILE_ATTRIBUTES &&
        (attrsError == ERROR_FILE_NOT_FOUND || attrsError == ERROR_ACCESS_DENIED);
    wprintf(L"shell_exit=%lu\nread_host_file_exit=%lu\nwrite_host_file_exit=%lu\nhost_secret_preserved=%s\nwrite_artifact_absent=%s\n",
        shellResult, readResult, writeResult,
        hostSecretPreserved ? L"true" : L"false",
        writeWasBlocked ? L"true" : L"false");
    wprintf(L"profile_write_exit=%lu\nprofile_write_visible=%s\n",
        allowedWriteResult, allowedWorkspaceWriteVisible ? L"true" : L"false");
    wprintf(L"staged_copy_exit=%lu\nstaged_input_copied=%s\nstaged_roundtrip_matches=%s\n",
        stagedReadResult, stagedInputCopied ? L"true" : L"false",
        stagedInputRoundTripMatches ? L"true" : L"false");
    wprintf(L"curl_version_exit=%lu\ncurl_local_health_exit=%lu\n",
        curlVersionResult, curlNetworkResult);
    if (!jobReady || shellResult != 0 || readResult == 0 || writeResult == 0 ||
        !hostSecretPreserved || !writeWasBlocked || !workspaceFolderReady ||
        allowedWriteResult != 0 || !allowedWorkspaceWriteVisible ||
        !stagedInputCopied || stagedReadResult != 0 || !stagedInputRoundTripMatches ||
        !curlCopied || curlVersionResult != 0 || curlNetworkResult == 0) return 20;
    return 0;
}