#include <util/wintypes.h>
#include <util/rapitypes.h>
#include <util/HiresTimer.h>
#include "streammultiplex.h"
#include "stringutils.h"
#include "debug.h"
#include <stdio.h>

const char* itsmethodnames[]= {
    "ITCallAddress",                // 00
    "ITColdboot",                   // 01
    "ITDialNumber",                 // 02
    "ITDiskProtect",                // 03
    "ITFindWindow",                 // 04
    "ITFlashROM",                   // 05
    "ITGetContext",                 // 06
    "ITGetForegroundWindow",        // 07
    "ITGetHandles",                 // 08
    "ITGetModuleList",              // 09
    "ITGetPartitionList",           // 0a
    "ITGetProcessHandle",           // 0b
    "ITGetProcessList",             // 0c
    "ITGetProcessUsageList",        // 0d
    "ITGetSTRGHandleList",          // 0e
    "ITGetStoreMgrList",            // 0f
    "ITGetSystemTime",              // 10
    "ITGetThreadTimes",             // 11
    "ITGetThreadUsageList",         // 12
    "ITGetVersion",                 // 13
    "ITGetWindowList",              // 14
    "ITGsmControl",                 // 15
    "ITLogDiskInfo",                // 16
    "ITReadDisk",                   // 17
    "ITReadGsmMemory",              // 18
    "ITReadProcessMemory",          // 19
    "ITReadSDCard",                 // 1a
    "ITReboot",                     // 1b
    "ITRegistryFlush",              // 1c
    "ITSDCardInfo",                 // 1d
    "ITSearchMemory",               // 1e
    "ITSendMessage",                // 1f
    "ITSendSms",                    // 20
    "ITSetSystemTime",              // 21
    "ITShowWindow",                 // 22
    "ITStartDebuggingProcess",      // 23
    "ITStopDebuggingProcess",       // 24
    "ITTFFSGetInfo",                // 25
    "ITTerminateProcess",           // 26
    "ITTerminateThread",            // 27
    "ITTraceWindow",                // 28
    "ITWaitForDebugEvent",          // 29
    "ITWaitForProcess",             // 2a
    "ITWriteDisk",                  // 2b
    "ITWriteProcessMemory",         // 2c
    "ITWriteSDCard",                // 2d
    "IT_POOM_AddRecord",            // 2e
    "IT_POOM_FindRecords",          // 2f
    "IT_POOM_GetRecord",            // 30
};
#ifdef _WIN32_WCE
#include "itsutils.h"
#include "DiskReader.h"
#include "PoomInterface.h"
#include "wintrace.h"
RAPIEXT*itsmethods[]= {
    (RAPIEXT*)&ITCallAddress, (RAPIEXT*)&ITColdboot, (RAPIEXT*)&ITDialNumber, (RAPIEXT*)&ITDiskProtect,
    (RAPIEXT*)&ITFindWindow, (RAPIEXT*)&ITFlashROM, (RAPIEXT*)&ITGetContext, (RAPIEXT*)&ITGetForegroundWindow,
    (RAPIEXT*)&ITGetHandles, (RAPIEXT*)&ITGetModuleList, (RAPIEXT*)&ITGetPartitionList,
    (RAPIEXT*)&ITGetProcessHandle, (RAPIEXT*)&ITGetProcessList, (RAPIEXT*)&ITGetProcessUsageList,
    (RAPIEXT*)&ITGetSTRGHandleList, (RAPIEXT*)&ITGetStoreMgrList, (RAPIEXT*)&ITGetSystemTime,
    (RAPIEXT*)&ITGetThreadTimes, (RAPIEXT*)&ITGetThreadUsageList, (RAPIEXT*)&ITGetVersion,
    (RAPIEXT*)&ITGetWindowList, (RAPIEXT*)&ITGsmControl, (RAPIEXT*)&ITLogDiskInfo, (RAPIEXT*)&ITReadDisk,
    (RAPIEXT*)&ITReadGsmMemory, (RAPIEXT*)&ITReadProcessMemory, (RAPIEXT*)&ITReadSDCard, (RAPIEXT*)&ITReboot,
    (RAPIEXT*)&ITRegistryFlush, (RAPIEXT*)&ITSDCardInfo, (RAPIEXT*)&ITSearchMemory, (RAPIEXT*)&ITSendMessage,
    (RAPIEXT*)&ITSendSms, (RAPIEXT*)&ITSetSystemTime, (RAPIEXT*)&ITShowWindow, (RAPIEXT*)&ITStartDebuggingProcess,
    (RAPIEXT*)&ITStopDebuggingProcess, (RAPIEXT*)&ITTFFSGetInfo, (RAPIEXT*)&ITTerminateProcess,
    (RAPIEXT*)&ITTerminateThread, (RAPIEXT*)&ITTraceWindow, (RAPIEXT*)&ITWaitForDebugEvent,
    (RAPIEXT*)&ITWaitForProcess, (RAPIEXT*)&ITWriteDisk, (RAPIEXT*)&ITWriteProcessMemory, (RAPIEXT*)&ITWriteSDCard,
    (RAPIEXT*)&IT_POOM_AddRecord, (RAPIEXT*)&IT_POOM_FindRecords, (RAPIEXT*)&IT_POOM_GetRecord,
};
#endif

class streamerror {
    HRESULT _hr;
    std::string _msg;
public:
    streamerror(HRESULT hr, const std::string& msg) : _hr(hr), _msg(msg) { }
    ~streamerror() { printf("ITSUTILS stream-ERROR: %08x: %s", _hr, _msg.c_str()); }
};
struct alloc_error {
    ~alloc_error() { printf("ITSUTILS alloc error\n"); }
};
template<typename T>
T *checkedalloc(size_t size)
{
#ifdef _WIN32
    T *p= (T*)LocalAlloc(LPTR, size);
#else
    T *p= (T*)malloc(size);
#endif
    if (p==NULL) {
        printf("ERROR in LocalAlloc(%08x)\n", (int)size);
        throw alloc_error();
    }
    return p;
}
void checkedfree(void *p)
{
#ifdef _WIN32
    LocalFree(p);
#else
    free(p);
#endif
}


streammx::streammx(IRAPIStream *stream) : _stream(stream), _apiversion(0)
{
    for(size_t i=0 ; i<sizeof(itsmethodnames)/sizeof(*itsmethodnames) ; i++)
    {
        _mth2id.insert(str2int_map::value_type(itsmethodnames[i], i));
    }
}
streammx::~streammx()
{
    if (_stream) {
        writedword(-1);
        writedword(0);
    }
    deletestream();
}
void streammx::deletestream()
{
    if (_stream) {
        _stream->Release();
        _stream= NULL;
    }
}

void streammx::writedata(const void *p, size_t n)
{
    size_t totalwritten=0;
    BYTE *pb= (BYTE*)p;
    while (_stream && totalwritten<n) {
        DWORD w=0;
        // limit max chunk to 64k, to avoid WSAENOBUFS
        DWORD want= std::min((size_t)0x10000, (size_t)(n-totalwritten));
        HiresTimer t;
        HRESULT hr=_stream->Write(pb+totalwritten, want, &w);
        if (hr) {
            printf("writedata(%d of %d) : error %08lx (t=%08lx(%+5d))\n", n-totalwritten, n, hr, int(t.getstamp()/1000), t.msecelapsed());
            throw streamerror(hr, "Write");
        }
        //debugt("writedata(%d of %d) : (t=%08lx->%08lx) : %s\n", n-totalwritten, n, t0, GetTickCount(), hexdump(pb+totalwritten, w).c_str());
        totalwritten+=w;
    }
}
void streammx::readdata(void *p, size_t n)
{
    size_t totalread=0;
    BYTE *pb= (BYTE*)p;
    while (_stream && totalread<n) {
        DWORD r=0;
        //DWORD t0= GetTickCount();
        HRESULT hr= _stream->Read(pb+totalread, n-totalread, &r);
        if (hr) {
            //debugt("readdata(%d of %d) : error %08lx (t=%08lx->%08lx)\n", n-totalread, n, hr, t0, GetTickCount());
            throw streamerror(hr, "Read");
        }
        //debugt("readdata(%d of %d) : (t=%08lx->%08lx) : %s\n", n-totalread, n, t0, GetTickCount(), hexdump(pb+totalread, r).c_str());
        totalread+=r;
    }
}
void streammx::writedword(DWORD dw)
{
    writedata(&dw, sizeof(dw));
}

DWORD streammx::readdword()
{
    DWORD dw;
    readdata(&dw, sizeof(dw));
    return dw;
}
HRESULT streammx::invoke(const std::string& methodname, DWORD insize, BYTE *inbuf, DWORD *poutsize, BYTE **outbuf)
{
    BYTE *out= NULL;
    DWORD result;
    try {
        int id;
        writedword(id= _mth2id[methodname]);
        writedword(insize);
        writedata(inbuf, insize);

        //debugt("sent request %03d:%ls (in=%d)\n", id, itsmethodnames[id], insize);

        result= readdword();
        *poutsize= readdword();
        out= checkedalloc<BYTE>(*poutsize);
        readdata(out, *poutsize);

        //debugt("received answer: out=%d, res=%08lx\n", *poutsize, result);
    }
    catch(...) {
        deletestream();
        result= 0x80090041;     // just made this error code up
    }
    if (outbuf)
        *outbuf= out;
    else if (out) {
        checkedfree(out);
        out=NULL;
    }

    return result;
}
DWORD streammx::getversion()
{
    if (_apiversion==0)
    {
        try {
        _apiversion= readdword();
        }
        catch(...)
        {
        }
    }
    return _apiversion;
}
#ifdef _WIN32_WCE
void streammx::service()
{
    BYTE *outbuf= NULL;
    try {
        int id= readdword();
        if (id==-1) {
            //printf("!itsutils - received EOF id: -1\n");
            deletestream();
            return;
        }
        DWORD insize= readdword();
        BYTE *inbuf= checkedalloc<BYTE>(insize);
        readdata(inbuf, insize);

        DWORD outsize=0;

        if (id<0 || id>=sizeof(itsmethods)/sizeof(*itsmethods)) {
            printf("!itsutils: invalid function id: %d ( max=%d )\n", id, sizeof(itsmethods)/sizeof(*itsmethods));
            writedword(0x80090007);
            writedword(0);
            return;
        }

        //debugt("invoking %03d:%ls (in=%d)\n", id, itsmethodnames[id], insize);

        // note: the invoked method is required to free inbuf
        DWORD result= itsmethods[id](insize, inbuf, &outsize, &outbuf, NULL);

        //debugt("result: out=%d, res=%08lx\n", outsize, result);

        writedword(result);
        writedword(outsize);
        if (outbuf==NULL) {
            outbuf= checkedalloc<BYTE>(outsize);
            memset(outbuf, 0, outsize);
        }
        writedata(outbuf, outsize);
    }
    catch(const char*msg) {
        printf("ERROR in itsutils: %s\n", msg);
        deletestream();
    }
    catch(...) {
        printf("ERROR in itsutils\n");
        deletestream();
    }
    if (outbuf)
        checkedfree(outbuf);
    outbuf= NULL;
    //printf("done with stream service()\n");
}
#endif
bool streammx::eof()
{
    return _stream==NULL;
}
void streammx::sendversion(DWORD version)
{
    try {
    writedword(version);
    }
    catch(...)
    {
    }
}