#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "device.h"
#include "mem.h"
#include "mem_seg.h"
#include "runner.h"
#include "addrs.h"


uint64_t memSegToGlobal(device_specs_t* spec, uint8_t seg, uint64_t localaddr)
{
    uint64_t offset = 0;
    for (uint8_t i = 0; i < seg; i++)
    {
        offset += spec->memSpecs[i]->len * spec->memSpecs[i]->wordLen;
    }
    return offset + localaddr;
}


void freeDevMem(device_mem_t* devMem)
{
    free(devMem->memsegShifts);
    free(devMem->rawCells);
    free(devMem->memreadCellAddrs);
    free(devMem->memwriteCellAddrs);
    free(devMem->memwriteCellSegments);
    free(devMem->memwriteValues[0]);
    free(devMem->memwriteValues);
    free(devMem->cells);
    free(devMem->smartAddrReadHandlers);
    free(devMem->smartAddrWriteHandlers);
    free(devMem);
}

void freeDevSpec(void* _specs)
{
    device_specs_t* specs = _specs;
    for (uint8_t i = 0; i < specs->memSpecsCount; i++)
    {
        free(specs->memSpecs[i]);
    }
    free(specs->memSpecs);
    free(specs->executableSegments);
    free(specs);
}


device_mem_t* genDevMem(device_specs_t* devSpec, char* errbuf)
{
    if (devSpec->memSpecsCount < 2)
    {
        sprintf(errbuf, "invalid amount of mem specs: %u", devSpec->memSpecsCount);
        return NULL;
    }

    device_mem_t* devMem = (device_mem_t*)calloc(1, sizeof(device_mem_t));
    if (devMem == NULL)
    {
        sprintf(errbuf, "unable to allocate dev memory struct");
        return NULL;
    }

    uint64_t memTotalSize = 0;
    for (uint8_t i = 0; i < devSpec->memSpecsCount; i++)
    {
        size_t tmp = devSpec->memSpecs[i]->start + devSpec->memSpecs[i]->len * devSpec->memSpecs[i]->wordLen;
        if (memTotalSize < tmp)
        {
            memTotalSize = tmp;
        }
    }

    void* rawCells = (void*)calloc(memTotalSize, sizeof(void*));
    if (rawCells == NULL)
    {
        sprintf(errbuf, "unable to allocate raw memory buf");
        free(devMem);
        return NULL;
    }

    for (size_t i = 0; i < memTotalSize; i++)
    {
        ((uint8_t*)rawCells)[i] = 0;
    }

    devMem->memsegShifts = calloc(devSpec->memSpecsCount, sizeof(uint64_t));

    if(devMem->memsegShifts == NULL)
    {
        sprintf(errbuf, "unable to allocate segment shift buffers");
        free(rawCells);
        free(devMem);
        return NULL;
    }

    void** cells = calloc(devSpec->memSpecsCount, sizeof(void*));

    if (cells == NULL)
    {
        sprintf(errbuf, "unable to allocate segment pointers");
        free(rawCells);
        free(devMem->memsegShifts);
        free(devMem);
        return NULL;
    }

    for (uint8_t i = 0; i < devSpec->memSpecsCount; i++)
    {
        cells[i] = (void*)((size_t)rawCells + (size_t)devSpec->memSpecs[i]->start);
    }

    char** cellNames = calloc(devSpec->memSpecsCount, sizeof(char*));

    if(cellNames == NULL)
    {
        sprintf(errbuf, "unable to allocate segment names map");
        free(devMem->memsegShifts);
        free(devMem);
        free(rawCells);
        free(cells);
        return NULL;
    }

    for(size_t i = 0; i < devSpec->memSpecsCount; i++)
    {
        cellNames[i] = devSpec->memSpecs[i]->name;
    }

    uint64_t* memreadCellAddrs = calloc(MEM_ACCESS_INTERCEPT_BUF_SIZE, sizeof(uint64_t));
    if (memreadCellAddrs == NULL)
    {
        sprintf(errbuf, "unable to allocate read interception addrs");
        free(devMem->memsegShifts);
        free(devMem);
        free(rawCells);
        free(cells);
        free(cellNames);
        return NULL;
    }
    uint64_t* memwriteCellAddrs = calloc(MEM_ACCESS_INTERCEPT_BUF_SIZE, sizeof(uint64_t));

    if (memwriteCellAddrs == NULL)
    {
        sprintf(errbuf, "unable to allocate write interception addrs");
        free(devMem->memsegShifts);
        free(devMem);
        free(rawCells);
        free(memreadCellAddrs);
        free(cells);
        free(cellNames);
        return NULL;
    }

    uint8_t* memwriteCellSegments = calloc(MEM_ACCESS_INTERCEPT_BUF_SIZE, sizeof(uint8_t));

    if(memwriteCellSegments == NULL)
    {
        sprintf(errbuf, "unable to allocate write interception addrs");
        free(devMem->memsegShifts);
        free(devMem);
        free(rawCells);
        free(memreadCellAddrs);
        free(cells);
        free(cellNames);
        free(memwriteCellAddrs);
        return NULL;
    }

    void** memwriteValues = calloc(MEM_ACCESS_INTERCEPT_BUF_SIZE, sizeof(void*));

    if(memwriteValues == NULL)
    {
        sprintf(errbuf, "unable to allocate write interception addrs");
        free(devMem->memsegShifts);
        free(devMem);
        free(rawCells);
        free(memreadCellAddrs);
        free(cells);
        free(cellNames);
        free(memwriteCellAddrs);
        free(memwriteCellSegments);
        return NULL;
    }

    uint64_t* memwriteValuesContainers = calloc(MEM_ACCESS_INTERCEPT_BUF_SIZE, sizeof(uint64_t));
    if(memwriteValuesContainers == NULL)
    {
        sprintf(errbuf, "unable to allocate write interception addrs");
        free(devMem->memsegShifts);
        free(devMem);
        free(rawCells);
        free(memreadCellAddrs);
        free(cells);
        free(cellNames);
        free(memwriteCellAddrs);
        free(memwriteCellSegments);
        free(memwriteValues);
        return NULL;
    }

    for(size_t i = 0; i < 64; i++)
    {
        memwriteValues[i] = &memwriteValuesContainers[i];
    }

    uint64_t smartAddrReadMask = 0;
    uint64_t smartAddrWriteMask = 0;
    mem_h_read_handler* smartAddrReadHandlers = calloc(memTotalSize, sizeof(mem_h_read_handler));


    if (smartAddrReadHandlers == NULL)
    {
        sprintf(errbuf, "unable to allocate read interception handlers");
        free(devMem->memsegShifts);
        free(devMem);
        free(rawCells);
        free(memreadCellAddrs);
        free(cells);
        free(cellNames);
        free(memwriteCellAddrs);
        free(memwriteCellSegments);
        free(memwriteValues);
        free(memwriteValuesContainers);
        return NULL;
    }

    mem_h_write_handler* smartAddrWriteHandlers = calloc(memTotalSize, sizeof(mem_h_write_handler));


    if (smartAddrWriteHandlers == NULL)
    {
        sprintf(errbuf, "unable to allocate write interception handlers");
        free(smartAddrReadHandlers);
        free(devMem->memsegShifts);
        free(devMem);
        free(rawCells);
        free(memreadCellAddrs);
        free(cells);
        free(cellNames);
        free(memwriteCellAddrs);
        free(memwriteCellSegments);
        free(memwriteValues);
        free(memwriteValuesContainers);
        return NULL;
    }


    for(uint64_t i = 0; i < memTotalSize; i++) 
    {
        smartAddrReadHandlers[i].func = default_addr_read_handler;
        smartAddrReadHandlers[i].ident = 0;

        smartAddrWriteHandlers[i].func = default_addr_write_handler;
        smartAddrWriteHandlers[i].ident = 0;
    }


    for(uint64_t i = 0; i < memTotalSize; i++) 
    {
        if((i & smartAddrReadMask) == smartAddrReadMask)
        {
            smartAddrReadHandlers[i].func = default_addr_read_handler;
        }
    }


    if (devSpec->smartReadSpecsCount > 0)
    {
        for(uint64_t i = 0; i < devSpec->smartReadSpecsCount; i++)
        {
            smart_read_spec_t t = devSpec->smartReadSpecs[i];
            uint64_t addr;
            if (t.addrType == SMART_ADDR_TYPE_GLOBAL)
            {
                addr = t.addr;
            }
            else
            {
                addr = memSegToGlobal(devSpec, t.segno, t.localAddr);
            }
            smartAddrReadHandlers[addr].func = *t.handler;
            smartAddrReadHandlers[addr].ident = t.ident;
            smartAddrReadMask |= addr;
        }
    }


    if (devSpec->smartWriteSpecsCount > 0)
    {
        for(uint64_t i = 0; i < devSpec->smartWriteSpecsCount; i++)
        {
            smart_write_spec_t t = devSpec->smartWriteSpecs[i];

            uint64_t addr;
            if (t.addrType == SMART_ADDR_TYPE_GLOBAL)
            {
                addr = t.addr;
            }
            else
            {
                addr = memSegToGlobal(devSpec, t.segno, t.localAddr);
            }
            smartAddrWriteHandlers[addr].func = *t.handler;
            smartAddrWriteHandlers[addr].ident = t.ident;
            smartAddrWriteMask |= t.addr;
        }
    }


    devMem->cells = cells;
    devMem->rawCells = rawCells;
    devMem->memreadCellAddrs = memreadCellAddrs;
    devMem->memwriteCellAddrs = memwriteCellAddrs;
    devMem->memwriteCellSegments = memwriteCellSegments;
    devMem->memwriteValues = memwriteValues;
    devMem->memreadLen = 0;
    devMem->memwriteLen = 0;
    devMem->smartAddrReadMask = smartAddrReadMask;
    devMem->smartAddrReadHandlers = smartAddrReadHandlers;
    devMem->smartAddrWriteMask = smartAddrWriteMask;
    devMem->smartAddrWriteHandlers = smartAddrWriteHandlers;
    devMem->memsegNames = cellNames;

    for(uint8_t i = 0; i < devSpec->memSpecsCount; i++)
    {
        devMem->memsegShifts[i] = memSegToGlobal(devSpec, i, 0);
    }

    return devMem;
}


uint8_t makeDeviceTick(device_public_context_t* devContext)
{
    device_info_t* devInfo = (device_info_t*)devContext->deviceInfo;
    return makeTick(devInfo->deviceMem);
}

device_info_t* initSpecs(device_specs_t* specs, char* errbuf)
{
    device_info_t* devInfo = calloc(1, sizeof(device_info_t));
    if (devInfo == NULL)
    {
        sprintf(errbuf, "unable to allocate dev info");
        return NULL;
    }

    char genErrBuf[200];
    device_mem_t* devMem = genDevMem(specs, genErrBuf);

    if (devMem == NULL)
    {
        sprintf(errbuf, "unable to generate device memory: %s", genErrBuf);
        free(devInfo);
        return NULL;
    }

    devInfo->specs = specs;
    devInfo->deviceMem = devMem;

    return devInfo;
}

void error(const char* err, char* errbuf)
{
    strcpy(errbuf, err);
}

device_specs_t* parseSpecsFromConfig(const conf_dev_t* devConf, char* errbuf)
{
    conf_mem_seg_t** segments = devConf->memConf->memSegConfs;
    device_specs_t* specs = calloc(1, sizeof(device_specs_t));
    if (specs == NULL)
    {
        error("unable to allocate mem specs struct", errbuf);
        return NULL;
    }


    memseg_metadata_t requiredSegments[] = MEMSEG_DEFINES;
    uint8_t* requiredSegmentsFoundMap = (uint8_t*)calloc(sizeof(requiredSegments) / sizeof(memseg_metadata_t), sizeof(uint8_t));

    if(requiredSegmentsFoundMap == NULL)
    {
        error("unable to allocate found map", errbuf);
        free(specs);
        return NULL;
    }


    uint8_t specCount = 0;
    specs->executableSegmentsCount = 0;

    while (segments[specCount] != NULL)
    {
        if(segments[specCount]->isExecutable)
        {
            specs->executableSegmentsCount++;
        }
        specCount++;
    }

    if(specCount < sizeof(requiredSegments)/sizeof(memseg_metadata_t))
    {
        specCount = sizeof(requiredSegments)/sizeof(memseg_metadata_t);
    }

    specs->memSpecsCount = specCount;

    specs->memSpecs = calloc(specCount, sizeof(memseg_spec_t*));
    if (specs->memSpecs == NULL)
    {
        free(specs);
        free(requiredSegmentsFoundMap);
        error("unable to allocate mem segment specs", errbuf);
        return NULL;
    }

    specs->executableSegments = calloc(specs->executableSegmentsCount, sizeof(uint8_t));
    if (specs->executableSegments == NULL)
    {
        error("unable to allocate mem executable specs", errbuf);
        free(specs->memSpecs);
        free(specs);
        free(requiredSegmentsFoundMap);
        return NULL;
    }


    for (uint8_t i = 0; i < specCount; i++) 
    {
        memseg_spec_t* spec = (memseg_spec_t*)calloc(1, sizeof(memseg_spec_t));
        if (spec == NULL)
        {
            sprintf(errbuf, "unable to allocate spec %d", i);

            for (uint8_t j = 0; j < i; j++)
            {
                free(specs->memSpecs[j]);
            }
            free(specs->memSpecs);
            free(specs->executableSegments);
            free(specs);
            free(requiredSegmentsFoundMap);
            return NULL;
        }
        spec->name = NULL;

        specs->memSpecs[i] = spec;
    }

    uint8_t executableSegmentsFound = 0;

    for (uint8_t i = 0; i < specCount; i++)
    {
        uint8_t specNum = 0xFF;
        uint8_t is_error = 0;
        uint8_t seek_found = 0;

        for(uint8_t j = 0; j < sizeof(requiredSegments) / sizeof(memseg_metadata_t); j++)
        {
            const memseg_metadata_t seg_def = requiredSegments[j];

            if(strcmp(seg_def.name, segments[i]->name))
            {
                const uint8_t is_found = requiredSegmentsFoundMap[j];
                if(is_found)
                {
                    is_error = 1;
                    sprintf(errbuf, "duplicate segment %s", seg_def.name);
                }
                else
                {
                    if(seg_def.is_executable == segments[i]->isExecutable || segments[i]->isExecutable)
                    {
                        requiredSegmentsFoundMap[j] = 1;
                        specNum = j;
                        seek_found = 1;
                    }
                    else
                    {
                        is_error = 1;
                        sprintf(errbuf, "segment %s must be executable", seg_def.name);
                    }
                }
                break;
            }
        }

        if(seek_found == 0)
        {
            is_error = 1;
            sprintf(errbuf, "unsupported memory segment: %s", segments[i]->name);
        }

        if(is_error)
        {
            freeDevSpec(specs);
            free(requiredSegmentsFoundMap);
            return NULL;
        }
        specs->memSpecs[specNum]->name = calloc(strlen(segments[i]->name), sizeof(char));
        if(specs->memSpecs[specNum]->name == NULL)
        {
            sprintf(errbuf, "unable to allocate spec %d name", i);
            freeDevSpec(specs);
            free(requiredSegmentsFoundMap);
            return NULL;
        }

        strcpy(specs->memSpecs[specNum]->name, segments[i]->name);
        specs->memSpecs[specNum]->start = segments[i]->start;
        specs->memSpecs[specNum]->len = segments[i]->len;
        specs->memSpecs[specNum]->wordLen = segments[i]->wordLen;
        if(segments[i]->isExecutable)
        {
            specs->executableSegments[executableSegmentsFound] = specNum;
            executableSegmentsFound++;
        }
    }

    for(uint8_t i = 0; i < sizeof(requiredSegments)/sizeof(memseg_metadata_t); i++)
    {
        if(requiredSegmentsFoundMap[i] == 0)
        {
            const memseg_metadata_t seg_def = requiredSegments[i];

            specs->memSpecs[seg_def.seg_id]->start = seg_def.default_addr;
            specs->memSpecs[seg_def.seg_id]->len = seg_def.default_size;
            specs->memSpecs[seg_def.seg_id]->wordLen = seg_def.word_len;

            specs->memSpecs[seg_def.seg_id]->name = calloc(strlen(seg_def.name), sizeof(char));

            if(specs->memSpecs[seg_def.seg_id]->name == NULL)
            {
                sprintf(errbuf, "unable to allocate spec %d name", i);
                freeDevSpec(specs);
                free(requiredSegmentsFoundMap);
                return NULL;
            }

            strcpy(specs->memSpecs[seg_def.seg_id]->name, seg_def.name);


            if(seg_def.is_executable)
            {
                specs->executableSegments[executableSegmentsFound] = seg_def.seg_id;
                executableSegmentsFound++;
            }
        }
    }

    if(executableSegmentsFound < specs->executableSegmentsCount)
    {
        sprintf(errbuf, "Not all executable segments found");
        freeDevSpec(specs);
        free(requiredSegmentsFoundMap);
        return NULL;
    }

    #ifdef MEMSEG_PC_SEG_NUM 
        #ifndef MEMSEG_PC_ADDR
            specs->pcAddr = memSegToGlobal(specs, MEMSEG_PC_SEG_NUM, 0);
        #else
            specs->pcAddr = memSegToGlobal(specs, MEMSEG_PC_SEG_NUM, MEMSEG_PC_ADDR);
        #endif
    #endif
    free(requiredSegmentsFoundMap);

    return specs;
}

void freeDevSpecs(device_specs_t* specs)
{
    for(size_t k = 0; k < specs->memSpecsCount; k++)
    {
        if(specs->memSpecs[k]->name != NULL)
        {
            free(specs->memSpecs[k]->name);
        }
        free(specs->memSpecs[k]);
    }
    free(specs->executableSegments);
    free(specs->memSpecs);
    free(specs);
}

void fillSmartReadSpecs(device_specs_t* specs, smart_read_spec_t* smartReadSpecs, uint64_t smartReadSpecsCount)
{
    specs->smartReadSpecs = smartReadSpecs;
    specs->smartReadSpecsCount = smartReadSpecsCount;
}

void fillSmartWriteSpecs(device_specs_t* specs, smart_write_spec_t* smartWriteSpecs, uint64_t smartWriteSpecsCount)
{
    specs->smartWriteSpecs = smartWriteSpecs;
    specs->smartWriteSpecsCount = smartWriteSpecsCount;
}

device_public_context_t* initDefault(smart_read_spec_t* smartReadSpecs, uint64_t smartReadSpecsCount, smart_write_spec_t* smartWriteSpecs, uint64_t smartWriteSpecsCount, char* errbuf)
{
    return NULL;
    //TODO
    device_specs_t* specs = calloc(1, sizeof(device_specs_t));
    if (specs == NULL)
    {
        return NULL;
    }



    specs->memSpecsCount = 4;
    specs->memSpecs = calloc(specs->memSpecsCount, sizeof(memseg_spec_t*));
    if (specs->memSpecs == NULL)
    {
        sprintf(errbuf, "unable to allocate default mem segment specs");
        free(specs);
        return NULL;
    }


    specs->executableSegmentsCount = 1;
    specs->executableSegments = calloc(1, sizeof(uint8_t));
    if (specs->executableSegments == NULL)
    {
        sprintf(errbuf, "unable to allocate default executable segments");
        free(specs->memSpecs);
        free(specs);
        return NULL;
    }
    specs->executableSegments[0] = 0;

    for (uint8_t i = 0; i < specs->memSpecsCount; i++) 
    {
        specs->memSpecs[i] = calloc(1, sizeof(memseg_spec_t));
        if (specs->memSpecs[i] == NULL)
        {
            sprintf(errbuf, "unable to allocate default mem seg spec %u", i);
            for (uint8_t j = 0; j < i; j++)
            {
                free(specs->memSpecs[j]);
            }
            free(specs->memSpecs);
            free(specs->executableSegments);
            free(specs);
            return NULL;
        }
    }

    specs->smartReadSpecs = smartReadSpecs;
    specs->smartReadSpecsCount = smartReadSpecsCount;

    specs->smartWriteSpecs = smartWriteSpecs;
    specs->smartWriteSpecsCount = smartWriteSpecsCount;

    specs->memSpecs[0]->len = 255;
    specs->memSpecs[0]->start = 0;
    specs->memSpecs[0]->wordLen = 1;

    specs->memSpecs[1]->len = 1;
    specs->memSpecs[1]->start = 256;
    specs->memSpecs[1]->wordLen = 1;


    char initErrbuf[200];
    device_public_context_t* ret = init(specs, initErrbuf);

    if (ret == NULL)
    {
        sprintf(errbuf, "unable to init default: %s", initErrbuf);
        freeDevSpec(specs);
        return NULL;
    }
    free(specs->memSpecs);
    free(specs->executableSegments);
    free(specs);
    
    return ret;
}


device_public_context_t* init(device_specs_t* specs, char* errbuf)
{
    device_public_context_t* pubDevContext = calloc(1, sizeof(device_public_context_t));
    if (pubDevContext == NULL) 
    {
        sprintf(errbuf, "unable to allocate public context");
        return NULL;
    }

    char initErrbuf[200];
    device_info_t* devInfo = initSpecs(specs, initErrbuf);
    if (devInfo == NULL)
    {
        sprintf(errbuf, "unable to init specs: %s", initErrbuf);
        free(pubDevContext);
        return NULL;
    }

    pubDevContext->deviceInfo   = (void*)devInfo;
    pubDevContext->deviceMem    = devInfo->deviceMem;

    return pubDevContext;
}

uint8_t pubDeviceType() 
{
    return DEVICE_TYPE;
}