#include "behaviour_tree.h"

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

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

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

    BzBTNodeType type;
    const char   *name;
    union {
        struct {
            i32 n;
        } repeat;
        struct {
            f32 ms;
        } delay;
        struct {
            BzBTActionFn fn;
        } action;
    } as;
};

struct BzBTNodeState {
    const BzBTNode *node;
    BzBTNodeState *next;
    BzBTNodeState *prev;

    union {
        struct {
            BzBTNode *running;
        } composite;
        struct {
            i32 iter;
        } repeat;
        struct {
            f32 elapsed;
        } delay;
    } as;
};


size_t bzBTGetNodeSize() {
    return sizeof(BzBTNode);
}
size_t  bzBTGetNodeStateSize() {
    return sizeof(BzBTNodeState);
}

const char *bzBTNodeTypeToStr(BzBTNodeType type) {
    switch (type) {
        case BZ_BT_COMP_SELECTOR:
            return "SELECTOR";
        case BZ_BT_COMP_SEQUENCE:
            return "SEQUENCE";
        case BZ_BT_COMP_PARALLEL_SELECTOR:
            return "PARALLEL_SELECTOR";
        case BZ_BT_COMP_PARALLEL_SEQUENCE:
            return "PARALLEL_SEQUENCE";
        case BZ_BT_DECOR_DUMMY:
            return "DUMMY";
        case BZ_BT_DECOR_SUCCESS:
            return "SUCCESS";
        case BZ_BT_DECOR_FAIL:
            return "FAIL";
        case BZ_BT_DECOR_INVERT:
            return "INVERT";
        case BZ_BT_DECOR_UNTIL_SUCCESS:
            return "UNTIL_SUCCESS";
        case BZ_BT_DECOR_UNTIL_FAIL:
            return "UNTIL_FAIL";
        case BZ_BT_DECOR_REPEAT:
            return "REPEAT";
        case BZ_BT_DECOR_DELAY:
            return "DELAY";
        case BZ_BT_ACTION:
            return "ACTION";
    }
}

static BzBTNode *bzBTNodeMake(BzObjectPool *nodePool, BzBTNode *parent, BzBTNodeType type) {
    BZ_ASSERT(nodePool);
    BZ_ASSERT(bzObjectPoolGetObjectSize(nodePool) == bzBTGetNodeSize());
    BzBTNode *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;
}

BzBTNode *bzBTMakeRoot(BzObjectPool *nodePool) {
    return bzBTNodeMake(nodePool, NULL, BZ_BT_DECOR_DUMMY);
}

void bzBTDestroyRoot(BzObjectPool *nodePool, BzBTNode *node) {
    BZ_ASSERT(node);
    BzBTNode *pNode = node;

    while (pNode) {
        BzBTNode *next = pNode->next;
        bzBTDestroyRoot(nodePool, pNode);
        pNode = next;
    }

    bzObjectPoolRelease(nodePool, node);
}

BzBTNode *bzBTCompSelector(BzObjectPool *nodePool, BzBTNode *parent, bool parallel) {
    BzBTNodeType type = parallel ?
                        BZ_BT_COMP_PARALLEL_SELECTOR :
                        BZ_BT_COMP_SELECTOR;
    return bzBTNodeMake(nodePool, parent, type);
}
BzBTNode *bzBTCompSequence(BzObjectPool *nodePool, BzBTNode *parent, bool parallel) {
    BzBTNodeType type = parallel ?
                        BZ_BT_COMP_PARALLEL_SEQUENCE :
                        BZ_BT_COMP_SEQUENCE;
    return bzBTNodeMake(nodePool, parent, type);
}

BzBTNode *bzBTDecorDummy(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_DECOR_DUMMY);
}
BzBTNode *bzBTDecorSuccess(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_DECOR_SUCCESS);
}
BzBTNode *bzBTDecorFail(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_DECOR_FAIL);
}
BzBTNode *bzBTDecorInvert(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_DECOR_INVERT);
}
BzBTNode *bzBTDecorUntilSuccess(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_DECOR_UNTIL_SUCCESS);
}
BzBTNode *bzBTDecorUntilFail(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_DECOR_UNTIL_FAIL);
}
BzBTNode *bzBTDecorRepeat(BzObjectPool *nodePool, BzBTNode *parent, i32 n) {
    BzBTNode *node = bzBTNodeMake(nodePool, parent, BZ_BT_DECOR_REPEAT);
    node->as.repeat.n = n;
    return node;
}
BzBTNode *bzBTDecorDelay(BzObjectPool *nodePool, BzBTNode *parent, f32 ms) {
    BzBTNode *node = bzBTNodeMake(nodePool, parent, BZ_BT_DECOR_DELAY);
    node->as.delay.ms = ms;
    return node;
}

BzBTNode *bzBTAction(BzObjectPool *nodePool, BzBTNode *parent, BzBTActionFn fn) {
    BzBTNode *node = bzBTNodeMake(nodePool, parent, BZ_BT_ACTION);
    node->as.action.fn = fn;
    return node;
}

void bzBTNodeSetName(BzBTNode *node, const char *name) {
    node->name = name;
}
const char *bzBTNodeGetName(const BzBTNode *node) {
    return node->name;
}

i32 bzBTDecorGetRepeat(const BzBTNode *node) {
    BZ_ASSERT(node->type == BZ_BT_DECOR_REPEAT);
    return node->as.repeat.n;
}
f32 bzBTDecorGetDelay(const BzBTNode *node) {
    BZ_ASSERT(node->type == BZ_BT_DECOR_DELAY);
    return node->as.delay.ms;
}

BzBTActionFn bzBTActionGetFn(const BzBTNode *node) {
    BZ_ASSERT(node->type == BZ_BT_ACTION);
    return node->as.action.fn;
}

BzBTNodeType bzBTGetNodeType(const BzBTNode *node) {
    return node->type;
}
BzBTNode *bzBTNodeChild(const BzBTNode *node) {
    return node->first;
}
BzBTNode *bzBTNodeNext(const BzBTNode *node) {
    return node->next;
}

const BzBTNodeState *bzBTNodeStateNext(const BzBTNodeState *state) {
    BZ_ASSERT(state);
    return state->next;
}
bool bzBTNodeMatchesState(const BzBTNode *node, const BzBTNodeState *state) {
    return state && state->node == node;
}

BzBTNode *bzBTCompStateGetRunningChild(const BzBTNodeState *state) {
    BZ_ASSERT(state->node);
    BzBTNodeType type = state->node->type;
    bool isComposite = type == BZ_BT_COMP_SELECTOR ||
                       type == BZ_BT_COMP_PARALLEL_SELECTOR ||
                       type == BZ_BT_COMP_SEQUENCE ||
                       type == BZ_BT_COMP_PARALLEL_SEQUENCE;
    BZ_ASSERT(isComposite);
    return state->as.composite.running;
}

i32 bzBTRepeatStateGetIter(const BzBTNodeState *state) {
    BZ_ASSERT(state->node && state->node->type == BZ_BT_DECOR_REPEAT);
    return state->as.repeat.iter;
}
f32 bzBTDelayStateGetElapsed(const BzBTNodeState *state) {
    BZ_ASSERT(state->node && state->node->type == BZ_BT_DECOR_DELAY);
    return state->as.delay.elapsed;
}

BzBTState bzBTCreateState(const BzBTStateDesc *desc) {
    BZ_ASSERT(desc->pool);
    BZ_ASSERT(bzObjectPoolGetObjectSize(desc->pool) == bzBTGetNodeStateSize());
    BZ_ASSERT(desc->root);
    return (BzBTState) {
            .root = desc->root,
            ._first = NULL,
            ._last = NULL,
            .nodeStatePool = desc->pool,
            .userData = desc->userData
    };
}
void bzBTDestroyState(BzBTState *state) {
    BzBTNodeState *pNodeState = state->_first;
    while (pNodeState) {
        BzBTNodeState *next = pNodeState->next;
        bzObjectPoolRelease(state->nodeStatePool, pNodeState);
        pNodeState = next;
    }
    bzMemSet(state, 0, sizeof(*state));
}

void bzBTStateAppend(BzBTState *state, BzBTNodeState *nodeState) {
    nodeState->next = NULL;
    nodeState->prev = state->_last;
    if (state->_last)
        state->_last->next = nodeState;
    else
        state->_first = nodeState;
    state->_last = nodeState;
}
void bzBTStatePop(BzBTState *state, BzBTNodeState *nodeState) {
    if (state->_first == nodeState) state->_first = nodeState->next;
    if (state->_last == nodeState) state->_last = nodeState->prev;
    BzBTNodeState *next = nodeState->next;
    BzBTNodeState *prev = nodeState->prev;
    if (nodeState->prev)
        nodeState->prev->next = next;
    if (nodeState->next)
        nodeState->next->prev = prev;
    nodeState->next = NULL;
    nodeState->prev = NULL;
}
void bzBTStateRenew(BzBTState *oldState, BzBTState *newState, BzBTNodeState *nodeState) {
    // Pop nodeState and transfer it to the back
    bzBTStatePop(oldState, nodeState);
    bzBTStateAppend(newState, nodeState);
}
BzBTNodeState *bzBTStatePool(BzBTState *state, const BzBTNode *node) {
    BzBTNodeState *nodeState = bzObjectPool(state->nodeStatePool);
    nodeState->next = NULL;
    nodeState->prev = NULL;
    nodeState->node = node;
    return nodeState;
}
void bzBTStateRelease(BzBTState *state, BzBTNodeState *nodeState) {
    bzBTStatePop(state, nodeState);
    bzObjectPoolRelease(state->nodeStatePool, nodeState);
}

bool nodeMatchesState(const BzBTNode *node, const BzBTNodeState *state) {
    return bzBTNodeMatchesState(node, state);
}
BzBTNodeState *getNextNodeState(const BzBTNode *node, BzBTNodeState *nodeState) {
    if (nodeState && nodeMatchesState(node, nodeState))
        return nodeState->next;
    return nodeState;
}

static inline BzBTStatus bzBTExecuteNode(const BzBTNode *node, f32 dt,
                                         BzBTNodeState *nodeState,
                                         BzBTState *oldState, BzBTState *newState);
static inline BzBTStatus bzBTExecuteComposite(const BzBTNode *node, f32 dt,
                                              BzBTNodeState *nodeState,
                                              BzBTState *oldState, BzBTState *newState) {
    BzBTNodeState *nextState = getNextNodeState(node, nodeState);
    BzBTNode *start = node->first;
    bool isParallel = node->type == BZ_BT_COMP_PARALLEL_SEQUENCE ||
                      node->type == BZ_BT_COMP_PARALLEL_SELECTOR;

    if (!isParallel && nodeMatchesState(node, nodeState))
        start = nodeState->as.composite.running;

    // Always push dummy state
    if (nodeMatchesState(node, nodeState)) {
        bzBTStateRenew(oldState, newState, nodeState);
    } else {
        nodeState = bzBTStatePool(oldState, node);
        bzBTStateAppend(newState, nodeState);
    }
    i32 numRunning = 0;
    i32 numSuccessful = 0;
    i32 numFailed = 0;
    i32 numChildren = 0;
    BzBTStatus status = BZ_BT_ERROR;
    BzBTNode *child = start;
    for (;child; child = child->next) {
        BzBTStatus childStatus = bzBTExecuteNode(child, dt, nextState, oldState, newState);
        numChildren++;
        switch (childStatus) {
            case BZ_BT_RUNNING:
                numRunning++;
                break;
            case BZ_BT_SUCCESS:
                numSuccessful++;
                break;
            case BZ_BT_FAIL:
                numFailed++;
                break;
            default:
                break;
        }
        switch (node->type) {
            case BZ_BT_COMP_SELECTOR:
            case BZ_BT_COMP_PARALLEL_SELECTOR:
                if (childStatus == BZ_BT_SUCCESS)
                    status = BZ_BT_SUCCESS;
                break;
            case BZ_BT_COMP_SEQUENCE:
            case BZ_BT_COMP_PARALLEL_SEQUENCE:
                if (childStatus == BZ_BT_FAIL)
                    status = BZ_BT_FAIL;
                break;
            default:
                break;
        }
        if (status == BZ_BT_FAIL || status == BZ_BT_SUCCESS)
            break;
        if (numRunning > 0 && !isParallel) {
            status = BZ_BT_RUNNING;
            break;
        }
    }
    switch (node->type) {
        case BZ_BT_COMP_SELECTOR:
        case BZ_BT_COMP_PARALLEL_SELECTOR:
            if (numFailed == numChildren)
                status = BZ_BT_FAIL;
            break;
        case BZ_BT_COMP_SEQUENCE:
        case BZ_BT_COMP_PARALLEL_SEQUENCE:
            if (numSuccessful == numChildren)
                status = BZ_BT_SUCCESS;
            break;
        default:
            break;
    }

    if (status == BZ_BT_ERROR) {
        bzBTStateRelease(newState, nodeState);
        return BZ_BT_ERROR;
    }

    bool finished = status == BZ_BT_SUCCESS ||
                    status == BZ_BT_FAIL;
    if (finished) {
        // Dummy state is no longer needed
        bzBTStateRelease(newState, nodeState);
    } else {
        BZ_ASSERT(status == BZ_BT_RUNNING);
        nodeState->as.composite.running = child;
    }
    return status;
}
static inline BzBTStatus bzBTExecuteDecorator(const BzBTNode *node, f32 dt,
                                              BzBTNodeState *nodeState,
                                              BzBTState *oldState, BzBTState *newState) {
    // Ensure decorator has only one child, if any
    BZ_ASSERT(!node->first || node->first == node->last);
    BzBTNodeState *nextState = getNextNodeState(node, nodeState);

    switch (node->type) {
        case BZ_BT_DECOR_REPEAT:
            if (!nodeMatchesState(node, nodeState)) {
                BzBTNodeState *newNodeState = bzBTStatePool(oldState, node);
                newNodeState->as.repeat.iter = 0;
                bzBTStateAppend(newState, newNodeState);
                nodeState = newNodeState;
            } else {
                bzBTStateRenew(oldState, newState, nodeState);
            }
            break;
        case BZ_BT_DECOR_DELAY:
            if (!nodeMatchesState(node, nodeState)) {
                BzBTNodeState *newNodeState = bzBTStatePool(oldState, node);
                newNodeState->as.delay.elapsed = dt;
                bzBTStateAppend(newState, newNodeState);
                return BZ_BT_RUNNING;
            }
            nodeState->as.delay.elapsed += dt;
            if (nodeState->as.delay.elapsed < node->as.delay.ms) {
                bzBTStateRenew(oldState, newState, nodeState);
                return BZ_BT_RUNNING;
            }
            break;
        default:
            break;
    }

    // Implicit success, if no children are present
    BzBTStatus inStatus = BZ_BT_SUCCESS;
    if (node->first)
        inStatus = bzBTExecuteNode(node->first, dt, nextState, oldState, newState);

    // Propagate ERROR, RUNNING up
    if (inStatus == BZ_BT_ERROR)
        return BZ_BT_ERROR;
    if (inStatus == BZ_BT_RUNNING)
        return BZ_BT_RUNNING;

    BzBTStatus status = BZ_BT_ERROR;
    switch (node->type) {
        case BZ_BT_DECOR_DUMMY:
        case BZ_BT_DECOR_DELAY: // Delay already handled
            status = inStatus;
            break;
        case BZ_BT_DECOR_SUCCESS:
            status = BZ_BT_SUCCESS;
            break;
        case BZ_BT_DECOR_FAIL:
            status = BZ_BT_FAIL;
            break;
        case BZ_BT_DECOR_INVERT:
            if (inStatus == BZ_BT_FAIL)
                status = BZ_BT_SUCCESS;
            if (inStatus == BZ_BT_SUCCESS)
                status = BZ_BT_FAIL;
            break;
        case BZ_BT_DECOR_UNTIL_SUCCESS:
            if (inStatus == BZ_BT_SUCCESS)
                status = BZ_BT_SUCCESS;
            else
                status = BZ_BT_RUNNING;
            break;
        case BZ_BT_DECOR_UNTIL_FAIL:
            if (inStatus == BZ_BT_FAIL)
                status = BZ_BT_SUCCESS;
            else
                status = BZ_BT_RUNNING;
            break;
        case BZ_BT_DECOR_REPEAT:
            BZ_ASSERT(nodeState->node == node);
            nodeState->as.repeat.iter++;
            if (nodeState->as.repeat.iter >= node->as.repeat.n) {
                bzBTStateRelease(newState, nodeState);
                status = inStatus;
                break;
            }
            status = BZ_BT_RUNNING;
            break;
        default:
            break;
    }
    return status;
}
static inline BzBTStatus bzBTExecuteNode(const BzBTNode *node, f32 dt,
                                         BzBTNodeState *nodeState,
                                         BzBTState *oldState, BzBTState *newState) {
    BzBTStatus status = BZ_BT_ERROR;
    switch (node->type) {
        case BZ_BT_COMP_SELECTOR:
        case BZ_BT_COMP_SEQUENCE:
        case BZ_BT_COMP_PARALLEL_SELECTOR:
        case BZ_BT_COMP_PARALLEL_SEQUENCE:
            status = bzBTExecuteComposite(node, dt, nodeState, oldState, newState);
            break;
        case BZ_BT_DECOR_DUMMY:
        case BZ_BT_DECOR_SUCCESS:
        case BZ_BT_DECOR_FAIL:
        case BZ_BT_DECOR_INVERT:
        case BZ_BT_DECOR_UNTIL_SUCCESS:
        case BZ_BT_DECOR_UNTIL_FAIL:
        case BZ_BT_DECOR_REPEAT:
        case BZ_BT_DECOR_DELAY:
            status = bzBTExecuteDecorator(node, dt, nodeState, oldState, newState);
            break;
        case BZ_BT_ACTION:
            BZ_ASSERT(node->as.action.fn);
            return node->as.action.fn(oldState->userData, dt);
    }
    return status;
}

BzBTStatus bzBTExecute(BzBTState *state, f32 dt) {
    BZ_ASSERT(state->nodeStatePool);
    BZ_ASSERT(bzObjectPoolGetObjectSize(state->nodeStatePool) == bzBTGetNodeStateSize());
    BZ_ASSERT(state);
    if (state->root == NULL) {
        return BZ_BT_FAIL;
    }

    BzBTState newState = {
            ._first = NULL,
            ._last = NULL,
            .nodeStatePool = state->nodeStatePool
    };

    BzBTNodeState *first = state->_first;
    const BzBTNode *firstNode = first ? first->node : state->root;
    BzBTStatus status = bzBTExecuteNode(firstNode, dt, first, state, &newState);

    // Release leftover states
    BzBTNodeState *pState = state->_first;
    while (pState) {
        BzBTNodeState *next = pState->next;
        bzBTStateRelease(state, pState);
        pState = next;
    }
    state->_first = newState._first;
    state->_last = newState._last;

    switch (status) {
        case BZ_BT_SUCCESS:
            state->onSuccess && state->onSuccess(state->userData, dt);
            break;
        case BZ_BT_FAIL:
            state->onFailure && state->onFailure(state->userData, dt);
            break;
        case BZ_BT_ERROR:
            state->onError && state->onError(state->userData, dt);
            break;
        default:
            break;
    }

    return status;
}