// PatchUpdate.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include <iostream>
#include <string>
#include <vector>
#include <windows.h>
#include <wuapi.h>
#include <comutil.h>
#include <urlmon.h>
#include <comdef.h>

std::string tempPath = "c:\\temp";
std::string wsusPath = tempPath + "\\WSUS2";
std::string patchscancab = tempPath + "\\patchscan.cab";

void DownloadFile(const std::string& url, const std::string& path) {
    HRESULT hr = URLDownloadToFile(NULL, url.c_str(), path.c_str(), 0, NULL);
    if (SUCCEEDED(hr)) {
        std::cout << "... Downloaded: " << path << std::endl;
    }
    else {
        std::cerr << "Failed to download: " << url << std::endl;
    }
}

void DeleteAndCreateDirectory(const std::string& path) {
    WIN32_FIND_DATAA findFileData;
    HANDLE hFind = FindFirstFileA((path + "\\*").c_str(), &findFileData);
    if (hFind != INVALID_HANDLE_VALUE) {
        do {
            const std::string fileOrDir = path + "\\" + findFileData.cFileName;
            if (findFileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
                if (strcmp(findFileData.cFileName, ".") != 0 && strcmp(findFileData.cFileName, "..") != 0) {
                    DeleteAndCreateDirectory(fileOrDir); // Rekursiv löschen
                }
            }
            else {
                DeleteFileA(fileOrDir.c_str());
            }
        } while (FindNextFileA(hFind, &findFileData) != 0);
        FindClose(hFind);
        RemoveDirectoryA(path.c_str());
        std::cout << "directory deleted: " << path << std::endl;
    }

    CreateDirectoryA(path.c_str(), NULL);
}


void ProcessDownloads(IUpdate* update) {
    IUpdateDownloadContentCollection* downloadContents = nullptr;
    HRESULT hr = update->get_DownloadContents(&downloadContents);
    if (FAILED(hr)) {
        std::cerr << "Failed to get download contents." << std::endl;
        return;
    }

    long count = 0;
    downloadContents->get_Count(&count);
    std::cout << "Check Downloads - found " << count << " downloads" << std::endl;

    bool skipDownloads = count > 20;
    if (skipDownloads) {
        std::cout << "More than 20 downloads - skip download!" << std::endl;
    }

    DeleteAndCreateDirectory(wsusPath);

    for (long i = 0; i < count; ++i) {
        IUpdateDownloadContent* downloadContent = nullptr;
        downloadContents->get_Item(i, &downloadContent);

        bstr_t downloadUrl;
        downloadContent->get_DownloadUrl(downloadUrl.GetAddress());

        if (!(!downloadUrl)) {
            std::cout << "##### Download URL: " << (char *)downloadUrl << std::endl;

            if (!skipDownloads) {
                IStringCollection* pathCollection = nullptr;
                hr = CoCreateInstance(__uuidof(StringCollection), NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&pathCollection));
                if (FAILED(hr)) {
                    std::cerr << "Failed to create StringCollection." << std::endl;
                    downloadContents->Release();
                    return;
                }

                std::string fileName = (char *)downloadUrl;
                size_t pos = fileName.find_last_of("/\\");
                if (pos != std::string::npos) {
                    fileName = fileName.substr(pos + 1);
                }
                std::string path = wsusPath + fileName;
                DownloadFile((char *)downloadUrl, path);

                long collIndex;
                hr = pathCollection->Add(_bstr_t(path.c_str()), &collIndex);
                if (SUCCEEDED(hr))
                {
                    IUpdate2 *update2;
                    hr = update->QueryInterface(IID_PPV_ARGS(&update2));
                    if (SUCCEEDED(hr))
                        hr = update2->CopyToCache(pathCollection);

                    if (FAILED(hr))
                        std::cerr << "!!!!!!!!!!!!!!!! Failed to copy to cache" << std::endl;
                    else
                        std::cout << "copied to cache successful" << std::endl;

                }
            }
            SysFreeString(downloadUrl);
        }
        downloadContent->Release();
    }

    downloadContents->Release();
}

void ProcessUpdate(IUpdate* update) 
{
    IUpdateCollection* bundledUpdates = nullptr;
    HRESULT hr = update->get_BundledUpdates(&bundledUpdates);
    if (SUCCEEDED(hr)) {
        long count = 0;
        bundledUpdates->get_Count(&count);

        for (long i = 0; i < count; ++i) {
            IUpdate* bundledUpdate = nullptr;
            bundledUpdates->get_Item(i, &bundledUpdate);

            BSTR title;
            bundledUpdate->get_Title(&title);
            std::cout << "    bundled update: " << title << std::endl;
            SysFreeString(title);

            ProcessUpdate(bundledUpdate);
            bundledUpdate->Release();
        }
        bundledUpdates->Release();
    }
    ProcessDownloads(update);
}

int main() {

    CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);

    IUpdateSession* updateSession = nullptr;
    auto hr = CoCreateInstance(__uuidof(UpdateSession), NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&updateSession));
    if (FAILED(hr)) {
        std::cerr << "Failed to create UpdateSession." << std::endl;
        return 1;
    }
    updateSession->put_ClientApplicationID(_bstr_t("PatchTester"));

    IUpdateServiceManager* updateServiceManager = nullptr;
    hr = CoCreateInstance(__uuidof(UpdateServiceManager), NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&updateServiceManager));
    if (FAILED(hr)) {
        std::cerr << "failed to create updateServiceManager " << hr << std::endl;
        CoUninitialize();
        return 1;
    }

    IUpdateService* updateService = nullptr;
    hr = updateServiceManager->AddScanPackageService(_bstr_t("Custom Update Service"), _bstr_t(patchscancab.c_str()), 0, &updateService);
    if (FAILED(hr)) {
        updateService->Release();
        std::cerr << "error adding update service: " << hr << std::endl;
        return 1;
    }

    IUpdateSearcher* updateSearcher = nullptr;
    hr = updateSession->CreateUpdateSearcher(&updateSearcher);
    if (FAILED(hr)) {
        std::cerr << "Failed to create UpdateSearcher." << std::endl;
        updateSession->Release();
        return 1;
    }

    updateSearcher->put_ServerSelection(ssOthers);
    bstr_t serviceId;
    updateService->get_ServiceID(serviceId.GetAddress());
    updateSearcher->put_ServiceID(serviceId);

    ISearchResult* searchResult = nullptr;
    hr = updateSearcher->Search(_bstr_t("IsInstalled=0 and Type='Software'"), &searchResult);
    if (FAILED(hr)) {
        std::cerr << "Failed to search for updates." << std::endl;
        updateSearcher->Release();
        updateSession->Release();
        updateService->Release();
        return 1;
    }

    IUpdateCollection* updates = nullptr;
    hr = searchResult->get_Updates(&updates);
    if (FAILED(hr)) {
        std::cerr << "Failed to get updates." << std::endl;
        updateSearcher->Release();
        updateSession->Release();
        updateService->Release();
        return 1;
    }

    long count = 0;
    updates->get_Count(&count);

    if (count > 0) std::cout << "updates found" << std::endl;

    for (long i = 0; i < count; ++i) {
        IUpdate* update = nullptr;
        updates->get_Item(i, &update);

        BSTR title;
        update->get_Title(&title);
        std::cout << "main update: " << title << std::endl;
        SysFreeString(title);

        ProcessUpdate(update);
        update->Release();
    }
    updates->Release();
    updateService->Release();
    searchResult->Release();
    updateSearcher->Release();
    updateSession->Release();
    CoUninitialize();
    return 0;
}