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

#include "udp_utils.h"
#include "tcp_utils.h"
#include "ipv4_utils.h"
#include "ipv6_utils.h"
#include "packet_priv.h"
#include "marsio.h"

/******************************************************************************
 * set and get from struct packet
 ******************************************************************************/

void packet_set_origin(struct packet *pkt, enum packet_origin origin)
{
    pkt->origin = origin;
}

enum packet_origin packet_get_origin(const struct packet *pkt)
{
    return pkt->origin;
}

int8_t packet_get_layers_number(const struct packet *pkt)
{
    return pkt->layers_used;
}

const struct packet_layer *packet_get_layer(const struct packet *pkt, int8_t idx)
{
    if (idx < 0 || idx >= pkt->layers_used)
    {
        return NULL;
    }
    return &pkt->layers[idx];
}

const char *packet_get_data(const struct packet *pkt)
{
    return pkt->data_ptr;
}

uint16_t packet_get_len(const struct packet *pkt)
{
    return pkt->data_len;
}

const char *packet_get_payload(const struct packet *pkt)
{
    if (pkt->layers_used == 0)
    {
        return NULL;
    }
    return pkt->layers[pkt->layers_used - 1].pld_ptr;
}

uint16_t packet_get_payload_len(const struct packet *pkt)
{
    if (pkt->layers_used == 0)
    {
        return 0;
    }
    return pkt->layers[pkt->layers_used - 1].pld_len;
}

void packet_set_action(struct packet *pkt, enum packet_action action)
{
    pkt->action = action;
}

enum packet_action packet_get_action(const struct packet *pkt)
{
    return pkt->action;
}

void packet_set_origin_ctx(struct packet *pkt, void *ctx)
{
    pkt->origin_ctx = ctx;
}

void *packet_get_origin_ctx(const struct packet *pkt)
{
    return pkt->origin_ctx;
}

int packet_is_fragment(const struct packet *pkt)
{
    return (pkt->frag_layer) ? 1 : 0;
}

struct packet *packet_new(uint16_t pkt_len)
{
    struct packet *pkt = (struct packet *)calloc(1, sizeof(struct packet) + pkt_len);
    if (pkt == NULL)
    {
        return NULL;
    }
    pkt->data_len = pkt_len;
    pkt->data_ptr = (const char *)pkt + sizeof(struct packet);
    pkt->origin = PACKET_ORIGIN_USERHEAP;

    return pkt;
}

void packet_free(struct packet *pkt)
{
    if (pkt && pkt->origin == PACKET_ORIGIN_USERHEAP)
    {
        free((void *)pkt);
    }
}

struct packet *packet_dup(const struct packet *pkt)
{
    if (pkt == NULL)
    {
        return NULL;
    }

    struct packet *dup_pkt = packet_new(pkt->data_len);
    if (dup_pkt == NULL)
    {
        return NULL;
    }

    dup_pkt->layers_used = pkt->layers_used;
    dup_pkt->layers_size = pkt->layers_size;

    memcpy((char *)dup_pkt->data_ptr, pkt->data_ptr, pkt->data_len);
    dup_pkt->data_len = pkt->data_len;

    dup_pkt->origin_ctx = NULL;
    dup_pkt->origin = PACKET_ORIGIN_USERHEAP;
    dup_pkt->action = PACKET_ACTION_DROP;

    for (int8_t i = 0; i < pkt->layers_used; i++)
    {
        dup_pkt->layers[i].type = pkt->layers[i].type;
        dup_pkt->layers[i].hdr_ptr = dup_pkt->data_ptr + pkt->layers[i].hdr_offset;
        dup_pkt->layers[i].pld_ptr = dup_pkt->data_ptr + pkt->layers[i].hdr_offset + pkt->layers[i].hdr_len;
        dup_pkt->layers[i].hdr_offset = pkt->layers[i].hdr_offset;
        dup_pkt->layers[i].hdr_len = pkt->layers[i].hdr_len;
        dup_pkt->layers[i].pld_len = pkt->layers[i].pld_len;
    }

    // update frag_layer
    if (pkt->frag_layer)
    {
        dup_pkt->frag_layer = &dup_pkt->layers[pkt->frag_layer - pkt->layers];
    }

    return dup_pkt;
}

/******************************************************************************
 * set and get from mbuff
 ******************************************************************************/

void packet_set_ctrl(struct packet *pkt)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        marsio_buff_set_ctrlbuf(mbuff);
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to set ctrl");
    }
}

int packet_is_ctrl(const struct packet *pkt)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        return marsio_buff_is_ctrlbuf(mbuff);
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to check ctrl");
        return 0;
    }
}

void packet_set_direction(struct packet *pkt, enum packet_direction dir)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        if (marsio_buff_set_metadata(mbuff, MR_BUFF_DIR, &dir, sizeof(dir)) != 0)
        {
            PACKET_LOG_ERROR("failed to set direction");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to set direction");
    }
}

enum packet_direction packet_get_direction(const struct packet *pkt)
{
    enum packet_direction dir = PACKET_DIRECTION_INCOMING;
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        if (marsio_buff_get_metadata(mbuff, MR_BUFF_DIR, &dir, sizeof(dir)) <= 0)
        {
            PACKET_LOG_ERROR("failed to get direction");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to get direction");
    }

    return dir;
}

void packet_set_session_id(struct packet *pkt, uint64_t sess_id)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        if (marsio_buff_set_metadata(mbuff, MR_BUFF_SESSION_ID, &sess_id, sizeof(sess_id)) != 0)
        {
            PACKET_LOG_ERROR("failed to set session id");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to set session id");
    }
}

uint64_t packet_get_session_id(const struct packet *pkt)
{
    uint64_t sess_id = 0;
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        if (marsio_buff_get_metadata(mbuff, MR_BUFF_SESSION_ID, &sess_id, sizeof(sess_id)) <= 0)
        {
            PACKET_LOG_ERROR("failed to get session id");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to get session id");
    }

    return sess_id;
}

void packet_set_domain(struct packet *pkt, uint64_t domain)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
// TODO
#if 0
        if (marsio_buff_set_metadata(mbuff, MR_BUFF_DOMAIN, &domain, sizeof(domain)) != 0)
        {
            PACKET_LOG_ERROR("failed to set domain");
        }
#endif
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to set domain");
    }
}

uint64_t packet_get_domain(const struct packet *pkt)
{
    uint64_t domain = 0;
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
// TODO
#if 0
        if (marsio_buff_get_metadata(mbuff, MR_BUFF_DOMAIN, &domain, sizeof(domain)) <= 0)
        {
            PACKET_LOG_ERROR("failed to get domain");
        }
#endif
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to get domain");
    }

    return domain;
}

void packet_set_route_ctx(struct packet *pkt, const struct route_ctx *ctx)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        if (marsio_buff_set_metadata(mbuff, MR_BUFF_ROUTE_CTX, (void *)ctx->data, ctx->used) != 0)
        {
            PACKET_LOG_ERROR("failed to set route ctx");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to set route ctx");
    }
}

void packet_get_route_ctx(const struct packet *pkt, struct route_ctx *ctx)
{
    ctx->used = 0;
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        ctx->used = marsio_buff_get_metadata(mbuff, MR_BUFF_ROUTE_CTX, ctx->data, sizeof(ctx->data));
        if (ctx->used <= 0)
        {
            PACKET_LOG_ERROR("failed to get route ctx");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to get route ctx");
    }
}

void packet_set_sid_list(struct packet *pkt, const struct sid_list *list)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        if (marsio_buff_set_sid_list(mbuff, (sid_t *)list->sid, list->used) != 0)
        {
            PACKET_LOG_ERROR("failed to set sid list");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to set sid list");
    }
}

void packet_get_sid_list(const struct packet *pkt, struct sid_list *list)
{
    list->used = 0;
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        list->used = marsio_buff_get_sid_list(mbuff, (sid_t *)list->sid, sizeof(list->sid) / sizeof(list->sid[0]));
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to get sid list");
    }
}

void packet_prepend_sid_list(struct packet *pkt, const struct sid_list *list)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        if (marsio_buff_prepend_sid_list(mbuff, (sid_t *)list->sid, list->used) != 0)
        {
            PACKET_LOG_ERROR("failed to prepend sid list");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to prepend sid list");
    }
}

void packet_append_sid_list(struct packet *pkt, const struct sid_list *list)
{
    if (packet_get_origin(pkt) == PACKET_ORIGIN_MARSIO)
    {
        marsio_buff_t *mbuff = (marsio_buff_t *)packet_get_origin_ctx(pkt);
        assert(mbuff != NULL);
        if (marsio_buff_append_sid_list(mbuff, (sid_t *)list->sid, list->used) != 0)
        {
            PACKET_LOG_ERROR("failed to append sid list");
        }
    }
    else
    {
        PACKET_LOG_WARN("packet origin is not marsio, failed to append sid list");
    }
}

int packet_get_fingerprint(const struct packet *pkt, char *buff, int size)
{
    int used = 0;
    char tuple_str[256] = {0};

    struct tuple6 tuple;
    const struct tcphdr *tcp_hdr = NULL;
    const struct udphdr *udp_hdr = NULL;
    const struct ip *ipv4_hdr = NULL;
    const struct ip6_hdr *ipv6_hdr = NULL;
    const struct packet_layer *l3_layer = NULL;
    const struct packet_layer *l4_layer = NULL;
    memset(buff, 0, size);

    if (packet_get_innermost_tuple6(pkt, &tuple) != 0)
    {
        return used;
    }
    tuple6_to_str(&tuple, tuple_str, sizeof(tuple_str));
    used += snprintf(buff + used, size - used, "tuple6: %s, ", tuple_str);

    l3_layer = packet_get_innermost_layer(pkt, LAYER_TYPE_L3);
    if (l3_layer == NULL || l3_layer->hdr_ptr == NULL)
    {
        return used;
    }
    switch (l3_layer->type)
    {
    case LAYER_TYPE_IPV4:
        ipv4_hdr = (const struct ip *)l3_layer->hdr_ptr;
        used += snprintf(buff + used, size - used, "ip_ttl: %u, ip_id: %u, ", ipv4_hdr_get_ttl(ipv4_hdr), ipv4_hdr_get_ipid(ipv4_hdr));
        break;
    case LAYER_TYPE_IPV6:
        ipv6_hdr = (const struct ip6_hdr *)l3_layer->hdr_ptr;
        used += snprintf(buff + used, size - used, "ip6_hlim: %u, ", ipv6_hdr_get_hop_limit(ipv6_hdr));
        break;
    default:
        break;
    }

    l4_layer = packet_get_innermost_layer(pkt, LAYER_TYPE_L4);
    if (l4_layer == NULL || l4_layer->hdr_ptr == NULL)
    {
        return used;
    }
    switch (l4_layer->type)
    {
    case LAYER_TYPE_TCP:
        tcp_hdr = (const struct tcphdr *)l4_layer->hdr_ptr;
        used += snprintf(buff + used, size - used, "tcp_seq: %u, tcp_ack: %u, tcp_checksum: %u, ", tcp_hdr_get_seq(tcp_hdr), tcp_hdr_get_ack(tcp_hdr), tcp_hdr_get_checksum(tcp_hdr));
        break;
    case LAYER_TYPE_UDP:
        udp_hdr = (const struct udphdr *)l4_layer->hdr_ptr;
        used += snprintf(buff + used, size - used, "udp_checksum: %u, ", udp_hdr_get_checksum(udp_hdr));
        break;
    default:
        break;
    }

    return used;
}