#include "hv_drv_Osd.h"

#define BITS_PER_LONG   32

static inline UCHAR8 find_first_bit(UINT32 uiWord)
{
    UCHAR8 ucNum = 0;

    if ((uiWord & 0xffff) == 0)
    {
        ucNum += 16;
        uiWord >>= 16;
    }
    if ((uiWord & 0xff) == 0)
    {
        ucNum += 8;
        uiWord >>= 8;
    }
    if ((uiWord & 0xf) == 0)
    {
        ucNum += 4;
        uiWord >>= 4;
    }
    if ((uiWord & 0x3) == 0)
    {
        ucNum += 2;
        uiWord >>= 2;
    }
    if ((uiWord & 0x1) == 0)
    {
        ucNum += 1;
    }
    return ucNum;
}

static inline UCHAR8 find_last_bit(UINT32 uiWord)
{
    UCHAR8 ucNum = BITS_PER_LONG - 1;

    if (!(uiWord & (~0ul << (BITS_PER_LONG-16))))
    {
        ucNum -= 16;
        uiWord <<= 16;
    }
    if (!(uiWord & (~0ul << (BITS_PER_LONG-8))))
    {
        ucNum -= 8;
        uiWord <<= 8;
    }
    if (!(uiWord & (~0ul << (BITS_PER_LONG-4))))
    {
        ucNum -= 4;
        uiWord <<= 4;
    }
    if (!(uiWord & (~0ul << (BITS_PER_LONG-2))))
    {
        ucNum -= 2;
        uiWord <<= 2;
    }
    if (!(uiWord & (~0ul << (BITS_PER_LONG-1))))
    {
        ucNum -= 1;
    }

    return ucNum;
}

static VOID Hv_Drv_OsdRadix_FreeSlots(OsdRadixTree *pstTree, USHORT16 usIndex, UCHAR8 ucSlots)
{
    USHORT16 *pusNodes = NULL;
    USHORT16 usInnerSlots = 0;

    usInnerSlots = (1 << pstTree->ucHeight) - 1;
    usIndex += usInnerSlots;
    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);
    while (1)
    {
        pusNodes[usIndex] += ucSlots;
        if (0 == usIndex)
        {
            break;
        }

        usIndex = (usIndex - 1) / 2;
    }
}

VOID Hv_Drv_OsdRadix_FreeSlot(OsdRadixTree *pstTree, USHORT16 usOff)
{
    USHORT16 *pusNodes = NULL;
    UINT32 *puiLeaf = NULL;
    USHORT16 usLeafOff = 0;
    UCHAR8 ucBit = 0;

    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);
    usLeafOff = usOff >> 5;
    ucBit = usOff & 0x1f;

    puiLeaf = pstTree->puiBuffer + usLeafOff;

    HV_ASSERT_TRUE_VOID((*puiLeaf & (1 << ucBit)) != 0);

    *puiLeaf &= ~(1 << ucBit);

    Hv_Drv_OsdRadix_FreeSlots(pstTree, usLeafOff, 1);

    OSD_DRV_LOG_DEBUG("radix tree free slot: offset %d free slots left %d",
            usOff, pusNodes[0]);
}

static VOID Hv_Drv_OsdRadix_AllocSlots(OsdRadixTree *pstTree, USHORT16 usIndex, UCHAR8 ucSlots)
{
    USHORT16 *pusNodes = NULL;
    USHORT16 usInnerSlots = 0;
    UINT32 *puiLeaf = pstTree->puiBuffer + usIndex;

    if (32 == ucSlots)
    {
        *puiLeaf = 0xffffffff;
    }
    else
    {
        *puiLeaf |= ((1 << ucSlots) - 1);
    }
    usInnerSlots = (1 << pstTree->ucHeight) - 1;
    usIndex += usInnerSlots;
    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);

    while (1)
    {
        pusNodes[usIndex] -= ucSlots;
        if (0 == usIndex)
        {
            break;
        }

        usIndex = (usIndex - 1) / 2;
    }
}

/*
 * Mark first usCount slots as allocated
 */
VOID Hv_Drv_OsdRadix_MarkSlots(OsdRadixTree *pstTree, USHORT16 usSlots)
{
    USHORT16 usIndex = 0;
    UCHAR8 ucRmd = usSlots & 0x1f;
    UCHAR8 ucQuo = usSlots >> 5;

    for (usIndex = 0; usIndex < ucQuo; ++usIndex)
    {
        Hv_Drv_OsdRadix_AllocSlots(pstTree, usIndex, 32);
    }

    if (0 != ucRmd)
    {
        Hv_Drv_OsdRadix_AllocSlots(pstTree, usIndex, ucRmd);
    }
}

USHORT16 Hv_Drv_OsdRadix_GetLastAllocatedSlot(OsdRadixTree *pstTree)
{
    return pstTree->usLastAllocatedSlot;
}

BOOL Hv_Drv_OsdRadix_AllocSlot(OsdRadixTree *pstTree, BOOL bFromRightHalf, USHORT16 *pusSlot)
{
    USHORT16 *pusNodes = NULL;
    UCHAR8 ucIndex = 0;
    USHORT16 usLeft = 0;
    USHORT16 usRight = 0;
    USHORT16 usOff = 0;
    UINT32 *puiLeaf = NULL;
    USHORT16 usInnerSlots = (1 << pstTree->ucHeight) - 1;
    UCHAR8 ucFirstBit = 0;

    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);

    /* 根节点表示所有空闲的slot个数 */
    HV_ASSERT_TRUE_RET(pusNodes[0] > 0, HV_FALSE);
    if (pstTree->usTotalSlots > OSD_BANK_CAPACITY)
    {
        if (bFromRightHalf)
        {
            HV_ASSERT_TRUE_RET(pusNodes[2] > 0, HV_FALSE);

            usOff = 2;
        }
        else
        {
            HV_ASSERT_TRUE_RET(pusNodes[1] > 0, HV_FALSE);

            usOff = 1;
        }

        ucIndex = 1;
        --pusNodes[0];
    }

    for (; ucIndex < pstTree->ucHeight; ++ucIndex)
    {
        --pusNodes[usOff];
        if (pusNodes[2 * usOff + 1] > 0)
        {
            usOff = 2 * usOff + 1;
        }
        else
        {
            usOff = 2 * usOff + 2;
        }
    }
    /* arrives leaf node */
    --pusNodes[usOff];
    /* leaf offset */
    usOff -= usInnerSlots;
    puiLeaf = pstTree->puiBuffer + usOff;
    ucFirstBit = find_first_bit(~(*puiLeaf));
    *puiLeaf |= (1 << ucFirstBit);

    usOff = (usOff << 5) + ucFirstBit;
    if ((pstTree->usTotalSlots > OSD_BANK_CAPACITY) && bFromRightHalf)
    {
        usOff -= OSD_BANK_CAPACITY;
    }

    *pusSlot = pstTree->usLastAllocatedSlot = usOff;

    OSD_DRV_LOG_DEBUG("radix tree alloc slot: offset %d from %s child. free slots left %d",
            usOff, bFromRightHalf ? "right" : "left", pusNodes[0]);

    return HV_TRUE;
}

VOID Hv_Drv_OsdRadix_InitializeTree(OsdRadixTree *pstTree)
{
    USHORT16 usIndex = 0;
    USHORT16 usInnerSlots = 0;
    USHORT16 usLeafCounts = pstTree->usLeafCounts;
    USHORT16 usLen = 0;
    UCHAR8 ucSlots = 32;

    usInnerSlots = (1 << pstTree->ucHeight) - 1;
    usLen = (usLeafCounts * 4) + (usInnerSlots + usLeafCounts) * 2;
    HV_MEMSET(pstTree->puiBuffer, 0, usLen);

    for (usIndex = 0; usIndex < usLeafCounts; ++usIndex)
    {
        if ((usIndex == usLeafCounts - 1))
        {
            ucSlots = pstTree->usTotalSlots & 0x1f;
            if (0 == ucSlots)
            {
                ucSlots = 32;
            }
        }
        Hv_Drv_OsdRadix_FreeSlots(pstTree, usIndex, ucSlots);
    }
}

/* check which half has more empty slots */
BOOL Hv_Drv_OsdRadix_IsLeftLeaned(OsdRadixTree *pstTree)
{
    USHORT16 *pusNodes = NULL;

    if (pstTree->usTotalSlots <= OSD_BANK_CAPACITY)
    {
        return HV_TRUE;
    }

    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);
    return pusNodes[1] >= pusNodes[2];
}

USHORT16 Hv_Drv_OsdRadix_GetEmptySlots(OsdRadixTree *pstTree)
{
    USHORT16 *pusNodes = NULL;

    if (pstTree == NULL)
    {
        return 0;
    }

    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);
    return pusNodes[0];
}

USHORT16 Hv_Drv_OsdRadix_GetLeftEmptySlots(OsdRadixTree *pstTree)
{
    USHORT16 *pusNodes = NULL;

    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);
    if (pstTree->usTotalSlots <= OSD_BANK_CAPACITY)
    {
        return pusNodes[0];
    }

    return pusNodes[1];
}

USHORT16 Hv_Drv_OsdRadix_GetRightEmptySlots(OsdRadixTree *pstTree)
{
    USHORT16 *pusNodes = NULL;

    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);
    if (pstTree->usTotalSlots <= OSD_BANK_CAPACITY)
    {
        return pusNodes[0];
    }

    return pusNodes[2];
}

USHORT16 Hv_Drv_OsdRadix_GetTotalSlots(OsdRadixTree *pstTree)
{
    if (pstTree == NULL)
    {
        return 0;
    }

    return pstTree->usTotalSlots;
}

OsdRadixTree* Hv_Drv_OsdRadix_ExpandTree(OsdRadixTree *pstTree, USHORT16 usSlots)
{
    OsdRadixTree *pstNewTree = NULL;
    USHORT16 *pusNodes = NULL;
    USHORT16 *pusNewNodes = NULL;
    USHORT16 usInnerSlots = 0;
    USHORT16 usNewInnerSlots = 0;
    USHORT16 usLeafIndex = 0;
    USHORT16 usIndex = 0;
    UCHAR8 ucAllocatedSlots = 0;

    /* 1. build an empty tree */
    pstNewTree = Hv_Drv_OsdRadix_BuildTree(pstTree->usTotalSlots + usSlots);
    if (NULL == pstNewTree)
    {
        OSD_DRV_LOG_ERROR("could not expand tree, not enough space");
        return NULL;
    }

    /* 2. copy bitmaps from old tree */
    HV_MEMCPY((VOID*)pstNewTree->puiBuffer, pstTree->puiBuffer,
            pstTree->usLeafCounts * 4);

    /* 3. mark allocated slots */
    usInnerSlots = (1 << pstTree->ucHeight) - 1;
    usNewInnerSlots = (1 << pstNewTree->ucHeight) - 1;
    pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);
    pusNewNodes = (USHORT16*)(pstNewTree->puiBuffer + pstNewTree->usLeafCounts);
    for (usIndex = 0; usIndex < pstTree->usLeafCounts; ++usIndex)
    {
        if (usIndex == pstTree->usLeafCounts - 1)
        {
            /* last leaf maybe part-full */
            ucAllocatedSlots = (pstTree->usTotalSlots - (usIndex << 5)) - pusNodes[usInnerSlots + usIndex];
        }
        else
        {
            ucAllocatedSlots = 32 - pusNodes[usInnerSlots + usIndex];
        }

        usLeafIndex = usNewInnerSlots + usIndex;

        while (1)
        {
            pusNewNodes[usLeafIndex] -= ucAllocatedSlots;
            if (0 == usLeafIndex)
            {
                break;
            }

            usLeafIndex = (usLeafIndex - 1) / 2;
        }
    }

    OSD_DRV_LOG_DEBUG("expand radix tree from %d slots to %d slots",
            pstTree->usTotalSlots, pstNewTree->usTotalSlots);

    //Hv_Drv_OsdRadix_PrintTree(pstTree);
    //Hv_Drv_OsdRadix_PrintTree(pstNewTree);

    /* 4. destroy old tree */
    Hv_Drv_OsdRadix_DestroyTree(pstTree);

    return pstNewTree;
}

/*
 * buffer layout:
 *
 *   bitmap-0    bitmap-1 bitmap-2 ... bitmap-N treeNode-0 treeNode-1 ...
 *|<- 4bytes ->|
 *
 * tree layout:
 *              o     <- treeNode-0
 *             / \
 * treeNode-1 o   o   <- treeNode-2
 *           / \ / \
 *          o  o o  o
 */
OsdRadixTree *Hv_Drv_OsdRadix_BuildTree(USHORT16 usTotalSlots)
{
    UCHAR8 ucLastBit = 0;
    USHORT16 usCompleteSlots = 0;
    USHORT16 usInnerSlots = 0;
    USHORT16 usLeafCounts = 0;
    OsdRadixTree *pstTree = NULL;
    USHORT16 usLen = 0;
    UCHAR8 ucSlots = 0;
    USHORT16 usIndex = 0;

    HV_ASSERT_TRUE_RET(usTotalSlots != 0, NULL);

    usLeafCounts = ALIGN(usTotalSlots, 32) >> 5;
    /* 补成完全树 */
    ucLastBit = find_last_bit(usLeafCounts);
    if (usLeafCounts & (usLeafCounts - 1))
    {
        ucLastBit += 1;
    }

    usCompleteSlots = 1 << ucLastBit;
    usInnerSlots = usCompleteSlots - 1;

    pstTree = HV_MALLOC(sizeof(*pstTree));

    HV_ASSERT_TRUE_RET(pstTree != NULL, NULL);

    HV_MEMSET(pstTree, 0, sizeof(*pstTree));
    pstTree->usTotalSlots = usTotalSlots;
    pstTree->usLeafCounts = usLeafCounts;
    pstTree->ucHeight = ucLastBit;

    usLen = (usLeafCounts * 4) + (usInnerSlots + usLeafCounts) * 2;
    pstTree->puiBuffer = HV_MALLOC(usLen);

    Hv_Drv_OsdRadix_InitializeTree(pstTree);

    OSD_DRV_LOG_DEBUG("build radix tree: total slots: %d height %d leaf counts %d inner slots %d buffer length %d",
            usTotalSlots, pstTree->ucHeight, pstTree->usLeafCounts, usInnerSlots, usLen);

    return pstTree;
}

VOID Hv_Drv_OsdRadix_DestroyTree(OsdRadixTree *pstTree)
{
    if (NULL == pstTree)
    {
        return;
    }

    HV_FREE(pstTree->puiBuffer);
    HV_FREE(pstTree);
}

UCHAR8* Hv_Drv_OsdRadix_PrintBitmap(UINT32 uiData)
{
    static UCHAR8 aucBuffer[64];
    UCHAR8 *pucBuffer = aucBuffer;
    UCHAR8 ucIndex = 0;

    for (ucIndex = 0; ucIndex < 32; ++ucIndex)
    {
        if ((uiData >> ucIndex) & 0x01)
        {
            pucBuffer += HV_SPRINTF(pucBuffer, "1");
        }
        else
        {
            pucBuffer += HV_SPRINTF(pucBuffer, "0");
        }
    }

    return aucBuffer;
}

VOID Hv_Drv_OsdRadix_PrintTree(OsdRadixTree *pstTree)
{
    UCHAR8 ucIndex = 0;
    UCHAR8 ucRowLeafCount = 8;
    UCHAR8 ucRowIndex = 0;
    UCHAR8 aucBuffer[1024];
    UCHAR8 *pucBuffer = NULL;
    USHORT16 *pusNodes = (USHORT16*)(pstTree->puiBuffer + pstTree->usLeafCounts);
    UINT32 *puiBuffer = pstTree->puiBuffer;

    OSD_LOG_INFO("Radix Tree: total slots: %d free slots: %d height: %d leaf counts: %d last allocated slot: %d",
            pstTree->usTotalSlots, pusNodes[0], pstTree->ucHeight, pstTree->usLeafCounts, pstTree->usLastAllocatedSlot);

    for (ucIndex = 0; ucIndex < pstTree->ucHeight; ++ucIndex)
    {
        pucBuffer = aucBuffer;
        for (ucRowIndex = 0; ucRowIndex < (1 << ucIndex); ++ucRowIndex)
        {
            pucBuffer += HV_SPRINTF(pucBuffer, "%d  ", pusNodes[((1 << ucIndex) - 1) + ucRowIndex]);
        }
        OSD_LOG_INFO("%s", aucBuffer);
    }

    pucBuffer = aucBuffer;
    for (ucIndex = 0; ucIndex < pstTree->usLeafCounts; ++ucIndex)
    {
        pucBuffer += HV_SPRINTF(pucBuffer, "%d  ", pusNodes[((1 << pstTree->ucHeight) - 1) + ucIndex]);
    }
    OSD_LOG_INFO("%s", aucBuffer);

    pucBuffer = aucBuffer;
    for (ucIndex = 0; ucIndex < pstTree->usLeafCounts; ++ucIndex)
    {
        pucBuffer += HV_SPRINTF(pucBuffer, "%s  ", Hv_Drv_OsdRadix_PrintBitmap(puiBuffer[ucIndex]));

        if ((ucIndex & 0x7) == 7)
        {
            OSD_LOG_INFO("%s", aucBuffer);

            pucBuffer = aucBuffer;
        }
    }

    if (pucBuffer != aucBuffer)
    {
        OSD_LOG_INFO("%s", aucBuffer);
    }
}
