#include <windows.h>
#include "vectorutils.h"
#include "debug.h"
#include "cenk.h"
#include <vector>
#include <map>
#include <algorithm>
#include "kernelmisc.h"

// this module gets memory usage info per process, ideally the output would look like this:
//        heapuse heapfree  pages
// nk.exe ..      ..        ..      
// proc1  ..
// free  : total free
//
struct Section {
    MemBlock*blocks[512];
};
int GetPageUseCount(PHYSICAL_ADDRESS page)
{
    MEMORYINFO *mi= (MEMORYINFO *)KData.aInfo[KINX_MEMINFO];
    for (uint i=0 ; i<mi->cFi ; i++)
    {
        FREEINFO *fi= &mi->pFi[i];
        if (fi->paStart<=page && page<fi->paEnd) {
            BYTE *pm=fi->pUseMap;
            BYTE use= pm[(page-fi->paStart)/KData.aInfo[KINX_PAGESIZE]];
            if (use==0 || use==0xFF)
                return 1;
            return use;
        }
    }
    return 1;
}
DWORD CalcMemoryUsageForSection(int section)
{
    Section**slist= (Section**)KData.aInfo[KINX_SECTIONS];
    Section* s= section<64 ? slist[section] : (Section*)KData.aInfo[KINX_NKSECTION];
    if (s==0)
        return 0;

    DWORD pagesize= KData.aInfo[KINX_PAGESIZE];
    DWORD memusage=0;
    for (int block=0 ; block<512 ; block++)
    {
        MemBlock *pmb= s->blocks[block];
        if (pmb==(MemBlock*)0) continue;
        if (pmb==(MemBlock*)1) continue;
        DWORD *pages= pmb->aPages;
        for (int page=0 ; page<16 ; page++) {
            DWORD entry= pages[page];
            if (entry==~0xf) continue;
            if (entry==0) continue;

            // correct usage for shared pages
            memusage += pagesize/GetPageUseCount(entry&~0xfff);
        }
    }
    return memusage;
}
void calctotals()
{
    KernelMode _km;
    DWORD t0= GetTickCount();
    // maps physaddr -> list of sectionid's
    typedef std::map<DWORD,ByteVector> PhysSectionMap;
    PhysSectionMap phys;
    Section**slist= (Section**)KData.aInfo[KINX_SECTIONS];
    for (BYTE section=1 ; section<64  ; section++)
    {
        Section* s= slist[section];
        if (s==0)
            continue;
        for (int block=0 ; block<512 ; block++)
        {
            MemBlock *pmb= s->blocks[block];
            if (pmb==(MemBlock*)0) continue;
            if (pmb==(MemBlock*)1) continue;
            DWORD *pages= pmb->aPages;
            for (int page=0 ; page<16 ; page++) {
                DWORD entry= pages[page];
                if (entry==~0xf) continue;
                if (entry==0) continue;
                phys[entry&~0xfff].push_back(section);
            }
        }
    }
    DWORD t1= GetTickCount();
    typedef std::map<BYTE,DWORD> SectionTotalMap;
    SectionTotalMap stotal;
    for (PhysSectionMap::iterator i=phys.begin() ; i!=phys.end() ; ++i)
    {
        for (ByteVector::iterator j=(*i).second.begin() ; j!=(*i).second.end() ; ++j)
        {
            stotal[*j] += 0x1000 / (*i).second.size();
        }
    }
    DWORD t2= GetTickCount();
    debug("t0..t1 = %d  t1..t2 = %d\n", t1-t0, t2-t1);
    for (SectionTotalMap::iterator i=stotal.begin() ; i!=stotal.end() ; i++)
    {
        debug("%02x : %08lx\n", (*i).first, (*i).second);
    }
#if 0
    // .. output just for debugging
    for (BYTE section=1 ; section<64  ; section++)
    {
        Section* s= slist[section];
        if (s==0) {
            debug("%02x : %08lx\n", s);
            continue;
        }
        for (int block=0 ; block<512 ; block++)
        {
            MemBlock *pmb= s->blocks[block];
            if (pmb==(MemBlock*)0 || pmb==(MemBlock*)1) {
                debug("%02x : %08lx[%03x]=%08lx\n",
                        section, s, block, pmb);
                continue;
            }
            DWORD *pages= pmb->aPages;
            debug("%02x : %08lx[%03x]=%08lx:%08lx %02x %02x %04x %04x %04x %08lx:",
                    section, s, block, pmb, 
                    pmb->alk ,
                    pmb->cUses, pmb->flags, pmb->ixBase,
                    pmb->hPf  , pmb->cLocks,
                    pmb->aPages);
            for (int page=0 ; page<16 ; page++) {
                DWORD entry= pages[page];
                if (entry==~0xf || entry==0) {
                    debug(" %08lx[--]", entry);
                    continue;
                }
                debug(" %08lx[%02x]", entry, phys[entry&~0xfff].size());
            }
            debug("\n");
        }
    }
#endif
}
void calcfree()
{
    DWORD t0= GetTickCount();
    KernelMode _km;
    MEMORYINFO *pmi= (MEMORYINFO *)KData.aInfo[KINX_MEMINFO];
    debug("pmi=%08lx: kd=%08lx end=%08lx pfi=%08lx-%08lx\n", pmi, pmi->pKData, pmi->pKEnd, pmi->pFi, pmi->pFi+pmi->cFi);
    FreeInfo *pfi= pmi->pFi;
    DWORD totalfree=0;
    DWORD totalused=0;
    DWORD totalkernel=0;
    for (size_t i=0 ; i<pmi->cFi ; i++, pfi++)
    {
        DWORD nfree=0;
        DWORD nused=0;
        size_t npages= (pfi->paEnd-pfi->paStart)/0x1000;
        for(size_t i=0 ; i<npages ; i++)
        {
            if (pfi->pUseMap[i])
                nused++;
            else
                nfree++;
        }
        debug("%08lx-%08lx|%08lx  map=%08lx : free=%08lx used=%08lx\n", pfi->paStart, pfi->paEnd, pfi->paRealEnd, pfi->pUseMap, nfree*0x1000, nused*0x1000);
        totalfree += nfree;
        totalused += nused;
    }
    DWORD t1= GetTickCount();
    debug("t0..t1 = %d\n", t1-t0);
    debug("total: free=%08lx used=%08lx\n", totalfree*0x1000, totalused*0x1000);
}
#if 0
void findfuncdata(const DWORD *pbegin, const DWORD *pend, DWORD low, DWORD high, DwordVector& data)
{
    for (const DWORD *p= pbegin ; p<pend ; p++) {
        if (*p>=low && *p<high)
            data.push_back(*p);
    }
}
struct isnotheapptr {
    bool operator()(DWORD *ptr) {
        if (IsBadReadPtr(ptr, sizeof(DWORD)))
            return true;
        return *ptr!=0x50616548;
    }
};
void filter_non_heapptrs(DwordVector& list)
{
    DwordVector::iterator newend= std::remove_if(list.begin(), list.end(), isnotheapptr());
    list.erase(newend, list.end());
}
DWORD findheaplist()
{
    // todo: get these from module list.
    DWORD codebegin=0x03f00000;
    DWORD codeend=0x04000000;
    DWORD databegin=0x01f00000;
    DWORD dataend=0x02000000;

    DwordVector gph_list;
    findfuncdata((DWORD*)&GetProcessHeap, (DWORD*)codeend, databegin, dataend, gph_list);
    filter_non_heapptrs(gph_list);

    DwordVector ghs_list;
    findfuncdata((DWORD*)&GetHeapSnapshot, (DWORD*)codeend, databegin, dataend, ghs_list);
    filter_non_heapptrs(ghs_list);

    if (gph_list.size()!=1) {
        debug("expecting only one item in gph list\n");
    }
    if (ghs_list.size()!=2) {
        debug("expecting only two items in ghs list\n");
    }
    if (ghs_list[0]==gph_list[0])
        return ghs_list[1];
    else
        return ghs_list[0];
}
struct item*getnextitem(struct item* i)
{
    DWORD ptr=(DWORD)i;
    if (i->size<0) {
        ptr -= i->size;
    }
    else {
        ptr += i->size;
    }
    return (struct item*)i;
}
void heap_summary()
{
    DwordVector heapfree;
    DwordVector heapused;

    PROCESS *procarray= (PROCESS*)KData.aInfo[KINX_PROCARRAY];
    for (int i=0 ; i<32 ; i++) {
        if (procarray[i].dwVMBase) {
            for (struct heap *h= getvmheaplist(procarray[i].dwVMBase); h ; h=h->phpNext)
            {
                for (struct vaitem *v= h->pvaList ; v ; v=v->pvaFwd)
                {
                }
                for (struct region *r= &h->rgn ; r ; r= r->prgnNext)
                {
                    for (struct item *i= (struct item *)(r+1) ; i<r->pitLast ; i=getnextitem(i) )
                    {
                        if (i->size<0)
                            heapfree -= i->size;
                        else
                            heapused += i->size;
                    }
                }
            }
        }
    }
}
#endif
void CalcMemoryUsage(DWORD& dwMemoryUsed, DWORD& dwMemoryFree, DWORD& dwKernelMemory)
{
    MEMORYINFO *mi= (MEMORYINFO *)KData.aInfo[KINX_MEMINFO];
    dwKernelMemory= KData.aInfo[KINX_SYSPAGES] * KData.aInfo[KINX_PAGESIZE];
    dwMemoryUsed= 0;
    dwMemoryFree= 0;
    for (uint i=0 ; i<mi->cFi ; i++)
    {
        FREEINFO *fi= &mi->pFi[i];
        DWORD npages= (fi->paEnd - fi->paStart)/KData.aInfo[KINX_PAGESIZE];
        BYTE *pm=fi->pUseMap;
        for (uint j=0 ; j<npages ; j++, pm++)
            if (*pm>0 && *pm<255)
                dwMemoryUsed++;
            else
                dwMemoryFree++;
    }
    dwMemoryUsed *= KData.aInfo[KINX_PAGESIZE];
    dwMemoryFree *= KData.aInfo[KINX_PAGESIZE];
}