#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_PSELECTOR:
            return "P_SELECTOR";
        case BZ_BT_COMP_PSEQUENCE:
            return "P_SEQUENCE";
        case BZ_BT_DECOR_DUMMY:
            return "DUMMY";
        case BZ_BT_DECOR_SUCCESS:
            return "RETURN_SUCCESS";
        case BZ_BT_DECOR_FAIL:
            return "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) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_COMP_SELECTOR);
}
BzBTNode *bzBTCompSequence(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_COMP_SEQUENCE);
}
BzBTNode *bzBTCompPSelector(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_COMP_PSELECTOR);
}
BzBTNode *bzBTCompPSequence(BzObjectPool *nodePool, BzBTNode *parent) {
    return bzBTNodeMake(nodePool, parent, BZ_BT_COMP_PSEQUENCE);
}

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 BzBTNode *bzBTNodeStateGetNode(const BzBTNodeState *state) {
    BZ_ASSERT(state);
    return state->node;
}

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_PSELECTOR ||
                       type == BZ_BT_COMP_SEQUENCE ||
                       type == BZ_BT_COMP_PSEQUENCE;
    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));
}

typedef struct BzBTExecState {
    BzBTNodeState *first;
    BzBTNodeState *last;
    void *userData;
} BzBTExecState;
typedef struct BzBTExecReturn {
    BzBTStatus    status;
    BzBTExecState state;
} BzBTExecReturn;

BzBTNodeState *execStatePopFront(BzBTExecState *state) {
    BZ_ASSERT(state->first);
    BzBTNodeState *popped = state->first;

    if (state->first)
        state->first = state->first->next;
    if (state->first == NULL)
        state->last = NULL;

    popped->prev = NULL;
    popped->next = NULL;
    return popped;
}
void execStatePushFront(BzBTExecState *state, BzBTNodeState *nodeState) {
    BZ_ASSERT(state);

    nodeState->prev = NULL;
    nodeState->next = state->first;

    if (state->first) {
        state->first->prev = nodeState;
    } else {
        state->first = nodeState;
        state->last = nodeState;
    }
    state->first = nodeState;
}
void execStatePushBack(BzBTExecState *state, BzBTNodeState *nodeState) {
    BZ_ASSERT(state);

    nodeState->next = NULL;
    nodeState->prev = state->last;

    if (state->last) {
        state->last->next = nodeState;
    } else {
        state->first = nodeState;
        state->last = nodeState;
    }
    state->last = nodeState;
}
void execStateClear(BzBTExecState *state, BzObjectPool *pool) {
    BzBTNodeState *pState = state->first;
    while (pState) {
        BzBTNodeState *next = pState->next;
        bzObjectPoolRelease(pool, pState);
        pState = next;
    }
    state->first = NULL;
    state->last = NULL;
}
void execStateMerge(BzBTExecState *state, BzBTExecState *toMerge) {
    BzBTNodeState *pState = toMerge->first;
    while (pState) {
        BzBTNodeState *next = pState->next;
        execStatePushBack(state, pState);
        pState = next;
    }

    toMerge->first = NULL;
    toMerge->last = NULL;
}


/**
 *
 * @param node
 * @param state
 * @return
 */
bool nodeMatchesState(const BzBTNode *node, const BzBTExecState *state) {
    return state->first && state->first->node == node;
}
/**
 * Return BzBTNodeState for given node, if node matches first state node.
 */
static BzBTNodeState *getNodeState(const BzBTNode *node, BzBTExecState *state) {
    if (nodeMatchesState(node, state)) {
        return execStatePopFront(state);
    }
    return NULL;
}
static BzBTNodeState *ensureValidNodeState(const BzBTNode *node, BzBTNodeState *existing,
                                           BzObjectPool *pool) {
    if (existing)
        return existing;

    existing = bzObjectPool(pool);
    existing->node = node;
    existing->prev = NULL;
    existing->next = NULL;
    bzMemSet(&existing->as, 0, sizeof(existing->as));
    return existing;
}
static void releaseNodeState(BzBTNodeState *nodeState, BzObjectPool *pool) {
    bzObjectPoolRelease(pool, nodeState);
}

static inline void bzBTClearNodeState(const BzBTNode *node, BzBTExecState *state,
                                      BzObjectPool *pool) {
    BzBTNodeState *nodeState = getNodeState(node, state);
    if (nodeState == NULL) return;
    releaseNodeState(nodeState, pool);

    const BzBTNode *first = node->first;
    while (first) {
        bzBTClearNodeState(first, state, pool);
        first = first->next;
    }
}
static inline BzBTExecReturn bzBTExecuteNode(const BzBTNode *node, f32 dt,
                                             BzBTExecState *state, BzObjectPool *statePool);
static inline BzBTExecReturn bzBTExecuteComposite(const BzBTNode *node, f32 dt,
                                                  BzBTExecState *state, BzObjectPool *statePool) {
    BzBTNodeState *nodeState = getNodeState(node, state);

    bool isParallel = node->type == BZ_BT_COMP_PSEQUENCE ||
                      node->type == BZ_BT_COMP_PSELECTOR;

    BzBTNode *start = node->first;
    if (!isParallel && nodeState)
        start = nodeState->as.composite.running;

    i32 numRunning = 0;
    i32 numSuccessful = 0;
    i32 numFailed = 0;
    i32 numChildren = 0;
    BzBTExecReturn execReturn = { .status = BZ_BT_ERROR };
    BzBTNode *child = start;
    for (;child; child = child->next) {
        BzBTExecReturn childReturn = bzBTExecuteNode(child, dt, state, statePool);
        if (childReturn.status == BZ_BT_RUNNING) {
            execStateMerge(&execReturn.state, &childReturn.state);
        }
        BZ_ASSERT(childReturn.state.first == NULL && childReturn.state.last == NULL);
        numChildren++;
        switch (childReturn.status) {
            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_PSELECTOR:
                if (childReturn.status == BZ_BT_SUCCESS)
                    execReturn.status = BZ_BT_SUCCESS;
                break;
            case BZ_BT_COMP_SEQUENCE:
            case BZ_BT_COMP_PSEQUENCE:
                if (childReturn.status == BZ_BT_FAIL)
                    execReturn.status = BZ_BT_FAIL;
                break;
            default:
                break;
        }
        if (execReturn.status == BZ_BT_FAIL || execReturn.status == BZ_BT_SUCCESS)
            break;
        if (numRunning > 0) {
            execReturn.status = BZ_BT_RUNNING;
            if (!isParallel)
                break;
        }
    }
    switch (node->type) {
        case BZ_BT_COMP_SELECTOR:
        case BZ_BT_COMP_PSELECTOR:
            if (numFailed == numChildren)
                execReturn.status = BZ_BT_FAIL;
            break;
        case BZ_BT_COMP_SEQUENCE:
        case BZ_BT_COMP_PSEQUENCE:
            if (numSuccessful == numChildren)
                execReturn.status = BZ_BT_SUCCESS;
            break;
        default:
            break;
    }

    // Propagate ERROR up
    if (execReturn.status == BZ_BT_ERROR) {
        execStateClear(&execReturn.state, statePool);
        execReturn.status = BZ_BT_ERROR;
        return execReturn;
    }

    bool finished = execReturn.status == BZ_BT_SUCCESS ||
                    execReturn.status == BZ_BT_FAIL;

    // Clean up, if exited prematurely
    const BzBTNode *pChild = child;
    while (finished && pChild) {
        bzBTClearNodeState(pChild, state, statePool);
        pChild = pChild->next;
    }

    if (finished) {
        BZ_ASSERT(execReturn.state.first == NULL);
        if (nodeState) releaseNodeState(nodeState, statePool);
        return execReturn;
    }
    BZ_ASSERT(execReturn.status == BZ_BT_RUNNING);
    // Parallels don't utilize state but still require it for correct tree traversal.
    nodeState = ensureValidNodeState(node, nodeState, statePool);
    nodeState->as.composite.running = child;
    execStatePushFront(&execReturn.state, nodeState);

    return execReturn;
}
static inline BzBTExecReturn bzBTExecuteDecorator(const BzBTNode *node, f32 dt,
                                                  BzBTExecState *state, BzObjectPool *statePool) {
    // Ensure decorator has only one child, if any
    BZ_ASSERT(!node->first || node->first == node->last);

    BzBTNodeState *nodeState = getNodeState(node, state);

    BzBTExecReturn execReturn = { .status = BZ_BT_ERROR };

    switch (node->type) {
        case BZ_BT_DECOR_REPEAT:
            nodeState = ensureValidNodeState(node, nodeState, statePool);
            execStatePushBack(&execReturn.state, nodeState);
            break;
        case BZ_BT_DECOR_DELAY:
            nodeState = ensureValidNodeState(node, nodeState, statePool);
            nodeState->as.delay.elapsed += dt;
            execStatePushBack(&execReturn.state, nodeState);
            if (nodeState->as.delay.elapsed < node->as.delay.ms) {
                execReturn.status = BZ_BT_RUNNING;
                return execReturn;
            }
            BZ_ASSERT(nodeState == execStatePopFront(&execReturn.state));
            releaseNodeState(nodeState, statePool);
            break;
        default:
            break;
    }

    // Implicit success, if no children are present
    BzBTExecReturn childReturn = { .status = BZ_BT_SUCCESS };
    if (node->first)
        childReturn = bzBTExecuteNode(node->first, dt, state, statePool);

    // Propagate ERROR, RUNNING up
    if (childReturn.status == BZ_BT_ERROR) {
        execStateClear(&childReturn.state, statePool);
        execStateClear(&execReturn.state, statePool);
        execReturn.status = BZ_BT_ERROR;
        return execReturn;

    }
    if (childReturn.status == BZ_BT_RUNNING) {
        execStateMerge(&execReturn.state, &childReturn.state);
        execReturn.status = BZ_BT_RUNNING;
        return execReturn;
    }
    BZ_ASSERT(childReturn.state.first == NULL);

    switch (node->type) {
        case BZ_BT_DECOR_DUMMY:
        case BZ_BT_DECOR_DELAY: // Delay already handled
            execReturn.status = childReturn.status;
            break;
        case BZ_BT_DECOR_SUCCESS:
            execReturn.status = BZ_BT_SUCCESS;
            break;
        case BZ_BT_DECOR_FAIL:
            execReturn.status = BZ_BT_FAIL;
            break;
        case BZ_BT_DECOR_INVERT:
            if (childReturn.status == BZ_BT_FAIL)
                execReturn.status = BZ_BT_SUCCESS;
            if (childReturn.status == BZ_BT_SUCCESS)
                execReturn.status = BZ_BT_FAIL;
            break;
        case BZ_BT_DECOR_UNTIL_SUCCESS:
            if (childReturn.status == BZ_BT_SUCCESS)
                execReturn.status = BZ_BT_SUCCESS;
            else
                execReturn.status = BZ_BT_RUNNING;
            break;
        case BZ_BT_DECOR_UNTIL_FAIL:
            if (childReturn.status == BZ_BT_FAIL)
                execReturn.status = BZ_BT_SUCCESS;
            else
                execReturn.status = BZ_BT_RUNNING;
            break;
        case BZ_BT_DECOR_REPEAT:
            BZ_ASSERT(nodeMatchesState(node, &execReturn.state));
            BZ_ASSERT(nodeState);
            nodeState->as.repeat.iter++;
            if (nodeState->as.repeat.iter >= node->as.repeat.n) {
                BZ_ASSERT(nodeState == execStatePopFront(&execReturn.state));
                releaseNodeState(nodeState, statePool);
                execReturn.status = childReturn.status;
                break;
            }
            execReturn.status = BZ_BT_RUNNING;
            break;
        default:
            break;
    }
    return execReturn;
}
static inline BzBTExecReturn bzBTExecuteNode(const BzBTNode *node, f32 dt,
                                             BzBTExecState *state, BzObjectPool *statePool) {
    BzBTExecReturn execReturn = { .status = BZ_BT_ERROR };
    switch (node->type) {
        case BZ_BT_COMP_SELECTOR:
        case BZ_BT_COMP_SEQUENCE:
        case BZ_BT_COMP_PSELECTOR:
        case BZ_BT_COMP_PSEQUENCE:
            execReturn = bzBTExecuteComposite(node, dt, state, statePool);
            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:
            execReturn = bzBTExecuteDecorator(node, dt, state, statePool);
            break;
        case BZ_BT_ACTION:
            BZ_ASSERT(node->as.action.fn);
            execReturn.status = node->as.action.fn(state->userData, dt);
            break;
    }
    return execReturn;
}

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;
    }

    BzBTNodeState *first = state->_first;
    const BzBTNode *firstNode = first ? first->node : state->root;
    BzBTExecState execState = {
            .first = state->_first,
            .last = state->_last,
            .userData = state->userData
    };
    BzBTExecReturn execReturn = bzBTExecuteNode(firstNode, dt, &execState, state->nodeStatePool);
    if (execReturn.status == BZ_BT_ERROR)
        execStateClear(&execReturn.state, state->nodeStatePool);
    BZ_ASSERT(execState.first == NULL && execState.last == NULL);

    state->_first = execReturn.state.first;
    state->_last = execReturn.state.last;

    switch (execReturn.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 execReturn.status;
}