#include "pkt_normalizer.h"
#include "u_async.h"
#include "settings.h"
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdio.h>

static void packer_handler(ll_queue_t* q, ll_entry_t* unused, void* arg);
static void unpacker_handler(ll_queue_t* q, ll_entry_t* unused, void* arg);
static void send_buf(pn_struct* pn);
static int get_header(uint8_t* header, size_t L);
static void fragment_timeout_cb(void* arg);
static void set_fragment_timeout(pn_struct* pn);
static void cancel_fragment_timeout(pn_struct* pn);

pn_struct* pkt_normalizer_init(int is_packer) {
    pn_struct* pn = malloc(sizeof(pn_struct));
    if (!pn) return NULL;

    pn->input = queue_new();
    if (!pn->input) {
        free(pn);
        return NULL;
    }
    pn->output = queue_new();
    if (!pn->output) {
        queue_free(pn->input);
        free(pn);
        return NULL;
    }
    pn->is_packer = is_packer;

    if (is_packer) {
        pn->u.packer.cap = settings.max_fragment_size;
        pn->u.packer.buf = malloc(pn->u.packer.cap);
        if (!pn->u.packer.buf) {
            queue_free(pn->input);
            queue_free(pn->output);
            free(pn);
            return NULL;
        }
        pn->u.packer.len = 0;
        pn->u.packer.error_count = 0;
        queue_set_callback(pn->input, packer_handler, pn);
    } else {
        pn->u.unpacker.buf = NULL;
        pn->u.unpacker.len = 0;
        pn->u.unpacker.total_len = 0;
        pn->u.unpacker.cap = 0;
        pn->u.unpacker.error_count = 0;
        pn->u.unpacker.in_fragment = 0;
        pn->u.unpacker.fragment_timeout_id = NULL;
        queue_set_callback(pn->input, unpacker_handler, pn);
    }
    return pn;
}

void pkt_normalizer_deinit(pn_struct* pn) {
    if (!pn) return;
    queue_free(pn->input);
    queue_free(pn->output);
    if (pn->is_packer) {
        free(pn->u.packer.buf);
    } else {
        cancel_fragment_timeout(pn);
        free(pn->u.unpacker.buf);
    }
    free(pn);
}

pkt_normalizer_pair* pkt_normalizer_pair_init(void) {
    pkt_normalizer_pair* pair = malloc(sizeof(pkt_normalizer_pair));
    if (!pair) return NULL;

    pair->packer = pkt_normalizer_init(1);
    if (!pair->packer) {
        free(pair);
        return NULL;
    }
    pair->unpacker = pkt_normalizer_init(0);
    if (!pair->unpacker) {
        pkt_normalizer_deinit(pair->packer);
        free(pair);
        return NULL;
    }
    return pair;
}

void pkt_normalizer_pair_deinit(pkt_normalizer_pair* pair) {
    if (!pair) return;
    pkt_normalizer_deinit(pair->packer);
    pkt_normalizer_deinit(pair->unpacker);
    free(pair);
}

static int get_header(uint8_t* header, size_t L) {
    if (L > 3839) return -1;
    if (L <= 239) {
        header[0] = (uint8_t)L;
        return 1;
    } else {
        uint8_t high = (uint8_t)(L >> 8);
        if (high > 14) return -1;
        header[0] = 0xF0 + high;
        header[1] = (uint8_t)(L & 0xFF);
        return 2;
    }
}

/* Сбросить состояние сборки фрагментов */
static void reset_fragment_state(pn_struct* pn) {
    if (!pn->is_packer) {
        cancel_fragment_timeout(pn);
        pn->u.unpacker.len = 0;
        pn->u.unpacker.total_len = 0;
        pn->u.unpacker.in_fragment = 0;
    }
}

/* Таймаут для сборки фрагментов */
static void fragment_timeout_cb(void* arg) {
    pn_struct* pn = (pn_struct*)arg;
    if (!pn || pn->is_packer) return;
    
    pn->u.unpacker.fragment_timeout_id = NULL;
    pn->u.unpacker.error_count++;
    reset_fragment_state(pn);
}

/* Установить таймаут для текущей сборки фрагментов */
static void set_fragment_timeout(pn_struct* pn) {
    if (!pn || pn->is_packer) return;
    
    cancel_fragment_timeout(pn);
    pn->u.unpacker.fragment_timeout_id = uasync_set_timeout(
        PKT_NORMALIZER_FRAGMENT_TIMEOUT, pn, fragment_timeout_cb);
}

/* Отменить таймаут сборки фрагментов */
static void cancel_fragment_timeout(pn_struct* pn) {
    if (!pn || pn->is_packer) return;
    
    if (pn->u.unpacker.fragment_timeout_id) {
        uasync_cancel_timeout(pn->u.unpacker.fragment_timeout_id);
        pn->u.unpacker.fragment_timeout_id = NULL;
    }
}

static void send_buf(pn_struct* pn) {
    if (pn->u.packer.len == 0) return;
    size_t payload_len = pn->u.packer.len;
    printf("[PN DEBUG] send_buf: packer len=%zu, output queue count=%d\n", payload_len, queue_entry_count(pn->output));
    ll_entry_t* out = queue_entry_new(2 + payload_len);
    if (!out) return;
    uint8_t* d = ll_entry_data(out);
    *(uint16_t*)d = (uint16_t)payload_len;
    memcpy(d + 2, pn->u.packer.buf, payload_len);
    queue_entry_put(pn->output, out);
    pn->u.packer.len = 0;
}

static void packer_handler(ll_queue_t* q, ll_entry_t* unused, void* arg) {
    (void)unused;
    pn_struct* pn = arg;
    size_t max = (size_t)settings.max_fragment_size;
    
    ll_entry_t* entry = queue_entry_get(q);
    if (!entry) {
        queue_resume_callback(q);
        return;
    }
    
    size_t L = ll_entry_size(entry);
    uint8_t* data = ll_entry_data(entry);
    uint8_t header[2];
    int hsize = get_header(header, L);
    size_t needed = (size_t)hsize + L;
    
    if (hsize < 0 || needed > max) {
        // Fragment
        if (pn->u.packer.len > 0) {
            send_buf(pn);
        }
        size_t remaining = L;
        size_t pos = 0;
        int fragment_count = 0;
        while (remaining > 0) {
            size_t chunk;
            size_t payload_len;
            ll_entry_t* fout;
            uint8_t* fd;
            uint8_t frag_header[2];
            int frag_hsize;
            
            if (fragment_count == 0) {
                // Первый фрагмент: FF + общая длина (2 байта)
                chunk = remaining > (max - 5) ? (max - 5) : remaining; // 2+1+2+chunk <= max
                payload_len = 1 + 2 + chunk; // FF + total_len + data
                fout = queue_entry_new(2 + payload_len);
                if (!fout) break;
                fd = ll_entry_data(fout);
                *(uint16_t*)fd = (uint16_t)payload_len;
                fd += 2;
                *fd++ = 0xFF;
                *fd++ = (uint8_t)(L >> 8);   // старший байт общей длины
                *fd++ = (uint8_t)(L & 0xFF); // младший байт общей длины
            } else {
                // Не первый фрагмент
                if (remaining <= max - 3) {
                    // Это последний возможный фрагмент (помещается в один пакет с префиксом FE)
                    // Пытаемся отправить как обычный блок
                    frag_hsize = get_header(frag_header, remaining);
                    if (frag_hsize > 0 && (size_t)frag_hsize + remaining + 2 <= max) {
                        // Успешно: обычный блок
                        payload_len = frag_hsize + remaining;
                        chunk = remaining;
                        fout = queue_entry_new(2 + payload_len);
                        if (!fout) break;
                        fd = ll_entry_data(fout);
                        *(uint16_t*)fd = (uint16_t)payload_len;
                        fd += 2;
                        memcpy(fd, frag_header, frag_hsize);
                        fd += frag_hsize;
                    } else {
                        // Не удалось отправить как обычный блок - ошибка
                        pn->u.packer.error_count++;
                        // Отправляем как FE (нарушение спецификации)
                        chunk = remaining > (max - 3) ? (max - 3) : remaining;
                        payload_len = 1 + chunk; // FE + data
                        fout = queue_entry_new(2 + payload_len);
                        if (!fout) break;
                        fd = ll_entry_data(fout);
                        *(uint16_t*)fd = (uint16_t)payload_len;
                        fd += 2;
                        *fd++ = 0xFE;
                    }
                } else {
                    // Промежуточный фрагмент, отправляем как FE
                    chunk = remaining > (max - 3) ? (max - 3) : remaining;
                    payload_len = 1 + chunk; // FE + data
                    fout = queue_entry_new(2 + payload_len);
                    if (!fout) break;
                    fd = ll_entry_data(fout);
                    *(uint16_t*)fd = (uint16_t)payload_len;
                    fd += 2;
                    *fd++ = 0xFE;
                }
            }
            memcpy(fd, data + pos, chunk);
            queue_entry_put(pn->output, fout);
            pos += chunk;
            remaining -= chunk;
            fragment_count++;
        }
    } else {
        if (pn->u.packer.len + needed > max) {
            send_buf(pn);
        }
        // Add to buffer
        uint8_t* p = pn->u.packer.buf + pn->u.packer.len;
        memcpy(p, header, (size_t)hsize);
        memcpy(p + hsize, data, L);
        pn->u.packer.len += needed;
    }
    
    queue_entry_free(entry);
    
    if (pn->u.packer.len > 0) {
        send_buf(pn);
    }
    
    queue_resume_callback(q);
}

int pkt_normalizer_get_error_count(const pn_struct* pn) {
    if (!pn) return 0;
    if (pn->is_packer) {
        return pn->u.packer.error_count;
    }
    return pn->u.unpacker.error_count;
}

void pkt_normalizer_reset_error_count(pn_struct* pn) {
    if (!pn) return;
    if (pn->is_packer) {
        pn->u.packer.error_count = 0;
    } else {
        pn->u.unpacker.error_count = 0;
    }
}

void pkt_normalizer_flush(pn_struct* pn) {
    if (!pn || !pn->is_packer) return;
    if (pn->u.packer.len > 0) {
        send_buf(pn);
    }
}

static void unpacker_handler(ll_queue_t* q, ll_entry_t* unused, void* arg) {
    (void)unused;
    pn_struct* pn = arg;
    while (queue_entry_count(q) > 0) {
        ll_entry_t* entry = queue_entry_get(q);
        uint8_t* data = ll_entry_data(entry);
        size_t total = ll_entry_size(entry);
        uint16_t payload_len = *(uint16_t*)data;
        if (total != 2 + (size_t)payload_len) {
            queue_entry_free(entry);
            continue;
        }
        uint8_t* cg = data + 2;
        size_t cg_pos = 0;
        size_t cg_len = (size_t)payload_len;
        while (cg_pos < cg_len) {
            uint8_t byte = cg[cg_pos++];
            
            if (byte == 0xFF) {
                /* Начало нового фрагментированного пакета */
                if (pn->u.unpacker.in_fragment) {
                    /* Не завершен предыдущий фрагмент - ошибка */
                    pn->u.unpacker.error_count++;
                    reset_fragment_state(pn);
                }
                /* Проверить, что есть 2 байта для общей длины */
                if (cg_pos + 2 > cg_len) {
                    pn->u.unpacker.error_count++;
                    goto err;
                }
                /* Прочитать общую длину */
                uint16_t total_len = ((uint16_t)cg[cg_pos] << 8) | cg[cg_pos + 1];
                cg_pos += 2;
                size_t chunk_len = cg_len - cg_pos;
                /* Выделить буфер при необходимости */
                size_t new_len = chunk_len;
                if (new_len > pn->u.unpacker.cap) {
                    size_t new_cap = pn->u.unpacker.cap ? pn->u.unpacker.cap * 2 : 4096;
                    if (new_cap < new_len) new_cap = new_len;
                    pn->u.unpacker.buf = realloc(pn->u.unpacker.buf, new_cap);
                    pn->u.unpacker.cap = new_cap;
                }
                memcpy(pn->u.unpacker.buf, cg + cg_pos, chunk_len);
                pn->u.unpacker.len = chunk_len;
                pn->u.unpacker.total_len = total_len;
                pn->u.unpacker.in_fragment = 1;
                set_fragment_timeout(pn);
                cg_pos += chunk_len;
                /* Проверить, не собрали ли уже весь пакет */
                if (pn->u.unpacker.len >= pn->u.unpacker.total_len) {
                    if (pn->u.unpacker.len == pn->u.unpacker.total_len) {
                        ll_entry_t* out = queue_entry_new(pn->u.unpacker.total_len);
                        if (out) {
                            memcpy(ll_entry_data(out), pn->u.unpacker.buf, pn->u.unpacker.total_len);
                            queue_entry_put(pn->output, out);
                        }
                    } else {
                        /* Слишком много данных - ошибка */
                        pn->u.unpacker.error_count++;
                    }
                    reset_fragment_state(pn);
                }
                continue;
            }
            
            if (byte == 0xFE) {
                /* Продолжение фрагментированного пакета */
                if (!pn->u.unpacker.in_fragment) {
                    /* Не было начала фрагмента - ошибка */
                    pn->u.unpacker.error_count++;
                    goto err;
                }
                size_t chunk_len = cg_len - cg_pos;
                size_t new_len = pn->u.unpacker.len + chunk_len;
                if (new_len > pn->u.unpacker.cap) {
                    size_t new_cap = pn->u.unpacker.cap ? pn->u.unpacker.cap * 2 : 4096;
                    if (new_cap < new_len) new_cap = new_len;
                    pn->u.unpacker.buf = realloc(pn->u.unpacker.buf, new_cap);
                    pn->u.unpacker.cap = new_cap;
                }
                memcpy(pn->u.unpacker.buf + pn->u.unpacker.len, cg + cg_pos, chunk_len);
                pn->u.unpacker.len = new_len;
                cg_pos += chunk_len;
                set_fragment_timeout(pn);
                /* Проверить, не собрали ли уже весь пакет */
                if (pn->u.unpacker.len >= pn->u.unpacker.total_len) {
                    if (pn->u.unpacker.len == pn->u.unpacker.total_len) {
                        ll_entry_t* out = queue_entry_new(pn->u.unpacker.total_len);
                        if (out) {
                            memcpy(ll_entry_data(out), pn->u.unpacker.buf, pn->u.unpacker.total_len);
                            queue_entry_put(pn->output, out);
                        }
                    } else {
                        /* Слишком много данных - ошибка */
                        pn->u.unpacker.error_count++;
                    }
                    reset_fragment_state(pn);
                }
                continue;
            }
            
            /* Обычная запись (не фрагмент) */
            size_t L;
            if (byte <= 0xEF) {
                L = byte;
            } else if (byte >= 0xF0 && byte <= 0xFE) {
                if (cg_pos >= cg_len) goto err;
                uint8_t ext = cg[cg_pos++];
                L = ((size_t)(byte - 0xF0) << 8) | ext;
            } else {
                /* Недопустимый байт */
                goto err;
            }
            
            if (cg_pos + L > cg_len) goto err;
            
            if (pn->u.unpacker.in_fragment) {
                /* Это последний фрагмент в виде обычной записи */
                size_t new_len = pn->u.unpacker.len + L;
                if (new_len > pn->u.unpacker.cap) {
                    size_t new_cap = pn->u.unpacker.cap ? pn->u.unpacker.cap * 2 : 4096;
                    if (new_cap < new_len) new_cap = new_len;
                    pn->u.unpacker.buf = realloc(pn->u.unpacker.buf, new_cap);
                    pn->u.unpacker.cap = new_cap;
                }
                memcpy(pn->u.unpacker.buf + pn->u.unpacker.len, cg + cg_pos, L);
                pn->u.unpacker.len = new_len;
                cg_pos += L;
                set_fragment_timeout(pn);
                
                /* Проверить, собрали ли весь пакет */
                if (pn->u.unpacker.len >= pn->u.unpacker.total_len) {
                    if (pn->u.unpacker.len == pn->u.unpacker.total_len) {
                        ll_entry_t* out = queue_entry_new(pn->u.unpacker.total_len);
                        if (out) {
                            memcpy(ll_entry_data(out), pn->u.unpacker.buf, pn->u.unpacker.total_len);
                            queue_entry_put(pn->output, out);
                        }
                    } else {
                        /* Слишком много данных - ошибка */
                        pn->u.unpacker.error_count++;
                    }
                    reset_fragment_state(pn);
                }
            } else {
                /* Обычная запись (не часть фрагмента) */
                ll_entry_t* out = queue_entry_new(L);
                if (out) {
                    memcpy(ll_entry_data(out), cg + cg_pos, L);
                    queue_entry_put(pn->output, out);
                }
                cg_pos += L;
            }
        }
    err:
        queue_entry_free(entry);
    }
    queue_resume_callback(q);
}