#include "behaviour_tree.h"

#include "../memory/memory.h"
#include "../util/object_pool.h"

//#define GET_NODE(idx) ((BzAIBTNode *) bzObjectPoolGetObject(bt->nodePool, idx))

struct BzAIBTNode {
    BzAIBTNode *parent;
    // Children
    BzAIBTNode *first;
    BzAIBTNode *last;
    // Siblings
    BzAIBTNode *prev;
    BzAIBTNode *next;

    BzAIBTNodeType type;
    union {
        struct {
            i32 n;
        } repeat;
        struct {
            f32 ms;
        } delay;
        struct {
            BzAIBTActionFn fn;
        } action;
    } as;

};

struct BzAIBTNodeState {
    const BzAIBTNode *node;
    BzAIBTNodeState *next;
    BzAIBTNodeState *prev;

    union {
        struct {
            i32 iter;
        } repeat;
        struct {
            f32 elapsed;
        } delay;
    } as;
};


size_t bzAIBTGetNodeSize() {
    return sizeof(BzAIBTNode);
}
size_t  bzAIBTGetNodeStateSize() {
    return sizeof(BzAIBTNodeState);
}

static BzAIBTNode *bzAIBTNodeMake(BzObjectPool *nodePool, BzAIBTNode *parent, BzAIBTNodeType type) {
    BZ_ASSERT(nodePool);
    BZ_ASSERT(bzObjectPoolGetObjectSize(nodePool) == bzAIBTGetNodeSize());
    BzAIBTNode *node = bzObjectPool(nodePool);
    bzMemSet(node, 0, sizeof(*node));

    node->type = type;

    if (parent && parent->last) {
        parent->last->next = node;
        node->prev = parent->last;
        parent->last = node;
    } else if (parent) {
        parent->first = node;
        parent->last = node;
    }
    node->parent = parent;

    return node;
}

BzAIBTNode *bzAIBTMakeRoot(BzObjectPool *nodePool) {
    return bzAIBTNodeMake(nodePool, NULL, BZ_AIBT_DECOR_DUMMY);
}

void bzAIBTDestroyRoot(BzObjectPool *nodePool, BzAIBTNode *node) {
    BZ_ASSERT(node);
    BzAIBTNode *pNode = node;

    while (pNode) {
        BzAIBTNode *next = pNode->next;
        bzAIBTDestroyRoot(nodePool, pNode);
        pNode = next;
    }

    bzObjectPoolRelease(nodePool, node);
}

BzAIBTNode *bzAIBTCompSelector(BzObjectPool *nodePool, BzAIBTNode *parent, bool parallel) {
    BzAIBTNodeType type = parallel ?
            BZ_AIBT_COMP_PARALLEL_SELECTOR :
            BZ_AIBT_COMP_SELECTOR;
    return bzAIBTNodeMake(nodePool, parent, type);
}
BzAIBTNode *bzAIBTCompSequence(BzObjectPool *nodePool, BzAIBTNode *parent, bool parallel) {
    BzAIBTNodeType type = parallel ?
                          BZ_AIBT_COMP_PARALLEL_SEQUENCE :
                          BZ_AIBT_COMP_SEQUENCE;
    return bzAIBTNodeMake(nodePool, parent, type);
}

BzAIBTNode *bzAIBTDecorDummy(BzObjectPool *nodePool, BzAIBTNode *parent) {
    return bzAIBTNodeMake(nodePool, parent, BZ_AIBT_DECOR_DUMMY);
}
BzAIBTNode *bzAIBTDecorSuccess(BzObjectPool *nodePool, BzAIBTNode *parent) {
    return bzAIBTNodeMake(nodePool, parent, BZ_AIBT_DECOR_SUCCESS);
}
BzAIBTNode *bzAIBTDecorFail(BzObjectPool *nodePool, BzAIBTNode *parent) {
    return bzAIBTNodeMake(nodePool, parent, BZ_AIBT_DECOR_FAIL);
}
BzAIBTNode *bzAIBTDecorInvert(BzObjectPool *nodePool, BzAIBTNode *parent) {
    return bzAIBTNodeMake(nodePool, parent, BZ_AIBT_DECOR_INVERT);
}
BzAIBTNode *bzAIBTDecorUntilSuccess(BzObjectPool *nodePool, BzAIBTNode *parent) {
    return bzAIBTNodeMake(nodePool, parent, BZ_AIBT_DECOR_UNTIL_SUCCESS);
}
BzAIBTNode *bzAIBTDecorUntilFail(BzObjectPool *nodePool, BzAIBTNode *parent) {
    return bzAIBTNodeMake(nodePool, parent, BZ_AIBT_DECOR_UNTIL_FAIL);
}
BzAIBTNode *bzAIBTDecorRepeat(BzObjectPool *nodePool, BzAIBTNode *parent, i32 n) {
    BzAIBTNode *node = bzAIBTNodeMake(nodePool, parent, BZ_AIBT_DECOR_REPEAT);
    node->as.repeat.n = n;
    return node;
}
BzAIBTNode *bzAIBTDecorDelay(BzObjectPool *nodePool, BzAIBTNode *parent, f32 ms) {
    BzAIBTNode *node = bzAIBTNodeMake(nodePool, parent, BZ_AIBT_DECOR_DELAY);
    node->as.delay.ms = ms;
    return node;
}

BzAIBTNode *bzAIBTAction(BzObjectPool *nodePool, BzAIBTNode *parent, BzAIBTActionFn fn) {
    BzAIBTNode *node = bzAIBTNodeMake(nodePool, parent, BZ_AIBT_ACTION);
    node->as.action.fn = fn;
    return node;
}

BzAIBTNodeType bzAIBTGetNodeType(BzAIBTNode *node) {
    return node->type;
}

BzAIBTState bzAIBTCreateState(const BzAIBTStateDesc *desc) {
    BZ_ASSERT(desc->pool);
    BZ_ASSERT(bzObjectPoolGetObjectSize(desc->pool) == bzAIBTGetNodeStateSize());
    BZ_ASSERT(desc->root);
    return (BzAIBTState) {
            .root = desc->root,
            .first = NULL,
            .last = NULL,
            .nodeStatePool = desc->pool,
            .userData = desc->userData
    };
}
void bzAIBTDestroyState(BzAIBTState *state) {
    BzAIBTNodeState *pNodeState = state->first;
    while (pNodeState) {
        BzAIBTNodeState *next = pNodeState->next;
        bzObjectPoolRelease(state->nodeStatePool, pNodeState);
        pNodeState = next;
    }
    bzMemSet(state, 0, sizeof(*state));
}

void bzAIBTStatePush(BzAIBTState *state, BzAIBTNodeState *nodeState,
                     const BzAIBTNodeState *desc) {
    BzAIBTNodeState *newState = bzObjectPool(state->nodeStatePool);
    BZ_ASSERT(newState && desc);
    *newState = *desc;

    newState->next = NULL;
    newState->prev = NULL;

    if (nodeState == NULL)
        nodeState = state->last;

    if (nodeState) {
        BzAIBTNodeState *next = nodeState->next;
        nodeState->next = newState;
        newState->prev = nodeState;
        newState->next = next;
        if (next)
            next->prev = newState;
        if (state->last == nodeState)
            state->last = newState;
    } else {
        newState->prev = state->last;
        state->last = newState;
        if (state->first == NULL) state->first = newState;
    }
}
void bzAIBTStatePop(BzAIBTState *state, BzAIBTNodeState *nodeState) {
    if (state->first == nodeState) state->first = nodeState->next;
    if (state->last == nodeState) state->last = nodeState->prev;
    BzAIBTNodeState *next = nodeState->next;
    BzAIBTNodeState *prev = nodeState->prev;
    if (nodeState->prev)
        nodeState->prev->next = next;
    if (nodeState->next)
        nodeState->next->prev = prev;
    bzObjectPoolRelease(state->nodeStatePool, nodeState);
}

static inline BzAIBTStatus bzAIBTExecuteNode(const BzAIBTNode *node, f32 dt,
                                      BzAIBTState *state, BzAIBTNodeState *nodeState);
/*
static inline BzAIBTStatus bzAIBTExecuteComposite(const BzObjectPool *nodePool, const BzAIBTNode *node,
                                           BzAIBTState *state, BzAIBTNodeState *nodeState) {
    switch (node->type) {
        case BZ_AIBT_COMP_SELECTOR:
            for (BzAIBTNode *child = node->first; child; child = child->next) {
                BzAIBTStatus status = bzAIBTExecuteNode(bt, child);
                if (status == BZ_AIBT_SUCCESS) return status;
                if (status == BZ_AIBT_RUNNING) return status;
            }
            return BZ_AIBT_FAIL;
        case BZ_AIBT_COMP_SEQUENCE:
            for (BzAIBTNode *child = node->first; child; child = child->next) {
                BzAIBTStatus status = bzAIBTExecuteNode(bt, child);
                if (status == BZ_AIBT_FAIL) return status;
                if (status == BZ_AIBT_RUNNING) return status;
            }
            return BZ_AIBT_SUCCESS;
        default:
            assert(false);
            return BZ_AIBT_ERROR;
    }
}
 */
static inline BzAIBTStatus bzAIBTExecuteDecorator(const BzAIBTNode *node, f32 dt,
                                           BzAIBTState *state, BzAIBTNodeState *nodeState) {
    // Ensure decorator has only one child
    BZ_ASSERT(node->first && node->first == node->last);
    BzAIBTNodeState *first = nodeState;
    if (nodeState && first->node == node) {
        first = first->next;
    }

    switch (node->type) {
        case BZ_AIBT_DECOR_REPEAT:
            if (!nodeState || nodeState->node != node) {
                bzAIBTStatePush(state, nodeState, &(BzAIBTNodeState) {
                        .node = node,
                        .as.repeat.iter = 0
                });
            }
            break;
        case BZ_AIBT_DECOR_DELAY:
            if (!nodeState || nodeState->node != node) {
                bzAIBTStatePush(state, nodeState, &(BzAIBTNodeState) {
                        .node = node,
                        .as.delay = {0.2f}
                });
                return BZ_AIBT_RUNNING;
            }
            nodeState->as.delay.elapsed += 0.2f;
            if (nodeState->as.delay.elapsed < node->as.delay.ms) {
                return BZ_AIBT_RUNNING;
            }
            bzAIBTStatePop(state, nodeState);
            break;
        default:
            break;
    }

    BzAIBTStatus inStatus = bzAIBTExecuteNode(node->first, dt, state, first);

    // ERROR, RUNNING are propagated up
    if (inStatus == BZ_AIBT_ERROR)
        return BZ_AIBT_ERROR;
    if (inStatus == BZ_AIBT_RUNNING)
        return BZ_AIBT_RUNNING;

    BzAIBTStatus status = BZ_AIBT_ERROR;
    switch (node->type) {
        case BZ_AIBT_DECOR_DUMMY:
        case BZ_AIBT_DECOR_DELAY: // Delay already handled
            status = inStatus;
            break;
        case BZ_AIBT_DECOR_SUCCESS:
            status = BZ_AIBT_SUCCESS;
            break;
        case BZ_AIBT_DECOR_FAIL:
            status = BZ_AIBT_FAIL;
            break;
        case BZ_AIBT_DECOR_INVERT:
            if (inStatus == BZ_AIBT_FAIL) status = BZ_AIBT_SUCCESS;
            if (inStatus == BZ_AIBT_SUCCESS) status = BZ_AIBT_FAIL;
            break;
        case BZ_AIBT_DECOR_UNTIL_SUCCESS:
            if (inStatus == BZ_AIBT_SUCCESS)
                status = BZ_AIBT_SUCCESS;
            else
                status = BZ_AIBT_RUNNING;
            break;
        case BZ_AIBT_DECOR_UNTIL_FAIL:
            if (inStatus == BZ_AIBT_FAIL)
                status = BZ_AIBT_SUCCESS;
            else
                status = BZ_AIBT_RUNNING;
            break;
        case BZ_AIBT_DECOR_REPEAT:
            BZ_ASSERT(nodeState->node == node);
            nodeState->as.repeat.iter++;
            if (nodeState->as.repeat.iter >= node->as.repeat.n) {
                bzAIBTStatePop(state, nodeState);
                status = inStatus;
                break;
            }
            status = BZ_AIBT_RUNNING;
            break;
        default:
            break;
    }
    return status;
}
static inline BzAIBTStatus bzAIBTExecuteNode(const BzAIBTNode *node, f32 dt,
                                      BzAIBTState *state, BzAIBTNodeState *nodeState) {
    BzAIBTStatus status = BZ_AIBT_ERROR;
    switch (node->type) {
        case BZ_AIBT_COMP_SELECTOR:
        case BZ_AIBT_COMP_SEQUENCE:
        case BZ_AIBT_COMP_PARALLEL_SELECTOR:
        case BZ_AIBT_COMP_PARALLEL_SEQUENCE:
            //status = bzAIBTExecuteComposite(bt, node, state, nodeState);
            break;
        case BZ_AIBT_DECOR_DUMMY:
        case BZ_AIBT_DECOR_SUCCESS:
        case BZ_AIBT_DECOR_FAIL:
        case BZ_AIBT_DECOR_INVERT:
        case BZ_AIBT_DECOR_UNTIL_SUCCESS:
        case BZ_AIBT_DECOR_UNTIL_FAIL:
        case BZ_AIBT_DECOR_REPEAT:
        case BZ_AIBT_DECOR_DELAY:
            status = bzAIBTExecuteDecorator(node, dt, state, nodeState);
            break;
        case BZ_AIBT_ACTION:
            BZ_ASSERT(node->as.action.fn);
            return node->as.action.fn(state->userData);
    }
    return status;
}

BzAIBTStatus bzAIBTExecute(BzAIBTState *state, f32 dt) {
    BZ_ASSERT(state->nodeStatePool);
    BZ_ASSERT(bzObjectPoolGetObjectSize(state->nodeStatePool) == bzAIBTGetNodeStateSize());
    BZ_ASSERT(state);
    BZ_ASSERT(state->root);
    BzAIBTNodeState *first = state->first;
    const BzAIBTNode *firstNode = first ? first->node : state->root;
    BzAIBTStatus status = bzAIBTExecuteNode(firstNode, dt, state, first);
    return status;
}