#include <stdio.h>
#include <stdint.h>
#include <malloc.h>

typedef struct _TrieTreeNode {
    uint8_t u8Lvl; // 用于调试打印
    uint8_t u8Val;
    uint8_t u8ChdNum;
    void (*pCbk)();
    struct _TrieTreeNode *pChd;
    struct _TrieTreeNode *pBro;
    struct _TrieTreeNode *pPar;
} TrieNode_t;

TrieNode_t g_stPrefixTreeHead = {
    .u8Lvl = 0,
    .u8Val = 0x0,
    .u8ChdNum = 0,
    .pCbk = NULL,
    .pChd = NULL,
    .pBro = NULL,
    .pPar = NULL,
};

TrieNode_t *createNode()
{
    TrieNode_t *pstNode = (TrieNode_t *)malloc(sizeof(TrieNode_t));
    pstNode->u8Lvl = 0;
    pstNode->u8Val = 0x0;
    pstNode->u8ChdNum = 0;
    pstNode->pCbk = NULL;
    pstNode->pChd = NULL;
    pstNode->pBro = NULL;
    pstNode->pPar = NULL;

    return pstNode;
}

void addNode(int iLen, uint8_t *pu8Prefix, void(*pCbk)())
{
    int i = 0;
    TrieNode_t *pstNode = &g_stPrefixTreeHead;
    for (i = 0; i < iLen; i++) {
        if (pstNode->u8ChdNum == 0) {
            pstNode->pChd = createNode();
            pstNode->u8ChdNum += 1;
            pstNode->pChd->pPar = pstNode;
            pstNode->pChd->u8Lvl = i + 1;
            pstNode->pChd->u8Val = pu8Prefix[i];
            pstNode = pstNode->pChd;
        } else {
            pstNode = pstNode ->pChd;
            while (pstNode->u8Val != pu8Prefix[i] && pstNode->pBro != NULL)
                pstNode = pstNode->pBro;
            if (pstNode->u8Val != pu8Prefix[i] && pstNode->pBro == NULL) {
                pstNode->pBro = createNode();
                pstNode->pBro->pPar = pstNode->pPar;
                pstNode->pBro->u8Lvl = i + 1;
                pstNode->pBro->u8Val = pu8Prefix[i];
                pstNode = pstNode->pBro;
            }
        }
    }
    pstNode->pCbk = pCbk;
}

void searchNode(int iLen, uint8_t *pu8Data, void(**pCbk)())
{
    if (g_stPrefixTreeHead.u8ChdNum == 0)
        *pCbk = NULL;

    int i = 0;
    TrieNode_t *pstNode = g_stPrefixTreeHead.pChd;
    for (i = 0; i < iLen; i++) {
        while (pstNode->u8Val != pu8Data[i] && pstNode->pBro != NULL)
            pstNode = pstNode->pBro;
        if (pstNode->u8Val != pu8Data[i]) {
            pstNode = pstNode->pPar;
            break;
        } else if (pstNode->pChd == NULL) {
            break;
        }
        pstNode = pstNode->pChd;
    }
    while (pstNode->pCbk == NULL && pstNode->pPar != NULL)
        pstNode = pstNode->pPar;
    *pCbk = pstNode->pCbk;
}

void printTree(TrieNode_t *pstNode)
{
    int i = 0;
    for (i = 0; i < pstNode->u8Lvl; i++)
        printf("-");
    printf("[val=0x%x,fun=%p,chd=", pstNode->u8Val, pstNode->pCbk);
    if (pstNode->u8ChdNum != 0) {
        printf("\n");
        printTree(pstNode->pChd);
    } else {
        printf("%p", pstNode->pChd);
    }
    printf(",bro=");
    if (pstNode->pBro != NULL) {
        printf("\n");
        printTree(pstNode->pBro);
    } else {
        printf("%p", pstNode->pBro);
    }
    printf("]");
}

void fun1(){};
void fun2(){};
void fun3(){};

int main()
{
    printf("TrieTree\n");

    uint8_t pu8Prefix[] = {0x31, 0x01, 0xC1, 0x03};

    addNode(4, pu8Prefix, fun1);
    addNode(4, (uint8_t[]){0x31, 0x01, 0xC1, 0x04}, fun2);
    addNode(4, (uint8_t[]){0x31, 0x02, 0xC3, 0x13}, fun3);
    addNode(1, (uint8_t[]){0x22}, fun3);
    addNode(1, (uint8_t[]){0x3E}, fun3);
    addNode(2, (uint8_t[]){0x3E, 0x00}, fun1);

    printf("Add done!\n");

    printTree(&g_stPrefixTreeHead);

    printf("\n\n");

    void (*pCbk)();
    searchNode(6, (uint8_t[]){0x31, 0x02, 0xC3, 0x13, 0x01, 0x02}, &pCbk);
    printf("pCbk = %p\n", pCbk);
    searchNode(6, (uint8_t[]){0x31, 0x02, 0xC3, 0x14, 0x01, 0x02}, &pCbk);
    printf("pCbk = %p\n", pCbk);
    searchNode(0, NULL, &pCbk);
    printf("pCbk = %p\n", pCbk);
    searchNode(6, (uint8_t[]){0x3E, 0x02, 0xC3, 0x14, 0x01, 0x02}, &pCbk);
    printf("pCbk = %p\n", pCbk);

    return 0;
}