#include "stdafx.h"
#include "RemoteAdmin.h"
#include "ProcessUpdate.h"
#include "Command.h"
#include "RemoteAdminDoc.h"
#include "MachineView.h"
#include "Resource.h"
#include "RemoteAdminView.h"
#include "GlobalMFCHelperFunc.h"


extern BOOL g_bUpdateProcessList;
extern UINT g_iUpdateProcessDelay;
extern CRITICAL_SECTION g_CriticalSection;

// Thread function that updates the all machines process's
UINT __stdcall ProcessUpdate::UpdateProcessListForAllMachines(void* pParam)
{
    SDocView* pDocView = reinterpret_cast<SDocView*>(pParam);

    CRemoteAdminDoc* pDoc              = pDocView->pDoc;
    CRemoteAdminView* pRemoteAdminView = pDocView->pRemoteAdminView;
    CMachineView* pMachineView         = pDocView->pMachineView;

    CString* pstrConnctedMachinesIP = NULL;
    int iNumberOfConnectedMachines = 0;
	OVERLAPPED olOverlapped = {0};
    
    // Debug diagnostics
    #ifdef _DEBUG
        int count = 0;
        CString strDebug;
    #endif
    CProcessInfoList pilProcessInfo;
    while (TRUE)
    {   
        // << Debug diagnostics
        #ifdef _DEBUG
            strDebug.Format("\n\nThreadID %d : counter %d\n", ::GetCurrentThreadId(), ++count);
        #endif
        TRACE0(strDebug.GetBuffer(0));         


        if (g_bUpdateProcessList == FALSE)
        {
            return 0;
        }
        //::EnterCriticalSection(&g_CriticalSection);
        pDoc->GetConnectedMachinesIP(&pstrConnctedMachinesIP, &iNumberOfConnectedMachines);

		// Start to get the processes on the remote computer and diaplay it
		for (int iCounter = 0; iCounter < iNumberOfConnectedMachines; ++iCounter)
        {
            HANDLE hProcessInfoPipe = pDoc->GetRemoteAdminProcessInfoPipe(pstrConnctedMachinesIP[iCounter]);
            
            #ifdef _DEBUG
                strDebug.Format("IP %s\n", pstrConnctedMachinesIP[iCounter].GetBuffer(0));
                TRACE0(strDebug.GetBuffer(0));
            #endif

            if (hProcessInfoPipe != NULL)
            {
				// Tell the server application, that we want an updated process list from 
				// that machine
                TriggerRemoteServiceToPrepareForProcessInfoUpdate(hProcessInfoPipe);              

				// First get the process count on that machine
                UINT iProcessCount = ProcessUpdate::ReceiveProcessCountFromRemoteMachine(hProcessInfoPipe);
                
				// This has to be a n/w problem
				if (iProcessCount <= 0 || iProcessCount >= INT_MAX)  
				{
					if (pstrConnctedMachinesIP[iCounter] != _T(""))
						ProcessUpdate::DisconnectAndReconnectToRemoteService(pstrConnctedMachinesIP[iCounter], pDoc, pMachineView, pRemoteAdminView);					

					continue;
				}

                #ifdef _DEBUG
                    strDebug.Format("Total Processes %d\n",iProcessCount);
                    TRACE0(strDebug.GetBuffer(0));
                #endif

                // Create a new list with all the updated processes
                //CProcessInfoList pilProcessInfo;
                for (UINT i = 0; i < iProcessCount; ++i)
                {
					ProcessUpdate::ReceiveProcessFromRemoteMachine(hProcessInfoPipe, pilProcessInfo, pstrConnctedMachinesIP[iCounter], iProcessCount, i + 1);       
					
					// This may be that case of the user loading a new file for monitoring from the
					// open command or MRU list. So please quit. g_bUpdateProcessList may be made 
					// FALSE in CRemoteAdminDoc::OnOpenDocument()
					if (g_bUpdateProcessList == FALSE)
					{
						return 0;
					}
                }

                // Now update with new processes info forthis machine IP
                pDoc->RefreshProcessList(pstrConnctedMachinesIP[iCounter], pilProcessInfo);

                // Now free the list
				ProcessUpdate::FreeProcessInfoList(pilProcessInfo);
            }
        }
  
        CString strIP = pMachineView->GetSelectedItemTextInTreeView();

        // strIP will be zero for root, else it will give the IP of the selected
        // machine
        if (strIP != _T("0"))
        {
            pRemoteAdminView->RefreshProcesses(strIP);
        }
        
		// Free the memory allocated for machine information
		if (pstrConnctedMachinesIP != NULL)
		{
			delete[] pstrConnctedMachinesIP;
			pstrConnctedMachinesIP = NULL;
		}
        
        
        //::LeaveCriticalSection(&g_CriticalSection);
        ::Sleep(g_iUpdateProcessDelay);
    }

	if (pDocView)
	{
		delete pDocView;
	}
    return 0;
}


void __stdcall ProcessUpdate::ReceiveProcessFromRemoteMachine(HANDLE hProcessInfoPipe, CProcessInfoList& pilProcessInfo, CString strIP, UINT iProcessCount, UINT iCounter)
{
	::EnterCriticalSection(&g_CriticalSection);

	ASSERT(hProcessInfoPipe != NULL);
	ASSERT(hProcessInfoPipe != INVALID_HANDLE_VALUE);

	//PROCESSENTRY32* pPe = new PROCESSENTRY32;
    SProcessInfo* pPe = new SProcessInfo;
	ASSERT(pPe);

	// Debug diagnostics
    #ifdef _DEBUG
        int count = 0;
        CString strDebug;
    #endif
       
    if (pPe != NULL)
    {
		//::memset(pPe, 0, sizeof(PROCESSENTRY32));
        ::memset(pPe, 0, sizeof(SProcessInfo));

		#ifdef _DEBUG
			strDebug.Format("Before Process %d of %d IP %s\n", iCounter, iProcessCount, strIP.GetBuffer(0));
			TRACE0(strDebug.GetBuffer(0));
        #endif
						
		DWORD dwRead			= 0;
        DWORD dwWritten			= 0;
		OVERLAPPED olOverlapped = {0};
		
		//BOOL bOk = ::ReadFile(hProcessInfoPipe, pPe, sizeof(PROCESSENTRY32), &dwRead, &olOverlapped);
        BOOL bOk = ::ReadFile(hProcessInfoPipe, pPe, sizeof(SProcessInfo), &dwRead, &olOverlapped);
						
		// If not successful, give some more time for the i/o coperation to complete
		if (!bOk)
		{
			ProcessUpdate::WaitFor_IO_OperationToCompleteWithExtraTime(300, 200);
		}
		
        
		pilProcessInfo.AddTail(pPe);
	}
	else
    {
		::AfxMessageBox(IDS_NO_MEMORY_FOR_NEW_PROCESSINFO);
    }

	::LeaveCriticalSection(&g_CriticalSection);
}


void __stdcall ProcessUpdate::TriggerRemoteServiceToPrepareForProcessInfoUpdate(HANDLE hProcessInfoPipe)
{
	::EnterCriticalSection(&g_CriticalSection);

    DWORD dwWritten			= 0;
    SCommand cmd			= {0};
    OVERLAPPED olOverlapped = {0};

    cmd.m_bThreadExit = FALSE;
    
    TRACE0("Before write\n");

    BOOL bOk = ::WriteFile(hProcessInfoPipe, &cmd, sizeof(SCommand), &dwWritten, &olOverlapped);
	if (!bOk)
	{
		ProcessUpdate::WaitFor_IO_OperationToComplete(400);
	}

	::LeaveCriticalSection(&g_CriticalSection);
}


UINT __stdcall ProcessUpdate::ReceiveProcessCountFromRemoteMachine(HANDLE hProcessInfoPipe)
{
	TRACE0("Before Process count read\n");

	::EnterCriticalSection(&g_CriticalSection);

	UINT iProcessCount		= 0;
	DWORD dwRead			= 0;
	OVERLAPPED olOverlapped = {0};

    BOOL bOk = ::ReadFile(hProcessInfoPipe, &iProcessCount, sizeof(UINT), &dwRead, &olOverlapped);
				
	// If not successful, give some more time for the i/o coperation to complete
	if (!bOk)
	{
		ProcessUpdate::WaitFor_IO_OperationToCompleteWithExtraTime(1000, 2000);
	}

	::LeaveCriticalSection(&g_CriticalSection);

	return iProcessCount;
}


void __stdcall ProcessUpdate::DisconnectAndReconnectToRemoteService(CString strIP, CRemoteAdminDoc* pDoc, CMachineView* pMachineView, CRemoteAdminView* pRemoteAdminView)
{
	// Retrive the password before we delete machine information, as we require it for reconnection.
	//CString strPwd = pDoc->GetPasswordForMachine(strIP);

	// Had to do all these hackery because MFC CString fails here dues to reference counting.
	// Other wise we would have taken the password in a CString
	TCHAR szPwd[_MAX_PATH];

	::memset(szPwd, 0, _MAX_PATH);
	::strcpy(szPwd, (LPCTSTR)pDoc->GetPasswordForMachine(strIP));
	
	//if (strPwd == _T("Could not retrieve the password"))  // Only disconnect machine
	if (strcmp(szPwd, "Could not retrieve the password") == 0)  // Only disconnect machine
	{
		// Show a balloon the system tray, that this machine is being disconnected.
		(MFC_DocView::GetApp())->ShowBalloonMsgInTray(strIP, "Password retrieval failure! Disconnecting the machine....");

		// Since we could not retrieve the password for the machine, there is some problem, 
		// disconnect the machine. Since there is a password retrieval failure, no point in 
		// reconnecting
		pMachineView->DeleteMachineFromBeingMonitored(strIP, pRemoteAdminView);
	}
	else  // Disconnect and try to reconnect the machine
	{
        // Show a balloon the system tray, that this machine is being disconnected.
		(MFC_DocView::GetApp())->ShowBalloonMsgInTray(strIP, "Network error! Trying to reconnect.......");

		// Get the machine's password length for creating a buffer to hold the password
		//unsigned int iPwdLen = (strPwd.GetLength() + sizeof(TCHAR)); // + sizeof(TCHAR) to accomodate '\0'

		// Get the password in a buffer because of reference counting errors in 
		// MFC's CString.
		//TCHAR* szPwd = new TCHAR[iPwdLen];

		//ASSERT(szPwd != NULL);

	//	if (szPwd != NULL)
	//	{
			// Prepare the password buffer
//            ::memset(szPwd, 0, iPwdLen);

			// Copy the password to the buffer
//			::strcpy(szPwd, (LPCTSTR)strPwd);

			// Since process count is zero or less, there is some problem, disconnect the machine and
			// try to connect again
			pMachineView->DeleteMachineFromBeingMonitored(strIP, pRemoteAdminView);
	
			// Try adding the same machine again, as there was a problem, network or otherwise.
			//	CString strPwd = szPwd;
			pMachineView->AddMachine(strIP, szPwd, pDoc, pRemoteAdminView);

//			if (szPwd != NULL)
//			{
//				delete szPwd;
///			}
//		}
//		else
//		{
		//	pMachineView->DeleteMachineFromBeingMonitored(strIP, pRemoteAdminView);
//		}
	}
}


void __stdcall ProcessUpdate::FreeProcessInfoList(CProcessInfoList& pilProcessInfo)
{
	::EnterCriticalSection(&g_CriticalSection);

	//PROCESSENTRY32* pPe = NULL;
    SProcessInfo* pPe = NULL;
    
	POSITION pos = pilProcessInfo.GetHeadPosition();
    while (pos != (POSITION)0xcdcdcdcd && pos != NULL)
    {
		pPe = pilProcessInfo.GetNext(pos);
        delete pPe;
    }

    pilProcessInfo.RemoveAll();

	::LeaveCriticalSection(&g_CriticalSection);
}

void __stdcall ProcessUpdate::WaitFor_IO_OperationToComplete(UINT iTimeToComplete)
{
	DWORD dwError = ::GetLastError();

	switch (dwError)
	{
	case ERROR_IO_PENDING:
       	::Sleep(iTimeToComplete);
		break;

	default:
		ASSERT(0);   // Generally shouldn't reach here
		break;
	}
}
             

void __stdcall ProcessUpdate::WaitFor_IO_OperationToCompleteWithExtraTime(UINT iTimeToComplete, UINT iExtraTimeToComplete)
{
	DWORD dwError = ::GetLastError();

	switch (dwError)
	{
	case ERROR_IO_PENDING:
       	::Sleep(iTimeToComplete);
        WaitFor_IO_OperationToComplete(iExtraTimeToComplete);
	break;

	default:
		ASSERT(0);	// Generally shouldn't reach here
		break;
	}
}