#include <windows.h>
#include "wintrace.h"
#include "debug.h"
#include "stringutils.h"
#include <map>


struct WindowInfo {
    WNDPROC wndproc;
    HWND hWnd;
    bool has_cs;
    CRITICAL_SECTION cs;
};
typedef std::map<HWND,WindowInfo> WndInfoMap;

CRITICAL_SECTION cs_wi;
HINSTANCE hThisLib;
bool has_cs_wi= false;
WndInfoMap windowlist;

WindowInfo* findWindow(HWND hWnd)
{
    EnterCriticalSection(&cs_wi);
    WndInfoMap::iterator wii=  windowlist.find(hWnd);
    if (wii==windowlist.end()) {
        std::pair<WndInfoMap::iterator,bool> pi= windowlist.insert(WndInfoMap::value_type(hWnd,WindowInfo()));
        if (pi.second) {
            //debug("new window added to list\n");
        }
        wii= pi.first;
    }
    if (!(*wii).second.has_cs) {
        (*wii).second.has_cs= true;
        InitializeCriticalSection(&(*wii).second.cs);
    }
    LeaveCriticalSection(&cs_wi);
    return &((*wii).second);
}
void removeWindow(HWND hWnd)
{
    EnterCriticalSection(&cs_wi);
    WndInfoMap::iterator wii=  windowlist.find(hWnd);
    if (wii==windowlist.end())
        return;

    if ((*wii).second.has_cs) {
        (*wii).second.has_cs= false;
        DeleteCriticalSection(&(*wii).second.cs);
    }
    windowlist.erase(wii);
    LeaveCriticalSection(&cs_wi);
}
void trymemdump(const char*tag, const BYTE*ofs)
{
    if (!IsBadReadPtr(ofs, 64)) {
        if ((DWORD(ofs)&0xf)==0 && !IsBadReadPtr(ofs-8, 8+16)) {
            struct heapitem;
            struct heaphead {
                DWORD dwSig;
                DWORD pvaList;
                heaphead *phpNext;
                DWORD flOptions;
                DWORD cbMaximum;
                DWORD pfnAlloc;
                DWORD pfnFree;
                CRITICAL_SECTION cs;
            };
            struct heapregion {
                DWORD pitFree;
                heapregion* prgnLast;
                DWORD cbMaxFree;
                heapitem *pitLast;
                heaphead *owner;
                heapregion* prgnNext;
                DWORD dwRgnData;
                DWORD pad;
            };
            struct heapitem {
                LONG itemsize;
                heapregion* region;
            };
            heapitem *hdr= (heapitem*)(ofs-8);
            if (!IsBadReadPtr(hdr->region, sizeof(heapregion))) {
                if (!IsBadReadPtr(hdr->region->owner, sizeof(heaphead))) {
                    if (hdr->region->owner->dwSig==0x50616548) {
                        if (hdr->itemsize>0) {
                            debug("%s: [heap]%08x: %s\n", tag, ofs, hexdump(ofs, (hdr->itemsize-8)/4, 4).c_str());
                            return;
                        }
                    }
                }
            }
        }
        debug("%s: %08x: %s\n", tag, ofs, hexdump(ofs, 64).c_str());
    }
}

LRESULT spyproc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
{
    WindowInfo *wi= findWindow(hWnd);

    // todo handle special cases of extra data

    DWORD t0= GetTickCount();

    LRESULT lRes= CallWindowProcW(wi->wndproc, hWnd, message, wParam, lParam);

    DWORD t1= GetTickCount();

    // todo handle special cases of extra returned data

    debugt("%08lx  wnd%08lx %08lx %08lx %08lx  %08lx-%08lx(%+3d) -> %08lx\n", 
                GetCurrentThreadId(), hWnd, message, wParam, lParam, t0, t1, t1-t0, lRes);

    trymemdump("wparam", (const BYTE*)wParam);
    trymemdump("lparam", (const BYTE*)lParam);

    return lRes;
}
bool AddTrace(HWND hWnd)
{
    WindowInfo *wi= findWindow(hWnd);

    debug("addtrace(%08lx): %08lx\n", hWnd, wi->hWnd);

    EnterCriticalSection(&(wi->cs));

    wi->wndproc= (WNDPROC)GetWindowLong(hWnd, GWL_WNDPROC);
    wi->hWnd= hWnd;

    WNDPROC oldproc1= (WNDPROC)SetWindowLong(hWnd, GWL_WNDPROC, (LONG)&spyproc);

    if (((DWORD)oldproc1^(DWORD)wi->wndproc)&0x01ffffff) {
        debug("NOTE: procchange detected during spy install: first: %08lx, then: %08lx\n", wi->wndproc, oldproc1);
        wi->wndproc= oldproc1;
    }
    LeaveCriticalSection(&(wi->cs));

    return true;
}
bool DelTrace(HWND hWnd)
{
    WindowInfo *wi= findWindow(hWnd);
    debug("deltrace(%08lx): %08lx\n", hWnd, wi->hWnd);
    EnterCriticalSection(&(wi->cs));
    WNDPROC curproc= (WNDPROC)SetWindowLong(hWnd, GWL_WNDPROC, (LONG)wi->wndproc);
    if (curproc && (((DWORD)curproc^(DWORD)&spyproc)&0x01ffffff)!=0) {
        // curproc==NULL means that the window was destroyed.
        debug("NOTE: procchange detected during spying: spy=%08lx cur=%08lx\n", &spyproc, curproc);
    }
    wi->wndproc= NULL;
    LeaveCriticalSection(&(wi->cs));
    removeWindow(hWnd);

    return true;
}
bool AddProcessTrace(HANDLE hProc)
{
    DWORD*pheap= (DWORD*)0x7c070060;
    while (pheap[-2]) {
        int size= pheap[-2];
        if (size<0) {
            size= -size;
        }
        else {
            if (pheap[3]==0x574e4457 && pheap[0x22]==(DWORD)hProc) {
                AddTrace((HWND)pheap);
            }
        }
        pheap += size/sizeof(DWORD);
    }
    return true;
}
bool DelProcessTrace(HANDLE hProc)
{
    DWORD*pheap= (DWORD*)0x7c070060;
    while (pheap[-2]) {
        int size= pheap[-2];
        if (size<0) {
            size= -size;
        }
        else {
            if (pheap[3]==0x574e4457 && pheap[0x22]==(DWORD)hProc) {
                DelTrace((HWND)pheap);
            }
        }
        pheap += size/sizeof(DWORD);
    }
    return true;
}
void StartTrace(IRAPIStream *pStream)
{
    if (!has_cs_wi) {
        debug("init cswi\n");
        InitializeCriticalSection(&cs_wi);
        hThisLib= LoadLibrary(_T("itsutils.dll"));
        has_cs_wi= true;
    }
    else {
        debug("cswi already initialized\n");
    }
}
void StopTrace()
{
    debug("deleting traces\n");
    while (!windowlist.empty())
    {
        WndInfoMap::iterator i= windowlist.begin();
        DelTrace((*i).first);
    }
    if (has_cs_wi) {
        DeleteCriticalSection(&cs_wi);
        has_cs_wi= false;
        FreeLibrary(hThisLib);
        debug("deleted cswi\n");
    }
    else {
        debug("cswi not there\n");
    }
}


ITSUTILS_API HRESULT STDAPICALLTYPE ITTraceWindow(
        DWORD cbInput, TraceWindowParams *pbInput,
        DWORD *pcbOutput, TraceWindowResult **ppbOutput,
        IRAPIStream *pStream)
{
    *pcbOutput= sizeof(TraceWindowResult);
    TraceWindowResult *pOut= *ppbOutput= (TraceWindowResult*)LocalAlloc(LPTR, *pcbOutput);

    //debug("ITTraceWindow+(%08lx)\n", pbInput);
    if (pbInput==NULL)
        return 0;
    switch (pbInput->cmd)
    {
        case WND_START_TRACE:
            StartTrace(pStream);
            DebugSetLogfile("winmsgtrace.log");
            break;
        case WND_STOP_TRACE :
            StopTrace();
            break;
        case WND_ADD_WINDOW :
            AddTrace((HWND)pbInput->hWnd);
            break;
        case WND_DEL_WINDOW :
            DelTrace((HWND)pbInput->hWnd);
            break;
        case WND_ADD_PROCESS :
            AddProcessTrace((HANDLE)pbInput->hWnd);
            break;
        case WND_DEL_PROCESS :
            DelProcessTrace((HANDLE)pbInput->hWnd);
            break;

    }
    //debug("ITTraceWindow-\n");
    return 0;
}