#include <windows.h>
#include "emulateinsn.h"
#include "insndecoder.h"
#include "ThreadContext.h"
#include "ProcessContext.h"
#include "bkptmanager.h"
#include "stringutils.h"

#include <pkfuncs.h>
#include <pwinbase.h>
// note: ~/sources/wince500/PRIVATE/WINCEOS/COREOS/NK/KERNEL/ARM/mdarm.c
//    contains ExceptionDispatch which invokes debug exceptions ( search /ExceptionCode\s*=\s*\w  )
//    calls UserDbgTrap from schedule.c
//    also called: 
//
// ~/sources/wince500/PRIVATE/WINCEOS/COREOS/NK/KERNEL/loader.c:   - inits some structs
// ~/sources/wince500/PRIVATE/WINCEOS/COREOS/NK/KERNEL/printf.c:   - OUTPUT_DEBUG_STRING_EVENT
// ~/sources/wince500/PRIVATE/WINCEOS/COREOS/NK/KERNEL/schedule.c: - the rest, see DbgrNotify...
//    contain the other debug events
//
struct win32_error {
public:
    std::string name;
    DWORD err;
    win32_error(const std::string& name) : err(GetLastError()), name(name) { }
    ~win32_error() { printf("%08lx - %s\n", err, name.c_str()); }
};

// copied from kernelmisc.h
struct HDATA {
    uint32_t res0[6];
    uint32_t pvObj;
    uint32_t res1[1];
};
struct PROCESS {
    uint32_t res0[3];
    uint32_t dwVMBase;
};
HDATA *cvHandle2HDataPtr(HANDLE h)
{
//    return (HDATA *)(KData.handleBase+((DWORD)h&HANDLE_ADDRESS_MASK));
    return (HDATA *)(0x80000000+((DWORD)h&0x1ffffffc));
}

class Win32ProcessContext : public ProcessContext {
    HANDLE _pid;
    HANDLE _heap;
    HINSTANCE _lib;
#ifdef _WIN32_WCE
    uint32_t _vmbase;
#endif
public:
    Win32ProcessContext(DWORD pid)
        : _pid((HANDLE)pid)
    {
        _lib= LoadLibrary(_T("coredll.dll"));
        if (_lib==NULL)
            throw win32_error("LoadLibrary(coredll)");

        printf("heapcreate=%p  &heapcreate=%p coredll.heapcreate=%p\n", HeapCreate, &HeapCreate, GetProcAddress(_lib, _T("HeapCreate")));
        CALLBACKINFO cb;
        cb.hProc= _pid;
        cb.pfn= (FARPROC)GetProcAddress(_lib, _T("HeapCreate"));
        cb.pvArg0= 0; // options
        _heap= (HANDLE)PerformCallBack4(&cb, 0x10000, 0, 0);
        if (_heap==NULL)
            throw win32_error("PerformCallBack4(HeapCreate)");
        printf("remoteheap=%08lx\n", _heap);

#ifdef _WIN32_WCE
        HDATA *hd= cvHandle2HDataPtr(_pid);  // kernelmisc.cpp
        PROCESS *pp=(PROCESS*)hd->pvObj;
        _vmbase= pp->dwVMBase;
#endif
    }
    virtual ~Win32ProcessContext()
    {
        CALLBACKINFO cb;
        cb.hProc= _pid;
        cb.pfn= (FARPROC)GetProcAddress(_lib, _T("HeapDestroy"));
        cb.pvArg0= _heap;
        PerformCallBack4(&cb, 0, 0, 0);
        _heap= 0;
        FreeLibrary(_lib);
    }
    virtual uint32_t readdword(uint32_t offset)
    {
        uint32_t dw;
        DWORD n;
        if (!ReadProcessMemory(_pid, (const void*)offset, (void*)&dw, sizeof(uint32_t), &n))
            throw "invalid pointer";
        return dw;
    }
    virtual void writedword(uint32_t offset, uint32_t value)
    {
        DWORD n;
        if (!WriteProcessMemory(_pid, (void*)offset, (void*)&value, sizeof(uint32_t), &n))
            throw "invalid pointer";
        printf("proc %08lx.word= %08x\n", offset, value);        
    }

    virtual uint8_t readbyte(uint32_t offset)
    {
        uint8_t dw;
        DWORD n;
        if (!ReadProcessMemory(_pid, (const void*)offset, (void*)&dw, sizeof(uint8_t), &n))
            throw "invalid pointer";
        return dw;

    }
    virtual void writebyte(uint32_t offset, uint8_t value)
    {
        DWORD n;
        if (!WriteProcessMemory(_pid, (void*)offset, (void*)&value, sizeof(uint8_t), &n))
            throw "invalid pointer";
        printf("proc %08lx.byte= %02x\n", offset, value);
    }
    uint32_t alloc(uint32_t size)
    {
        uint32_t remote_p= (uint32_t)RemoteHeapAlloc(_pid, _heap, 0, size);
        printf("ralloc(0x%x) -> %08lx\n", size, remote_p);
        return remote_p;
    }
    void free(uint32_t ptr)
    {
        RemoteHeapFree(_pid, _heap, 0, (void*)ptr);
    }
    virtual uint32_t pid() const
    {
        return (uint32_t)_pid;
    }
    virtual uint32_t canonicaladdress(uint32_t a)
    {
#ifdef _WIN32_WCE
        if ((a&0xfe000000)==0)
            a+=_vmbase;
#endif
        // todo: think of something for win32: like return uint64_t

        return a;
    }
};
class Win32ProcessManager {
    typedef std::map<uint32_t,ProcessContext_ptr> procmap_t;
    procmap_t _procs;
public:
    ProcessContext_ptr forpid(uint32_t pid)
    {
        procmap_t::iterator i= _procs.find(pid);
        if (i==_procs.end())
        {
            std::pair<procmap_t::iterator,bool> j= _procs.insert(procmap_t::value_type(pid, ProcessContext_ptr(new Win32ProcessContext(pid))));
            return (*j.first).second;
        }
        return (*i).second;
    }
};
Win32ProcessManager pmgr;
class Win32ThreadContext : public ThreadContext {
    DWORD _pid, _tid;
public:
    Win32ThreadContext(DWORD pid, DWORD tid)
        : _pid(pid), _tid(tid)
    {
        CONTEXT ctx;
        // note: not handling floating point CONTEXT_FLOATING_POINT
        ctx.ContextFlags = CONTEXT_CONTROL | CONTEXT_INTEGER; 
        if (!GetThreadContext((HANDLE)tid, &ctx))
            throw win32_error("GetThreadContext");
        DWORD *regs= &ctx.R0;
        for (int i=0 ; i<16 ; i++)
            setreg(i, regs[i]);
        setcpsr(ctx.Psr);

    }
    virtual ~Win32ThreadContext()
    {
        CONTEXT ctx;
        DWORD *regs= &ctx.R0;
        for (int i=0 ; i<16 ; i++)
            regs[i]= getreg(i);
        ctx.Psr= getcpsr();
        ctx.ContextFlags = CONTEXT_CONTROL | CONTEXT_INTEGER;
        if (!SetThreadContext((HANDLE)_tid, &ctx))
            throw win32_error("SetThreadContext");
    }
    virtual uint32_t tid() const
    {
        return _tid;
    }
    virtual uint32_t pid() const
    {
        return _pid;
    }
    virtual ProcessContext_ptr proc()
    {
        return pmgr.forpid(_pid);
    }
    void dumpstack()
    {
        uint32_t sp= getreg(13);
        uint32_t tos= (sp|0xFFF)+1;
        ProcessContext_ptr p= proc();
        printf("%08lx: ", sp);
        while (sp<tos) {
            printf(" %08lx", p->readdword(sp));
            sp+=4;
        }
        printf("\n");
    }
};


void attach(DWORD pid)
{
    if (!DebugActiveProcess(pid))
        throw win32_error("DebugActiveProcess");
}
void loadexe(const std::string& exename, const std::string& cmdline, bool bSuspended, DWORD& pid, DWORD& tid)
{
    PROCESS_INFORMATION pi;

    if (!CreateProcess(ToWString(exename).c_str(), ToWString(cmdline).c_str(), NULL, NULL, 0, DEBUG_PROCESS|(bSuspended?CREATE_SUSPENDED:0), NULL, NULL, NULL, &pi))
        throw win32_error("CreateProcess");
    pid= pi.dwProcessId;
    tid= pi.dwThreadId;
}
void resume(DWORD tid)
{
    printf("resuming thread %08lx\n", tid);
    ResumeThread((HANDLE)tid);
    printf("resumed thread %08lx\n", tid);
}

// EXCEPTION_RECORD  is defined in z:\home\sources\wince500\PUBLIC\COMMON\SDK\INC\winnt.h
bool handle_exception     (breakpoint_manager& bpm, DWORD pid, DWORD tid, EXCEPTION_RECORD         *p)
{
    printf("debug %08lx:%08lx exception code %d, flags=%x, addr=%08lx nparm=%d\n", pid, tid,
            p->ExceptionCode, p->ExceptionFlags, p->ExceptionAddress, p->NumberParameters);
    if (p->ExceptionCode==EXCEPTION_BREAKPOINT) {
        printf("debug %08lx:%08lx exception breakpoint\n", pid, tid);
        Win32ThreadContext td(pid,tid);
        try {
            bpm.handle_bpt(td);
            return true;
        }
        catch(...)
        {
        }
    }
    else {
        printf("debug %08lx:%08lx exception code %d\n", p->ExceptionCode, pid, tid);
    }
    return false;
}
bool handle_create_thread (breakpoint_manager& bpm, DWORD pid, DWORD tid, CREATE_THREAD_DEBUG_INFO *p)
{
    printf("debug %08lx:%08lx createthread h=%08lx, tls=%08lx start=%08lx\n", pid, tid, p->hThread, p->lpThreadLocalBase, p->lpStartAddress);
    return false;
}
class singlestep_action : public breakpoint_action {
public:
    singlestep_action(breakpoint_manager& bpm, ThreadContext& td) : breakpoint_action(bpm,td) { }
    virtual void action() const
    {
        // emulateinsn
        //   ... todo: think of way to do this:
        //
        //  ss bpt
        //  ss action
        //  --- resume ---
        //  bpt in emu
        //  set new ss bpt at new 'pc'
        //  --- resume ---
        //  ss bpt
        //  ss action
        //
        //
        //  -> the 'ss' is not nescesary
        //
        //   ... todo: get 'original bytes value for memory from somewhere
    }
    virtual int type() const { return BPT_USERREQUEST; }
    virtual std::string description() const
    {
        return "SINGLESTEP";
    }
};

std::string debugname(void *p, bool fUnicode, int length=-1)
{
    if (length==-1) {
        if (fUnicode)
            return ToString((WCHAR*)p);
        else
            return ToString((char*)p);
    }
    else {
        if (fUnicode)
            return ToString(std::wstring((WCHAR*)p, length));
        else
            return ToString(std::string((char*)p, length));
    }
}
// other debug strings are in z:\home\sources\wince500\PUBLIC\COMMON\SDK\INC\winbase.h
bool handle_create_process(breakpoint_manager& bpm, DWORD pid, DWORD tid, CREATE_PROCESS_DEBUG_INFO*p)
{
    printf("debug %08lx:%08lx createprocess hfile=%08lx hproc=%08lx hthread=%08lx base=%08lx\n"
            "       dbginfo:%08lx/%08lx  tls=%08lx start=%08lx name=%08lx(u=%d):'%s'\n", pid, tid,
        p->hFile, p->hProcess, p->hThread, p->lpBaseOfImage,
        p->dwDebugInfoFileOffset, p->nDebugInfoSize, p->lpThreadLocalBase, p->lpStartAddress, p->lpImageName, p->fUnicode,
        debugname(p->lpImageName, p->fUnicode).c_str());
//    Win32ThreadContext td(pid,tid);
//    bpm.reg_bpt(VirtualAddress(td.proc(), (uint32_t)p->lpStartAddress), new singlestep_action(bpm, td));
    return false;
}
bool handle_exit_thread   (breakpoint_manager& bpm, DWORD pid, DWORD tid, EXIT_THREAD_DEBUG_INFO   *p)
{
    printf("debug %08lx:%08lx exitthread code=%d\n", pid, tid, p->dwExitCode);
    return false;
}
bool handle_exit_process  (breakpoint_manager& bpm, DWORD pid, DWORD tid, EXIT_PROCESS_DEBUG_INFO  *p)
{
    printf("debug %08lx:%08lx exitprocess code=%d\n", pid, tid, p->dwExitCode);
    return false;
}
bool handle_load_dll      (breakpoint_manager& bpm, DWORD pid, DWORD tid, LOAD_DLL_DEBUG_INFO      *p)
{
    printf("debug %08lx:%08lx loaddll hfile=%08lx base=%08lx dbginfo=%08lx/%08lx name=%08lx(u=%d):'%s'\n", pid, tid,
        p->hFile, p->lpBaseOfDll, p->dwDebugInfoFileOffset, p->nDebugInfoSize, p->lpImageName, p->fUnicode,
        debugname(p->lpImageName, p->fUnicode).c_str());
    return false;
}
bool handle_unload_dll    (breakpoint_manager& bpm, DWORD pid, DWORD tid, UNLOAD_DLL_DEBUG_INFO    *p)
{
    printf("debug %08lx:%08lx unloaddll base=%08lx\n", pid, tid, p->lpBaseOfDll);
    return false;
}
bool handle_string        (breakpoint_manager& bpm, DWORD pid, DWORD tid, OUTPUT_DEBUG_STRING_INFO *p)
{
    printf("debug %08lx:%08lx string str=%08lx(u=%d, l=%d):'%s'\n", pid, tid, p->lpDebugStringData, p->fUnicode, p->nDebugStringLength,
            debugname(p->lpDebugStringData, p->fUnicode, p->nDebugStringLength).c_str());
    return false;
}
bool handle_rip           (breakpoint_manager& bpm, DWORD pid, DWORD tid, RIP_INFO                 *p)
{
    printf("debug %08lx:%08lx RIP err=%08lx type=%08lx\n", pid, tid, p->dwError, p->dwType);
    return false;
}


bool dispatch_event(breakpoint_manager& bpm, DEBUG_EVENT *p)
{
    printf("debug event %d\n", p->dwDebugEventCode);
    Win32ThreadContext td(p->dwProcessId, p->dwThreadId);
    td.dumpstack();
    switch (p->dwDebugEventCode)
    {
        case EXCEPTION_DEBUG_EVENT     : return handle_exception     (bpm, p->dwProcessId, p->dwThreadId, &p->u.Exception.ExceptionRecord  );
        case CREATE_THREAD_DEBUG_EVENT : return handle_create_thread (bpm, p->dwProcessId, p->dwThreadId, &p->u.CreateThread     );
        case CREATE_PROCESS_DEBUG_EVENT: return handle_create_process(bpm, p->dwProcessId, p->dwThreadId, &p->u.CreateProcessInfo);
        case EXIT_THREAD_DEBUG_EVENT   : return handle_exit_thread   (bpm, p->dwProcessId, p->dwThreadId, &p->u.ExitThread       );
        case EXIT_PROCESS_DEBUG_EVENT  : return handle_exit_process  (bpm, p->dwProcessId, p->dwThreadId, &p->u.ExitProcess      );
        case LOAD_DLL_DEBUG_EVENT      : return handle_load_dll      (bpm, p->dwProcessId, p->dwThreadId, &p->u.LoadDll          );
        case UNLOAD_DLL_DEBUG_EVENT    : return handle_unload_dll    (bpm, p->dwProcessId, p->dwThreadId, &p->u.UnloadDll        );
        case OUTPUT_DEBUG_STRING_EVENT : return handle_string        (bpm, p->dwProcessId, p->dwThreadId, &p->u.DebugString      );
        case RIP_EVENT                 : return handle_rip           (bpm, p->dwProcessId, p->dwThreadId, &p->u.RipInfo          );
        default:
                                         printf("unknown debug event: %d\n", p->dwDebugEventCode);
    }

    return false;
}

void debug_loop(DWORD pid, DWORD tid)
{
    printf("starting debugloop for %08lx:%08lx\n", pid, tid);
    ProcessContext_ptr proc= pmgr.forpid(pid);
    codesnip_manager code;
    breakpoint_manager bpm(code);

    InstructionEmulator emu(bpm);

    // note: when creating suspended: initially R15 = MainThreadBaseFunc
    //       when creating nonsuspended: the initial R15== the R15 of the first createprocess debug event
    Win32ThreadContext td(pid, tid);
    td.dumpstack();
/*
    uint8_t insnbytes[4];
    printf("start ip=%08lx\n", td.getreg(15));
    proc->readmemory(td.getreg(15), 4, insnbytes);
    instruction insn(insnbytes, 4);
    emu.emulate_insn(td, insn);
*/
    resume(tid);

    printf("processing debug events\n");
    while (1)
    {
        DEBUG_EVENT ev;
        if (WaitForDebugEvent(&ev, 1000))
        {
            bool bHandled= dispatch_event(bpm, &ev);
            // if td.getpc is in proc space -> emulate_insn
            if (!ContinueDebugEvent(ev.dwProcessId, ev.dwThreadId, bHandled ? DBG_CONTINUE : DBG_EXCEPTION_NOT_HANDLED ))
                throw win32_error("ContinueDebugEvent");
        }
        else {
            //printf("no event\n");
        }
    }
}

void start()
{
    DWORD pid, tid;
    loadexe("/windows/test.exe", "", true, pid, tid);
    printf("loaded test.exe\n");
    debug_loop(pid, tid);
}



int WINAPI WinMain(HINSTANCE hInstance,
                   HINSTANCE hPrevInstance,
                   LPTSTR    lpCmdLine,
                   int       nCmdShow)
{
#ifdef _WIN32_WCE
        SetKMode(TRUE);
        SetProcPermissions(0xFFFFFFFF);
#endif
    try {
    start();
    }
    catch(const char*msg)
    {
        printf("ERROR: %s\n", msg);
    }
    catch(...)
    {
        printf("singlestepper EXCEPTION\n");
    }
    return 0;
}