#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;
            const char *name;
        } action;
    } as;

};

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

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


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

const char *bzAIBTNodeTypeToStr(BzAIBTNodeType type) {
    switch (type) {
        case BZ_AIBT_COMP_SELECTOR:
            return "SELECTOR";
        case BZ_AIBT_COMP_SEQUENCE:
            return "SEQUENCE";
        case BZ_AIBT_COMP_PARALLEL_SELECTOR:
            return "PARALLEL_SELECTOR";
        case BZ_AIBT_COMP_PARALLEL_SEQUENCE:
            return "PARALLEL_SEQUENCE";
        case BZ_AIBT_DECOR_DUMMY:
            return "DUMMY";
        case BZ_AIBT_DECOR_SUCCESS:
            return "SUCCESS";
        case BZ_AIBT_DECOR_FAIL:
            return "FAIL";
        case BZ_AIBT_DECOR_INVERT:
            return "INVERT";
        case BZ_AIBT_DECOR_UNTIL_SUCCESS:
            return "UNTIL_SUCCESS";
        case BZ_AIBT_DECOR_UNTIL_FAIL:
            return "UNTIL_FAIL";
        case BZ_AIBT_DECOR_REPEAT:
            return "REPEAT";
        case BZ_AIBT_DECOR_DELAY:
            return "DELAY";
        case BZ_AIBT_ACTION:
            return "ACTION";
    }
}

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,
                         const char *name) {
    BzAIBTNode *node = bzAIBTNodeMake(nodePool, parent, BZ_AIBT_ACTION);
    node->as.action.fn = fn;
    node->as.action.name = name;
    return node;
}

i32 bzAIBTDecorGetRepeat(const BzAIBTNode *node) {
    BZ_ASSERT(node->type == BZ_AIBT_DECOR_REPEAT);
    return node->as.repeat.n;
}
f32 bzAIBTDecorGetDelay(const BzAIBTNode *node) {
    BZ_ASSERT(node->type == BZ_AIBT_DECOR_DELAY);
    return node->as.delay.ms;
}

BzAIBTActionFn bzAIBTActionGetFn(const BzAIBTNode *node) {
    BZ_ASSERT(node->type == BZ_AIBT_ACTION);
    return node->as.action.fn;
}
const char *bzAIBTActionGetName(const BzAIBTNode *node) {
    BZ_ASSERT(node->type == BZ_AIBT_ACTION);
    return node->as.action.name;
}

BzAIBTNodeType bzAIBTGetNodeType(const BzAIBTNode *node) {
    return node->type;
}
BzAIBTNode *bzAIBTNodeChild(const BzAIBTNode *node) {
    return node->first;
}
BzAIBTNode *bzAIBTNodeNext(const BzAIBTNode *node) {
    return node->next;
}

const BzAIBTNodeState *bzAIBTNodeStateNext(const BzAIBTNodeState *state) {
    BZ_ASSERT(state);
    return state->next;
}
bool bzAIBTNodeMatchesState(const BzAIBTNode *node, const BzAIBTNodeState *state) {
    return state && state->node == node;
}

i32 bzAIBTRepeatStateGetIter(const BzAIBTNodeState *state) {
    BZ_ASSERT(state->node && state->node->type == BZ_AIBT_DECOR_REPEAT);
    return state->as.repeat.iter;
}
f32 bzAIBTDelayStateGetElapsed(const BzAIBTNodeState *state) {
    BZ_ASSERT(state->node && state->node->type == BZ_AIBT_DECOR_DELAY);
    return state->as.delay.elapsed;
}

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 bzAIBTStateAppend(BzAIBTState *state, BzAIBTNodeState *nodeState) {
    nodeState->next = NULL;
    nodeState->prev = state->last;
    if (state->last)
        state->last->next = nodeState;
    else
        state->first = nodeState;
    state->last = nodeState;
}
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;
    nodeState->next = NULL;
    nodeState->prev = NULL;
}
void bzAIBTStateRenew(BzAIBTState *oldState, BzAIBTState *newState, BzAIBTNodeState *nodeState) {
    // Pop nodeState and transfer it to the back
    bzAIBTStatePop(oldState, nodeState);
    bzAIBTStateAppend(newState, nodeState);
}
BzAIBTNodeState *bzAIBTStatePool(BzAIBTState *state, const BzAIBTNode *node) {
    BzAIBTNodeState *nodeState = bzObjectPool(state->nodeStatePool);
    nodeState->next = NULL;
    nodeState->prev = NULL;
    nodeState->node = node;
    return nodeState;
}
void bzAIBTStateRelease(BzAIBTState *state, BzAIBTNodeState *nodeState) {
    bzObjectPoolRelease(state->nodeStatePool, nodeState);
}

bool nodeMatchesState(const BzAIBTNode *node, const BzAIBTNodeState *state) {
    return bzAIBTNodeMatchesState(node, state);
}
BzAIBTNodeState *getNextNodeState(const BzAIBTNode *node, BzAIBTNodeState *nodeState) {
    if (nodeState && nodeMatchesState(node, nodeState))
        return nodeState->next;
    return nodeState;
}

static inline BzAIBTStatus bzAIBTExecuteNode(const BzAIBTNode *node, f32 dt,
                                             BzAIBTNodeState *nodeState,
                                             BzAIBTState *oldState, BzAIBTState *newState);
static inline BzAIBTStatus bzAIBTExecuteComposite(const BzAIBTNode *node, f32 dt,
                                                  BzAIBTNodeState *nodeState,
                                                  BzAIBTState *oldState, BzAIBTState *newState) {
    BzAIBTNodeState *nextState = getNextNodeState(node, nodeState);
    BzAIBTNode *start = node->first;
    bool isParallel = node->type == BZ_AIBT_COMP_PARALLEL_SEQUENCE ||
            node->type == BZ_AIBT_COMP_PARALLEL_SELECTOR;

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

    // Always push dummy state
    if (nodeMatchesState(node, nodeState)) {
        bzAIBTStateRenew(oldState, newState, nodeState);
    } else {
        nodeState = bzAIBTStatePool(oldState, node);
        bzAIBTStateAppend(newState, nodeState);
    }
    i32 numRunning = 0;
    i32 numSuccessful = 0;
    i32 numFailed = 0;
    i32 numChildren = 0;
    BzAIBTStatus status = BZ_AIBT_ERROR;
    BzAIBTNode *child = start;
    for (;child; child = child->next) {
        BzAIBTStatus childStatus = bzAIBTExecuteNode(child, dt, nextState, oldState, newState);
        numChildren++;
        switch (childStatus) {
            case BZ_AIBT_RUNNING:
                numRunning++;
                break;
            case BZ_AIBT_SUCCESS:
                numSuccessful++;
                break;
            case BZ_AIBT_FAIL:
                numFailed++;
                break;
            default:
                break;
        }
        switch (node->type) {
            case BZ_AIBT_COMP_SELECTOR:
            case BZ_AIBT_COMP_PARALLEL_SELECTOR:
                if (childStatus == BZ_AIBT_SUCCESS)
                    status = BZ_AIBT_SUCCESS;
                break;
            case BZ_AIBT_COMP_SEQUENCE:
            case BZ_AIBT_COMP_PARALLEL_SEQUENCE:
                if (childStatus == BZ_AIBT_FAIL)
                    status = BZ_AIBT_FAIL;
                break;
            default:
                break;
        }
        if (status == BZ_AIBT_FAIL || status == BZ_AIBT_SUCCESS)
            break;
        if (numRunning > 0 && !isParallel) {
            status = BZ_AIBT_RUNNING;
            break;
        }
    }
    switch (node->type) {
        case BZ_AIBT_COMP_SELECTOR:
        case BZ_AIBT_COMP_PARALLEL_SELECTOR:
            if (numFailed == numChildren)
                status = BZ_AIBT_FAIL;
            break;
        case BZ_AIBT_COMP_SEQUENCE:
        case BZ_AIBT_COMP_PARALLEL_SEQUENCE:
            if (numSuccessful == numChildren)
                status = BZ_AIBT_SUCCESS;
            break;
        default:
            break;
    }

    if (status == BZ_AIBT_ERROR) {
        bzAIBTStatePop(newState, nodeState);
        return BZ_AIBT_ERROR;
    }

    bool finished = status == BZ_AIBT_SUCCESS ||
            status == BZ_AIBT_FAIL;
    if (finished) {
        // Dummy state is no longer needed
        bzAIBTStatePop(newState, nodeState);
    } else {
        BZ_ASSERT(status == BZ_AIBT_RUNNING);
        nodeState->as.composite.running = child;
    }
    return status;
}
static inline BzAIBTStatus bzAIBTExecuteDecorator(const BzAIBTNode *node, f32 dt,
                                                  BzAIBTNodeState *nodeState,
                                                  BzAIBTState *oldState, BzAIBTState *newState) {
    // Ensure decorator has only one child, if any
    BZ_ASSERT(!node->first || node->first == node->last);
    BzAIBTNodeState *nextState = getNextNodeState(node, nodeState);

    switch (node->type) {
        case BZ_AIBT_DECOR_REPEAT:
            if (!nodeMatchesState(node, nodeState)) {
                BzAIBTNodeState *newNodeState = bzAIBTStatePool(oldState, node);
                newNodeState->as.repeat.iter = 0;
                bzAIBTStateAppend(newState, newNodeState);
                nodeState = newNodeState;
            } else {
                bzAIBTStateRenew(oldState, newState, nodeState);
            }
            break;
        case BZ_AIBT_DECOR_DELAY:
            if (!nodeMatchesState(node, nodeState)) {
                BzAIBTNodeState *newNodeState = bzAIBTStatePool(oldState, node);
                newNodeState->as.delay.elapsed = dt;
                bzAIBTStateAppend(newState, newNodeState);
                return BZ_AIBT_RUNNING;
            }
            nodeState->as.delay.elapsed += dt;
            if (nodeState->as.delay.elapsed < node->as.delay.ms) {
                bzAIBTStateRenew(oldState, newState, nodeState);
                return BZ_AIBT_RUNNING;
            }
            break;
        default:
            break;
    }

    // Implicit success, if no children are present
    BzAIBTStatus inStatus = BZ_AIBT_SUCCESS;
    if (node->first)
        inStatus = bzAIBTExecuteNode(node->first, dt, nextState, oldState, newState);

    // Propagate ERROR, RUNNING 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(newState, nodeState);
                status = inStatus;
                break;
            }
            status = BZ_AIBT_RUNNING;
            break;
        default:
            break;
    }
    return status;
}
static inline BzAIBTStatus bzAIBTExecuteNode(const BzAIBTNode *node, f32 dt,
                                             BzAIBTNodeState *nodeState,
                                             BzAIBTState *oldState, BzAIBTState *newState) {
    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(node, dt, nodeState, oldState, newState);
            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, nodeState, oldState, newState);
            break;
        case BZ_AIBT_ACTION:
            BZ_ASSERT(node->as.action.fn);
            return node->as.action.fn(oldState->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);

    BzAIBTState newState = {
            .first = NULL,
            .last = NULL,
    };

    BzAIBTNodeState *first = state->first;
    const BzAIBTNode *firstNode = first ? first->node : state->root;
    BzAIBTStatus status = bzAIBTExecuteNode(firstNode, dt, first, state, &newState);

    // Release leftover states
    BzAIBTNodeState *pState = state->first;
    while (pState) {
        BzAIBTNodeState *next = pState->next;
        bzAIBTStateRelease(state, pState);
        pState = next;
    }
    state->first = newState.first;
    state->last = newState.last;
    return status;
}