#include "xtlite/xtlite_engine.h"

#include <stdlib.h>
#include <string.h>

typedef struct {
    size_t source;
    size_t destination;
} xt_edge;

typedef struct {
    float gain;
} xt_gain_node;

struct xt_graph {
    size_t max_nodes;
    size_t node_count;
    size_t edge_count;
    xt_gain_node *nodes;
    xt_edge *edges;
};

xt_result xt_graph_create(size_t max_nodes, xt_graph **out_graph)
{
    xt_graph *graph;

    if (out_graph == NULL || max_nodes == 0U) {
        return XT_ERR_INVALID_ARGUMENT;
    }
    *out_graph = NULL;

    graph = (xt_graph *)calloc(1U, sizeof(*graph));
    if (graph == NULL) {
        return XT_ERR_OUT_OF_MEMORY;
    }

    graph->nodes = (xt_gain_node *)calloc(max_nodes, sizeof(*graph->nodes));
    graph->edges = (xt_edge *)calloc(max_nodes * max_nodes, sizeof(*graph->edges));
    if (graph->nodes == NULL || graph->edges == NULL) {
        free(graph->nodes);
        free(graph->edges);
        free(graph);
        return XT_ERR_OUT_OF_MEMORY;
    }

    graph->max_nodes = max_nodes;
    *out_graph = graph;
    return XT_OK;
}

void xt_graph_destroy(xt_graph *graph)
{
    if (graph != NULL) {
        free(graph->nodes);
        free(graph->edges);
        free(graph);
    }
}

xt_result xt_graph_add_gain(xt_graph *graph, float gain, size_t *out_node_id)
{
    if (graph == NULL || out_node_id == NULL) {
        return XT_ERR_INVALID_ARGUMENT;
    }
    if (graph->node_count >= graph->max_nodes) {
        return XT_ERR_CAPACITY;
    }

    graph->nodes[graph->node_count].gain = gain;
    *out_node_id = graph->node_count;
    graph->node_count++;
    return XT_OK;
}

xt_result xt_graph_connect(xt_graph *graph, size_t source_node,
                           size_t destination_node)
{
    size_t i;
    size_t edge_capacity;

    if (graph == NULL) {
        return XT_ERR_INVALID_ARGUMENT;
    }
    if (source_node >= graph->node_count || destination_node >= graph->node_count) {
        return XT_ERR_NOT_FOUND;
    }

    edge_capacity = graph->max_nodes * graph->max_nodes;
    if (graph->edge_count >= edge_capacity) {
        return XT_ERR_CAPACITY;
    }
    for (i = 0U; i < graph->edge_count; i++) {
        if (graph->edges[i].source == source_node &&
            graph->edges[i].destination == destination_node) {
            return XT_ERR_INVALID_STATE;
        }
    }

    graph->edges[graph->edge_count].source = source_node;
    graph->edges[graph->edge_count].destination = destination_node;
    graph->edge_count++;
    return XT_OK;
}

xt_result xt_graph_validate(const xt_graph *graph)
{
    size_t *indegree;
    size_t *queue;
    size_t head = 0U;
    size_t tail = 0U;
    size_t visited = 0U;
    size_t i;

    if (graph == NULL) {
        return XT_ERR_INVALID_ARGUMENT;
    }
    if (graph->node_count == 0U) {
        return XT_OK;
    }

    indegree = (size_t *)calloc(graph->node_count, sizeof(*indegree));
    queue = (size_t *)calloc(graph->node_count, sizeof(*queue));
    if (indegree == NULL || queue == NULL) {
        free(indegree);
        free(queue);
        return XT_ERR_OUT_OF_MEMORY;
    }

    for (i = 0U; i < graph->edge_count; i++) {
        if (graph->edges[i].source >= graph->node_count ||
            graph->edges[i].destination >= graph->node_count) {
            free(indegree);
            free(queue);
            return XT_ERR_INVALID_STATE;
        }
        indegree[graph->edges[i].destination]++;
    }

    for (i = 0U; i < graph->node_count; i++) {
        if (indegree[i] == 0U) {
            queue[tail++] = i;
        }
    }

    while (head < tail) {
        size_t node = queue[head++];
        visited++;
        for (i = 0U; i < graph->edge_count; i++) {
            if (graph->edges[i].source == node) {
                size_t dest = graph->edges[i].destination;
                if (indegree[dest] > 0U) {
                    indegree[dest]--;
                    if (indegree[dest] == 0U) {
                        queue[tail++] = dest;
                    }
                }
            }
        }
    }

    free(indegree);
    free(queue);
    return visited == graph->node_count ? XT_OK : XT_ERR_GRAPH_CYCLE;
}

xt_result xt_graph_process_offline(xt_graph *graph, const float *input,
                                   float *output, size_t frame_count,
                                   float input_gain)
{
    size_t i;
    float cumulative_gain = input_gain;

    if (graph == NULL || input == NULL || output == NULL) {
        return XT_ERR_INVALID_ARGUMENT;
    }
    if (xt_graph_validate(graph) != XT_OK) {
        return XT_ERR_INVALID_STATE;
    }

    /* Temporary starter behavior: apply gains in node insertion order.
       This is not yet graph-plan execution and is only for offline tests. */
    for (i = 0U; i < graph->node_count; i++) {
        cumulative_gain *= graph->nodes[i].gain;
    }
    for (i = 0U; i < frame_count; i++) {
        output[i] = input[i] * cumulative_gain;
    }
    return XT_OK;
}