#include <c_NlpNgramGraph.h>
#include <c_Memory.h>

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

/* FNV-1a 32位 滾動哈希，用於對變長 Unicode 陣列計算主桶索引 */
C_STATIC_FORCE_INLINE
c_size_t c_NlpNgram_HashContext(const c_ucs4_t* context, c_size_t len, c_size_t bucket_count) {
    unsigned int hash = 2166136261U;
    for (c_size_t i = 0; i < len; i++) {
        hash ^= context[i];
        hash *= 16777619U;
    }
    return (c_size_t)(hash % bucket_count);
}

/* 比較兩個上下文陣列是否完全一致 */
C_STATIC_FORCE_INLINE
c_bool_t c_NlpNgram_ContextEquals(const c_ucs4_t* ctx1, const c_ucs4_t* ctx2, c_size_t len) {
    for (c_size_t i = 0; i < len; i++) {
        if (ctx1[i] != ctx2[i]) return C_FALSE;
    }
    return C_TRUE;
}

/* 建立並初始化一個全新的歷史狀態節點 */
C_STATIC_FORCE_INLINE
c_NgramNode_t* c_NgramNode_Create(const c_ucs4_t* context, c_size_t ctx_len, c_size_t rel_cap) {
    c_NgramNode_t* node = (c_NgramNode_t*)C_ALLOC(sizeof(c_NgramNode_t));
    if (!node) return NULL;

    node->frequency = 0;
    node->relation_capacity = rel_cap;
    node->next = NULL;

    // 深拷貝上下文狀態陣列
    node->context = (c_ucs4_t*)C_ALLOC(sizeof(c_ucs4_t) * ctx_len);
    node->next_chars = (c_NgramEdge_t**)C_ALLOC(sizeof(c_NgramEdge_t*) * rel_cap);

    if (!node->context || !node->next_chars) {
        C_FREE(node->context);
        C_FREE(node->next_chars);
        C_FREE(node);
        return NULL;
    }

    memcpy(node->context, context, sizeof(c_ucs4_t) * ctx_len);
    for (c_size_t i = 0; i < rel_cap; i++) {
        node->next_chars[i] = NULL;
    }
    return node;
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

c_err_t c_NlpNgramGraph_Init(c_NlpNgramGraph_t* self, c_size_t n_value, c_size_t bucket_count, c_size_t relation_default_capacity) {
    if (!self || n_value < 2 || bucket_count == 0 || relation_default_capacity == 0) return C_ERR_PARAM;

    self->n_value = n_value;
    self->bucket_count = bucket_count;
    self->relation_default_capacity = relation_default_capacity;
    self->total_states = 0;

    self->buckets = (c_NgramNode_t**)C_ALLOC(sizeof(c_NgramNode_t*) * bucket_count);
    if (!self->buckets) return C_ERR_NOMEM;

    for (c_size_t i = 0; i < bucket_count; i++) {
        self->buckets[i] = NULL;
    }
    return C_ERR_OK;
}

/* 拉鏈法安全獲取或建立上下文狀態節點 */
static c_NgramNode_t* c_NlpNgramGraph_GetOrCreate(c_NlpNgramGraph_t* self, const c_ucs4_t* context) {
    c_size_t ctx_len = self->n_value - 1;
    c_size_t idx = c_NlpNgram_HashContext(context, ctx_len, self->bucket_count);

    c_NgramNode_t* curr = self->buckets[idx];
    c_NgramNode_t* prev = NULL;

    while (curr != NULL) {
        if (c_NlpNgram_ContextEquals(curr->context, context, ctx_len)) {
            return curr; // 命中已有歷史狀態
        }
        prev = curr;
        curr = curr->next;
    }

    // 未命中，建立新歷史狀態節點
    c_NgramNode_t* new_node = c_NgramNode_Create(context, ctx_len, self->relation_default_capacity);
    if (!new_node) return NULL;

    if (prev == NULL) {
        self->buckets[idx] = new_node;
    } else {
        prev->next = new_node;
    }
    self->total_states++;
    return new_node;
}

/* 核心注入：滑動視窗掃描整篇文章語料 */
c_err_t c_NlpNgramGraph_Process(c_NlpNgramGraph_t* self, const c_ucs4_t* unicode_array, c_size_t length) {
    if (!self || !self->buckets || !unicode_array || length == 0) return C_ERR_PARAM;

    c_size_t n = self->n_value;
    // 如果文章字數小於 N 元模型所需的長度，無法建立狀態
    if (length < n) return C_ERR_OK;

    // 滑動視窗上界為 length - n + 1
    for (c_size_t i = 0; i <= length - n; i++) {
        // 1. 當前滑動視窗的起點指標，指向長度為 N-1 的歷史上下文
        const c_ucs4_t* current_context = &unicode_array[i];

        // 2. 獲取或建立對應的多元歷史節點
        c_NgramNode_t* node = c_NlpNgramGraph_GetOrCreate(self, current_context);
        if (!node) return C_ERR_NOMEM;

        node->frequency++;

        // 3. 提取第 N 個字作為預測轉移目標
        c_ucs4_t target_char = unicode_array[i + n - 1];

        // 4. 將預測目標掛入該歷史節點的 next_chars 邊哈希拉鏈中
        c_size_t edge_idx = (c_size_t)(target_char % node->relation_capacity);
        c_NgramEdge_t* edge = node->next_chars[edge_idx];
        c_bool_t edge_found = C_FALSE;

        while (edge != NULL) {
            if (edge->target_cp == target_char) {
                edge->co_count++;
                edge_found = C_TRUE;
                break;
            }
            edge = edge->next;
        }

        if (!edge_found) {
            c_NgramEdge_t* new_edge = (c_NgramEdge_t*)C_ALLOC(sizeof(c_NgramEdge_t));
            if (!new_edge) return C_ERR_NOMEM;
            new_edge->target_cp = target_char;
            new_edge->co_count = 1;
            new_edge->next = node->next_chars[edge_idx];
            node->next_chars[edge_idx] = new_edge;
        }
    }
    return C_ERR_OK;
}

void c_NlpNgramGraph_Destroy(c_NlpNgramGraph_t* self) {
    if (!self || !self->buckets || self->bucket_count == 0) return;

    c_NgramNode_t** local_buckets = self->buckets;
    c_size_t local_count = self->bucket_count;

    // 搶先阻斷外部二次訪問漏洞
    self->buckets = NULL;
    self->bucket_count = 0;
    self->total_states = 0;

    for (c_size_t i = 0; i < local_count; i++) {
        c_NgramNode_t* node = local_buckets[i];

        while (node != NULL) {
            c_NgramNode_t* next_node = node->next;

            // 1. 釋放深層歷史上下文陣列
            if (node->context != NULL) {
                C_FREE(node->context);
            }

            // 2. 釋放後續預測鄰接邊拉鏈
            if (node->next_chars != NULL) {
                for (c_size_t k = 0; k < node->relation_capacity; k++) {
                    c_NgramEdge_t* edge = node->next_chars[k];
                    while (edge != NULL) {
                        c_NgramEdge_t* temp_edge = edge->next;
                        C_FREE(edge);
                        edge = temp_edge;
                    }
                }
                C_FREE(node->next_chars);
            }

            // 3. 釋放狀態節點本身
            C_FREE(node);
            node = next_node;
        }
    }
    C_FREE(local_buckets);
}

c_err_t c_NlpNgramGraph_PredictNext(c_NlpNgramGraph_t* self, const c_ucs4_t* history_context,
                                     c_ucs4_t* out_next_char, double* out_confidence) {
    if (!self || !self->buckets || !history_context || !out_next_char || !out_confidence) {
        return C_ERR_PARAM;
    }

    c_size_t ctx_len = self->n_value - 1;
    c_size_t idx = c_NlpNgram_HashContext(history_context, ctx_len, self->bucket_count);

    // 1. 定位多元歷史狀態
    c_NgramNode_t* node = self->buckets[idx];
    while (node != NULL) {
        if (c_NlpNgram_ContextEquals(node->context, history_context, ctx_len)) {
            break;
        }
        node = node->next;
    }

    // 如果語料庫中從未出現過這段連續的歷史短语（未命中）
    if (!node || node->frequency == 0) {
        return C_ERR_FAIL;
    }

    // 2. 統計該狀態下的邊權重，找出概率最高的轉移字符
    c_size_t total_co_count = 0;
    c_size_t max_co_count = 0;
    c_ucs4_t best_target_cp = 0;

    for (c_size_t i = 0; i < node->relation_capacity; i++) {
        c_NgramEdge_t* edge = node->next_chars[i];
        while (edge != NULL) {
            total_co_count += edge->co_count;
            if (edge->co_count > max_co_count) {
                max_co_count = edge->co_count;
                best_target_cp = edge->target_cp;
            }
            edge = edge->next;
        }
    }

    if (total_co_count == 0) return C_ERR_FAIL;

    // 输出預測最優解與馬爾可夫最大似然概率
    *out_next_char = best_target_cp;
    *out_confidence = (double)max_co_count / (double)total_co_count;

    return C_ERR_OK;
}