#include "iap/iap.h"
#include "uart_driver.h"
#include "flash_manager.h"
#include "gd32f4xx.h"
#include <cstring>
#include <cstdio>

#define YM_SOH 0x01
#define YM_STX 0x02
#define YM_EOT 0x04
#define YM_ACK 0x06
#define YM_NAK 0x15
#define YM_CAN 0x18
#define YM_CRC 0x43

#define YM_PACKET_SIZE_128 128
#define YM_PACKET_SIZE_1024 1024
#define YM_PACKET_OVERHEAD 5

__attribute__((section(".sdram"))) static uint8_t s_firmware_buffer[IAP_BUFFER_SIZE];

static uint16_t crc16_ccitt(const uint8_t *data, uint32_t len)
{
    uint16_t crc = 0;
    for (uint32_t i = 0; i < len; i++)
    {
        crc ^= (uint16_t)data[i] << 8;
        for (int j = 0; j < 8; j++)
        {
            if (crc & 0x8000)
                crc = (crc << 1) ^ 0x1021;
            else
                crc <<= 1;
        }
    }
    return crc;
}

Iap::Iap()
    : initialized_(false), firmware_size_(0)
{
    firmware_name_[0] = '\0';
}

Iap &Iap::instance()
{
    static Iap inst;
    return inst;
}

RetCode Iap::init()
{
    initialized_ = true;
    firmware_size_ = 0;
    firmware_name_[0] = '\0';
    return RET_OK;
}

RetCode Iap::receiveFirmware(USART &uart, uint32_t timeout_ms)
{
    if (!initialized_)
        return RET_NOT_INITIALIZED;

    firmware_size_ = 0;
    firmware_name_[0] = '\0';

    RetCode ret = ymodemReceive(uart, timeout_ms);
    if (ret != RET_OK)
    {
        return ret;
    }

    if (firmware_size_ == 0 || firmware_size_ > IAP_FIRMWARE_MAX_SIZE)
    {
        return RET_ERROR;
    }

    return RET_OK;
}

RetCode Iap::programFirmware(uint32_t flash_addr)
{
    if (!initialized_)
        return RET_NOT_INITIALIZED;
    if (firmware_size_ == 0)
        return RET_ERROR;

    FlashManager &flash = FlashManager::instance();

    if (!flash.is_initialized())
    {
        RetCode ret = flash.init();
        if (ret != RET_OK)
            return ret;
    }

    uint32_t sector_size = flash.sector_size();
    uint32_t start_aligned = flash_addr & ~(sector_size - 1);
    uint32_t end_addr = flash_addr + firmware_size_;
    uint32_t end_aligned = (end_addr + sector_size - 1) & ~(sector_size - 1);
    uint32_t erase_size = end_aligned - start_aligned;

    RetCode ret = flash.erase(start_aligned, erase_size);
    if (ret != RET_OK)
        return ret;

    ret = flash.write(flash_addr, s_firmware_buffer, firmware_size_);
    if (ret != RET_OK)
        return ret;

    return RET_OK;
}

RetCode Iap::verifyFirmware(uint32_t flash_addr)
{
    if (!initialized_)
        return RET_NOT_INITIALIZED;
    if (firmware_size_ == 0)
        return RET_ERROR;

    FlashManager &flash = FlashManager::instance();
    if (!flash.is_initialized())
    {
        RetCode ret = flash.init();
        if (ret != RET_OK)
            return ret;
    }

    uint8_t verify_buffer[256];
    uint32_t remaining = firmware_size_;
    uint32_t offset = 0;

    while (remaining > 0)
    {
        uint32_t chunk = (remaining > sizeof(verify_buffer)) ? sizeof(verify_buffer) : remaining;
        RetCode ret = flash.read(flash_addr + offset, verify_buffer, chunk);
        if (ret != RET_OK)
            return ret;

        if (memcmp(verify_buffer, &s_firmware_buffer[offset], chunk) != 0)
        {
            return RET_ERROR;
        }

        offset += chunk;
        remaining -= chunk;
    }

    return RET_OK;
}

void Iap::jumpToApp(uint32_t app_addr)
{
    __disable_irq();

    for (uint32_t i = 0; i < 8; i++)
    {
        NVIC->ICER[i] = 0xFFFFFFFF;
        NVIC->ICPR[i] = 0xFFFFFFFF;
    }

    __DSB();
    __ISB();

    SCB->VTOR = app_addr;

    __DSB();
    __ISB();

    uint32_t msp = *(volatile uint32_t *)app_addr;
    uint32_t reset_vector = *(volatile uint32_t *)(app_addr + 4);

    __set_MSP(msp);
    __set_PSP(msp);

    void (*app_entry)(void) = (void (*)(void))reset_vector;
    app_entry();

    while (1)
        ;
}

bool Iap::isValidApp(uint32_t app_addr)
{
    if (app_addr == 0 || (app_addr & 0x1FF) != 0)
    {
        return false;
    }

    uint32_t msp = *(volatile uint32_t *)app_addr;
    uint32_t reset = *(volatile uint32_t *)(app_addr + 4);

    bool sp_ok = (msp >= 0x10000000 && msp <= 0x10010000) ||
                 (msp >= 0x20000000 && msp <= 0x20030000);

    bool reset_ok = (reset >= 0x08000000 && reset <= 0x08100000);

    return sp_ok && reset_ok;
}

RetCode Iap::backupDomainEnable(void)
{
    rcu_periph_clock_enable(RCU_PMU);
    pmu_backup_write_enable();

    if (!(RCU_BDCTL & RCU_BDCTL_RTCEN))
    {
        RCU_BDCTL |= RCU_BDCTL_RTCEN;
    }

    return RET_OK;
}

void Iap::bootFlagWrite(uint32_t magic)
{
    backupDomainEnable();
    RTC_BKP0 = magic;
}

uint32_t Iap::bootFlagRead(void)
{
    rcu_periph_clock_enable(RCU_PMU);

    if (!(RCU_BDCTL & RCU_BDCTL_RTCEN))
    {
        RCU_BDCTL |= RCU_BDCTL_RTCEN;
    }

    return RTC_BKP0;
}

// ���������� Ymodem Receive ����������

void Iap::ymodemSend(USART &uart, uint8_t b)
{
    uart.send_byte(b);
}

int Iap::ymodemRead(USART &uart, uint32_t timeout_ms)
{
    uint8_t b;
    RetCode ret = uart.receive_byte(&b, timeout_ms);
    if (ret != RET_OK)
        return -1;
    return (int)b;
}

RetCode Iap::ymodemReceive(USART &uart, uint32_t timeout_ms)
{
    uint8_t pkt_buf[YM_PACKET_SIZE_1024 + YM_PACKET_OVERHEAD];
    uint32_t total_received = 0;
    uint8_t expected_seq = 1;
    bool first_packet = true;

    printf("IAP: Waiting for Ymodem transfer (timeout=%lu ms)...\r\n", timeout_ms);

    int remaining_ms = (int)timeout_ms;

    while (remaining_ms > 0)
    {
        ymodemSend(uart, YM_CRC);

        int byte = ymodemRead(uart, 3000);
        if (byte < 0)
        {
            remaining_ms -= 3000;
            continue;
        }

        int pkt_type = byte;
        int pkt_len = 0;

        if (pkt_type == YM_SOH)
        {
            pkt_len = YM_PACKET_SIZE_128;
        }
        else if (pkt_type == YM_STX)
        {
            pkt_len = YM_PACKET_SIZE_1024;
        }
        else if (pkt_type == YM_EOT)
        {
            ymodemSend(uart, YM_NAK);
            byte = ymodemRead(uart, 3000);
            if (byte == YM_EOT)
            {
                ymodemSend(uart, YM_ACK);
                printf("IAP: Transfer complete (%lu bytes)\r\n", total_received);
                firmware_size_ = total_received;
                return RET_OK;
            }
            continue;
        }
        else if (pkt_type == YM_CAN)
        {
            printf("IAP: Transfer cancelled by sender\r\n");
            return RET_ERROR;
        }
        else
        {
            continue;
        }

        pkt_buf[0] = (uint8_t)pkt_type;
        pkt_buf[1] = (uint8_t)ymodemRead(uart, 100);
        pkt_buf[2] = (uint8_t)ymodemRead(uart, 100);

        uint8_t block_num = pkt_buf[1];
        uint8_t block_comp = pkt_buf[2];

        if (block_num != (uint8_t)(~block_comp))
        {
            ymodemSend(uart, YM_NAK);
            remaining_ms -= 100;
            continue;
        }

        bool data_ok = true;
        for (int i = 0; i < pkt_len; i++)
        {
            int d = ymodemRead(uart, 100);
            if (d < 0)
            {
                data_ok = false;
                break;
            }
            pkt_buf[3 + i] = (uint8_t)d;
        }

        if (!data_ok)
        {
            ymodemSend(uart, YM_NAK);
            remaining_ms -= 100;
            continue;
        }

        int crc_h = ymodemRead(uart, 100);
        int crc_l = ymodemRead(uart, 100);
        if (crc_h < 0 || crc_l < 0)
        {
            ymodemSend(uart, YM_NAK);
            remaining_ms -= 100;
            continue;
        }

        uint16_t received_crc = (uint16_t)((uint16_t)crc_h << 8) | (uint16_t)crc_l;
        uint16_t calc_crc = crc16_ccitt(&pkt_buf[3], (uint32_t)pkt_len);

        if (received_crc != calc_crc)
        {
            ymodemSend(uart, YM_NAK);
            remaining_ms -= 100;
            continue;
        }

        if (first_packet)
        {
            if (block_num != 0)
            {
                ymodemSend(uart, YM_NAK);
                remaining_ms -= 100;
                continue;
            }

            const char *data_str = (const char *)&pkt_buf[3];
            const char *fn = data_str;
            uint32_t fn_len = 0;
            while (fn_len < (uint32_t)pkt_len && data_str[fn_len] != '\0')
                fn_len++;

            if (fn_len > 0 && fn_len < sizeof(firmware_name_))
            {
                memcpy(firmware_name_, fn, fn_len);
                firmware_name_[fn_len] = '\0';
            }

            uint32_t name_end = fn_len + 1;
            if (name_end < (uint32_t)pkt_len)
            {
                const char *size_str = data_str + name_end;
                firmware_size_ = 0;
                while (*size_str >= '0' && *size_str <= '9')
                {
                    firmware_size_ = firmware_size_ * 10 + (*size_str - '0');
                    size_str++;
                }
            }

            if (firmware_size_ == 0 || firmware_size_ > IAP_FIRMWARE_MAX_SIZE)
            {
                printf("IAP: Invalid firmware size: %lu\r\n", firmware_size_);
                ymodemSend(uart, YM_CAN);
                ymodemSend(uart, YM_CAN);
                return RET_INVALID_PARAM;
            }

            printf("IAP: Receiving '%s' (%lu bytes)\r\n", firmware_name_, firmware_size_);
            ymodemSend(uart, YM_ACK);
            first_packet = false;
            expected_seq = 1;
        }
        else
        {
            if (block_num != expected_seq)
            {
                ymodemSend(uart, YM_NAK);
                remaining_ms -= 100;
                continue;
            }

            if (total_received + (uint32_t)pkt_len > IAP_BUFFER_SIZE)
            {
                printf("IAP: Buffer overflow\r\n");
                ymodemSend(uart, YM_CAN);
                ymodemSend(uart, YM_CAN);
                return RET_ERROR;
            }

            memcpy(&s_firmware_buffer[total_received], &pkt_buf[3], (uint32_t)pkt_len);
            total_received += (uint32_t)pkt_len;

            ymodemSend(uart, YM_ACK);
            expected_seq++;
        }

        remaining_ms = (int)timeout_ms;
    }

    printf("IAP: Timeout waiting for Ymodem transfer\r\n");
    return RET_TIMEOUT;
}

int Iap::ymodemReadPacket(USART &uart, uint8_t *buf, uint32_t timeout_ms)
{
    int byte = ymodemRead(uart, timeout_ms);
    if (byte < 0)
        return -1;

    int pkt_type = byte;
    int data_len = 0;

    if (pkt_type == YM_SOH)
    {
        data_len = YM_PACKET_SIZE_128;
    }
    else if (pkt_type == YM_STX)
    {
        data_len = YM_PACKET_SIZE_1024;
    }
    else if (pkt_type == YM_EOT || pkt_type == YM_CAN)
    {
        return pkt_type;
    }
    else
    {
        return -1;
    }

    buf[0] = (uint8_t)pkt_type;

    for (int i = 1; i < data_len + YM_PACKET_OVERHEAD - 1; i++)
    {
        int d = ymodemRead(uart, 100);
        if (d < 0)
            return -1;
        buf[i] = (uint8_t)d;
    }

    return pkt_type;
}