#include "dcc/dn_dcc_proto.h"
#include "dcc/plat.h"
#include "devices.h"

#ifdef CDEFS
const char CFLAGS[] = "C:DumpNow DCC Loader. (c) 2026 Wrapper.;Compile flags: " CDEFS ";Compile Date: " __DATE__;
#endif

static uint8_t read_buffer[DCC_BUFFER_SIZE + 0x4000];
#if !USE_BREAKPOINTS && !DISABLE_COMPRESS
static uint8_t compress_buffer[DCC_BUFFER_SIZE + 0x8000];
#endif

size_t strlen(const char *str);

static inline uint8_t dcc_send_data(uint8_t *data, uint32_t size, uint8_t algo) {
    uint32_t dcc_comp_packet_size;

    switch (algo) {
        case CMD_READ_COMP_RLE:
        #if !USE_BREAKPOINTS && !DISABLE_COMPRESS
            dcc_comp_packet_size = DN_Packet_Compress(data, size, compress_buffer);
            DN_Packet_Send(compress_buffer, dcc_comp_packet_size);
            break;
        #endif

        case CMD_READ_COMP_NONE:
            DN_Packet_Send_DirectUncompressed(data, ALIGN4(size));
            break;
        
        #if HAVE_LZ4 && !USE_BREAKPOINTS && !DISABLE_COMPRESS
        case CMD_READ_COMP_LZ4:
            dcc_comp_packet_size = DN_Packet_Compress2(data, size, compress_buffer);
            DN_Packet_Send(compress_buffer, dcc_comp_packet_size);
            break;
        #endif

        #if HAVE_ZLIB && !USE_BREAKPOINTS && !DISABLE_COMPRESS
        case CMD_READ_COMP_ZLIB:
            dcc_comp_packet_size = DN_Packet_Compress3(data, size, compress_buffer);
            DN_Packet_Send(compress_buffer, dcc_comp_packet_size);
            break;
        #endif

        default:
            DN_Packet_Send_One(CMD_READ_RESP_FAIL(DCC_INVALID_ARGS));
            return 0;
    }

    return 1;
}

#if DEBUG_PORT
static int DN_DebugPort_Out(int ch, lwprintf_t* lwp) {
    if (ch != 0x00)
        WRITE_U32(0x44444440, ch);
    return ch;
}
#endif

// dcc code
void dcc_main(uint32_t StartAddress, uint32_t PageSize) {
    DCCMemory dev_mem[16] = { 0 };
    uint8_t mem_has_oob[16] = { 0 };
    uint32_t dcc_init_buf[2048];
    uint32_t dcc_init_offset = 0;
    uint32_t temp_ext_mem;
    DCC_RETURN res;

    /* 00 - Init debug stuff */
#if DEBUG_PORT
    lwprintf_init(DN_DebugPort_Out);
#endif

    /* 01 - Probe flash devices */
    PLAT_DBG_LOG("[DUMPNOW_DCC]: Begin initializing memory devices\n");

    for (int i = 0; i < 16; i++) {
        if (!devices[i].driver) break; // Break when reaching the end of list

        /* Probe device */
        res = devices[i].driver->initialize(&dev_mem[i], devices[i].base_offset, PageSize);
        PLAT_DBG_LOG("[DUMPNOW_DCC]: Memory device %d initialization returned 0x%02x\n", i + 1, res);

        if (res != DCC_OK) 
            dev_mem[i].type = MEMTYPE_NONE;

        /* Print appropriate value */
        switch (dev_mem[i].type) {
            /* Anything without spare */
            case MEMTYPE_NOR:
            case MEMTYPE_SUPERAND:
                temp_ext_mem = DCC_MEM_EXTENDED(1, dev_mem[i].page_size, dev_mem[i].block_size, dev_mem[i].size >> 20);
                mem_has_oob[i] = 0;

            /* Extended memory logic */
            WRITE_EXTMEM:
                /* Set name flag if memory name is defined */
                if (strlen(dev_mem[i].name)) temp_ext_mem |= 0x80;

                /* First device info */
                dcc_init_buf[dcc_init_offset++] = DCC_MEM_OK | (temp_ext_mem << 16);
                dcc_init_buf[dcc_init_offset++] = dev_mem[i].manufacturer | (dev_mem[i].device_id << 16);
                
                /* Print additional information */
                if (strlen(dev_mem[i].name)) {
                    int sLen = strlen(dev_mem[i].name);
                    uint8_t *dcc_init_buf_uint8 = (uint8_t *)dcc_init_buf;
                    
                    /* 8-bit length then name (WORD aligned) */
                    dcc_init_buf[dcc_init_offset] = sLen;
                    memcpy(dcc_init_buf_uint8 + (dcc_init_offset << 2) + 1, dev_mem[i].name, sLen);

                    dcc_init_offset += ALIGN4(1 + sLen) >> 2;
                }

                /* Second device info */
                dcc_init_buf[dcc_init_offset++] = temp_ext_mem;
                break;

            /* Regular NAND */
            case MEMTYPE_NAND:
                if (strlen(dev_mem[i].name)) goto NAND_EXTMEM; // Extended device info if we have additional information

                /* Device info */
                dcc_init_buf[dcc_init_offset++] = DCC_MEM_OK | (dev_mem[i].page_size << 16);
                dcc_init_buf[dcc_init_offset++] = dev_mem[i].manufacturer | (dev_mem[i].device_id << 16);
                mem_has_oob[i] = 1;
                break;

            /* Anything with spare */
            case MEMTYPE_ONENAND:
            case MEMTYPE_AND:
            case MEMTYPE_AG_AND:
            NAND_EXTMEM:
                temp_ext_mem = DCC_MEM_EXTENDED(0, dev_mem[i].page_size, dev_mem[i].block_size, dev_mem[i].size >> 20);
                mem_has_oob[i] = 1;
                goto WRITE_EXTMEM;

            /* When device probe fails, it goes here */
            default:
                dcc_init_buf[dcc_init_offset++] = DCC_MEM_OK | (DCC_MEM_NONE << 16);
                dcc_init_buf[dcc_init_offset++] = dev_mem[i].probe_error_code;
                mem_has_oob[i] = 0;

        }
    }

    /* 02 - Print buffer size */
    dcc_init_buf[dcc_init_offset++] = DCC_MEM_OK | (DCC_MEM_BUFFER(0) << 16);
    dcc_init_buf[dcc_init_offset++] = DCC_BUFFER_SIZE;
    
    DN_Packet_Send((uint8_t *)dcc_init_buf, dcc_init_offset << 2);

    uint32_t flashDevID;
    uint32_t dataOffset;
    uint32_t dataSize;
    uint32_t dataBlockSize;
    uint32_t oobBlockSize;

    /* 03 - The loop */
    PLAT_DBG_LOG("[DUMPNOW_DCC]: Ready to accept commands\n");

    while (1) {
        wdog_reset();
        uint32_t cmd = DN_Packet_DCC_Read();

        switch (cmd & 0xff) {
            /* Extended operation */
            case CMD_EXTENDED:
                for (int c = 0; c < (cmd >> 0x10); c += 4) {
                    DN_Packet_DCC_Read();
                }

                // TODO: Extended command support
                DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(0x38, 0x6)); 
                break;

            /* Get devices information */
            case CMD_GETINFO:
                DN_Packet_Send((uint8_t *)dcc_init_buf, dcc_init_offset << 2);
                break;
                
            /* Get memory size */
            case CMD_GETMEMSIZE:
                flashDevID = (cmd >> 8) & 0xff;
                if (flashDevID == 0) {
                    DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(0x21, 0));
                } else if (flashDevID <= 0x10 && dev_mem[flashDevID - 1].type != MEMTYPE_NONE) {
                    DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(0x21, dev_mem[flashDevID - 1].size >> 20));
                } else {
                    DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_FLASH_NOENT, flashDevID));
                }
                break;

            /* Flash read */
            case CMD_READ:
                flashDevID = (cmd >> 8) & 0xff;
                dataOffset = DN_Packet_DCC_Read();
                dataSize   = DN_Packet_DCC_Read();
                uint8_t  readCompAlgo   = (cmd >> 24) & 0xff;
                uint32_t readSize;

                /* Check for read size not exceeding buffer */
                if (dataSize > DCC_BUFFER_SIZE) {
                    DN_Packet_Send_One(CMD_READ_RESP_FAIL(DCC_INVALID_ARGS));
                    continue;
                }

                if (flashDevID == 0) { // Read from RAM
                Jump_Read_NOR:
                    if (!dcc_send_data((uint8_t *)(dataOffset), dataSize, readCompAlgo))
                        continue;
                    
                } else if (flashDevID <= 0x10 && dev_mem[flashDevID - 1].type != MEMTYPE_NONE) {
                    switch (dev_mem[flashDevID - 1].type) {
                        case MEMTYPE_NAND:
                        case MEMTYPE_ONENAND:
                        case MEMTYPE_SUPERAND:
                        case MEMTYPE_AND:
                        case MEMTYPE_AG_AND:
                            /* Perform driver read routines */
                            res = devices[flashDevID - 1].driver->read(&dev_mem[flashDevID - 1], dataOffset, dataSize, read_buffer, &readSize);
                            if (res != DCC_OK) { // Check if error
                                DN_Packet_Send_One(CMD_READ_RESP_FAIL(res));
                                continue;
                            }
                            
                            /* Compression */
                            if (!dcc_send_data(read_buffer, readSize, readCompAlgo))
                                continue;

                            break;
                        case MEMTYPE_NOR:
                        default:
                            /* NOR reads directly */
                            dataOffset &= (dev_mem[flashDevID - 1].size - 1);
                            dataOffset += dev_mem[flashDevID - 1].base_offset;
                            goto Jump_Read_NOR;
                    }
                } else {
                    /* No flash found */
                    DN_Packet_Send_One(CMD_READ_RESP_FAIL(DCC_FLASH_NOENT));
                }

                break;
            
            /* Flash erase */
            case CMD_ERASE:
                flashDevID = (cmd >> 8) & 0xff;
                dataOffset = DN_Packet_DCC_Read();
                dataSize = DN_Packet_DCC_Read();
                dataBlockSize = dev_mem[flashDevID - 1].block_size;

                if (flashDevID == 0) flashDevID = 1;

                if (flashDevID <= 0x10 && dev_mem[flashDevID - 1].type != MEMTYPE_NONE) {
                    if (((dataOffset % dataBlockSize) != 0) || ((dataSize % dataBlockSize) != 0)) {
                        DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_INVALID_ARGS, flashDevID));
                        continue;
                    }

                    if (devices[flashDevID - 1].driver->erase) {
                        res = devices[flashDevID - 1].driver->erase(&dev_mem[flashDevID - 1], dataOffset, dataSize);
                        DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(res, flashDevID));
                    } else {
                        DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_ROFS_ERROR, flashDevID));
                    }
                } else {
                    DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_FLASH_NOENT, flashDevID));
                }
                break;

            /* Flash write */
            case CMD_WRITE:
                flashDevID = (cmd >> 16) & 0xff;
                dataOffset = DN_Packet_DCC_Read();
                dataBlockSize = dev_mem[flashDevID - 1].block_size;
                oobBlockSize = mem_has_oob[flashDevID - 1] ? (dev_mem[flashDevID - 1].block_size >> 5) : 0;

                DecompressContext decompCtx;
                uint32_t writeCompAlgo      = DN_Packet_DCC_Read();
                uint8_t  writeDataType      = (cmd >> 8) & 0x7f;
                uint8_t  writeUseECC        = (cmd >> 8) & 0x80;
                uint32_t crcDataComputed    = 0xffffffff;
                uint32_t writeDataBufOffset = 0;

                if (flashDevID == 0) flashDevID = 1;

                if (flashDevID <= 0x10 && dev_mem[flashDevID - 1].type != MEMTYPE_NONE) {
                    if ((dataOffset % dataBlockSize) != 0) {
                        DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_INVALID_ARGS, flashDevID));
                        continue;
                    }

                    switch (writeCompAlgo) {
                        case CMD_WRITE_COMP_NONE:
                            if (writeDataType & 2) { // Data
                                DN_Packet_Read(read_buffer + writeDataBufOffset, dataBlockSize);
                                crcDataComputed = DN_Calculate_CRC32(crcDataComputed, read_buffer + writeDataBufOffset, dataBlockSize);
                                writeDataBufOffset += dataBlockSize;
                            }
                            if ((writeDataType & 1) && oobBlockSize > 0) { // OOB
                                DN_Packet_Read(read_buffer + writeDataBufOffset, oobBlockSize);
                                crcDataComputed = DN_Calculate_CRC32(crcDataComputed, read_buffer + writeDataBufOffset, oobBlockSize);
                                writeDataBufOffset += oobBlockSize;
                            }
                            break;

                        case CMD_WRITE_COMP_RLE:
                            dataSize = DN_Packet_DCC_Read();
                            DN_Packet_Reader_Init(&decompCtx, read_buffer, dataSize);
                            while (decompCtx.decomp_data_size > 0) {
                                DN_Packet_Reader_Process(&decompCtx, DN_Packet_DCC_Read());
                            }

                            if (decompCtx.out_buffer_offset < (dataBlockSize + oobBlockSize)) {
                                DN_Packet_Send_One(DCC_E_UNDERRUN(decompCtx.out_buffer_offset));
                                continue;
                            }

                            crcDataComputed = DN_Calculate_CRC32(crcDataComputed, read_buffer, dataSize);
                            break;

                        default:
                            DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_CHECKSUM_ERROR, flashDevID));
                            continue;
                    }

                    uint32_t checksum = DN_Packet_DCC_Read();

                    if (checksum != crcDataComputed) {
                        DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_CHECKSUM_ERROR, flashDevID));
                        continue;
                    }

                    if (devices[flashDevID - 1].driver->configure)
                        devices[flashDevID - 1].driver->configure(&dev_mem[flashDevID - 1], CONFIG_ECC, writeUseECC);

                    if (devices[flashDevID - 1].driver->write) {
                        res = devices[flashDevID - 1].driver->write(&dev_mem[flashDevID - 1], dataOffset, read_buffer, dataBlockSize);
                        DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(res, flashDevID));
                    } else {
                        DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_ROFS_ERROR, flashDevID));
                    }
                } else {
                    DN_Packet_Send_One(CMD_WRITE_ERASE_STATUS(DCC_FLASH_NOENT, flashDevID));
                }
                break;

            /* Catch-all for unknown commands */
            default:
                DN_Packet_Send_One(DCC_BAD_COMMAND(cmd & 0xff));
        }
    }
}