#include <iostream>
#include <Windows.h>
#include <string>
#include <WDBGEXTS.H>

#include "common.h"

#include "threadpool.h"
#include "client.h"

#include "../user/um/umanager.h"
#include "../user/km/kmanager.h"

DWORD WINAPI Init(HINSTANCE hinstDLL)
{
    AllocConsole();
    FILE* file;
    freopen_s( &file, "CONOUT$", "w", stdout );
    freopen_s( &file, "CONIN$", "r", stdin );

    std::this_thread::sleep_for( std::chrono::seconds( 1 ) );

    LPTSTR pipe_name = (LPTSTR)L"\\\\.\\pipe\\DonnaACPipe";
    LPCWSTR driver_name = L"\\\\.\\DonnaAC";

    std::shared_ptr<global::ThreadPool> thread_pool = std::make_shared<global::ThreadPool>( 4 );
    std::shared_ptr<global::Client> client_interface = std::make_shared<global::Client>( thread_pool, pipe_name );

    usermode::UManager umanager( thread_pool, client_interface );
    kernelmode::KManager kmanager( driver_name, thread_pool, client_interface);

    global::headers::SYSTEM_INFORMATION system_information;
    kmanager.SendClientHardwareInformation();

    global::report_structures::SYSTEM_INFORMATION_REQUEST_RESPONSE response;

    client_interface->ServerReceive( &response, sizeof( response ) );

    std::cout << "RequestID: " << response.RequestId << " CanUserProceed: " << 
        response.CanUserProceed << " Reason: " << response.reason << std::endl;

    while ( !GetAsyncKeyState( VK_DELETE ) )
    {
        //srand( time( NULL ) );
        //int seed = ( rand() % 7 );

        //std::cout << "Seed: " << seed << std::endl;

        //switch ( seed )
        //{
        //case 0:
        //    kmanager.EnumerateHandleTables();
        //    break;
        //case 1:
        //    kmanager.PerformIntegrityCheck();
        //    break;
        //case 2:
        //    kmanager.ScanPoolsForUnlinkedProcesses();
        //    break;
        //case 3:
        //    kmanager.VerifySystemModules();
        //    break;
        //case 4:
        //    kmanager.ValidateProcessModules();
        //    break;
        //case 5:
        //    kmanager.RunNmiCallbacks();
        //    break;
        //case 6:
        //    kmanager.CheckForAttachedThreads();
        //    break;
        //}

        //kmanager.MonitorCallbackReports();
        kmanager.ScanPoolsForUnlinkedProcesses();
        kmanager.MonitorCallbackReports();
        std::this_thread::sleep_for( std::chrono::seconds( 10 ) );
    }

    fclose( stdout );
    fclose( stdin );
    FreeConsole();

    FreeLibraryAndExitThread( hinstDLL, 0);
    return 0;
}

BOOL WINAPI DllMain(
    HINSTANCE hinstDLL,  // handle to DLL module
    DWORD fdwReason,     // reason for calling function
    LPVOID lpvReserved )  // reserved
{
    // Perform actions based on the reason for calling.
    switch ( fdwReason )
    {
    case DLL_PROCESS_ATTACH:

        DisableThreadLibraryCalls( hinstDLL );

        const auto thread = CreateThread(
            nullptr,
            0,
            reinterpret_cast< LPTHREAD_START_ROUTINE >( Init ),
            hinstDLL,
            0,
            nullptr
        );

        if ( thread )
            CloseHandle( thread );

        break;
    }
    return TRUE;  // Successful DLL_PROCESS_ATTACH.
}