#include <time.h>
#include <errno.h>
#include <assert.h>

#include "tcp_utils.h"
#include "udp_utils.h"
#include "ipv4_utils.h"
#include "ipv6_utils.h"
#include "packet_io.h"
#include "packet_priv.h"
#include "session_priv.h"
#include "stellar_priv.h"

#define INJECT_PACKET_LOG_ERROR(format, ...) LOG_ERROR("inject packet", format, ##__VA_ARGS__)
#define INJECT_PACKE_LOG_DEBUG(format, ...) LOG_DEBUG("inject packet", format, ##__VA_ARGS__)

struct tcp_fingerprint
{
    uint16_t ipid;
    uint8_t ttl;
    uint16_t win;
};

static uint16_t checksum(const char *data, uint16_t len)
{
    uint32_t sum = 0;
    const uint16_t *ip1 = (const uint16_t *)data;

    while (len > 1)
    {
        sum += *ip1++;
        if (sum & 0x80000000)
        {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }
        len -= 2;
    }

    while (sum >> 16)
    {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    return (~sum);
}

static uint16_t checksum_v4(const void *l4_hdr, uint16_t l4_total_len, uint8_t l4_proto, struct in_addr *src_addr, struct in_addr *dst_addr)
{
    uint16_t *ip_src = (uint16_t *)src_addr;
    uint16_t *ip_dst = (uint16_t *)dst_addr;
    const uint16_t *buffer = (u_int16_t *)l4_hdr;

    uint32_t sum = 0;
    size_t len = l4_total_len;

    while (len > 1)
    {
        sum += *buffer++;
        if (sum & 0x80000000)
        {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }
        len -= 2;
    }

    if (len & 1)
    {
        sum += *((uint8_t *)buffer);
    }

    sum += *(ip_src++);
    sum += *ip_src;
    sum += *(ip_dst++);
    sum += *ip_dst;
    sum += htons(l4_proto);
    sum += htons(l4_total_len);

    while (sum >> 16)
    {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    return ((uint16_t)(~sum));
}

static uint16_t checksum_v6(const void *l4_hdr, uint16_t l4_total_len, uint8_t l4_proto, struct in6_addr *src_addr, struct in6_addr *dst_addr)
{
    uint16_t *ip_src = (uint16_t *)src_addr;
    uint16_t *ip_dst = (uint16_t *)dst_addr;
    const uint16_t *buffer = (u_int16_t *)l4_hdr;

    uint32_t sum = 0;
    size_t len = l4_total_len;

    while (len > 1)
    {
        sum += *buffer++;
        if (sum & 0x80000000)
        {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }
        len -= 2;
    }

    if (len & 1)
    {
        sum += *((uint8_t *)buffer);
    }

    for (int i = 0; i < 8; i++)
    {
        sum += *ip_src;
        ip_src++;
    }

    for (int i = 0; i < 8; i++)
    {
        sum += *ip_dst;
        ip_dst++;
    }
    sum += htons(l4_proto);
    sum += htons(l4_total_len);

    while (sum >> 16)
    {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    return ((uint16_t)(~sum));
}

static void update_tcp_hdr(struct tcphdr *tcphdr, uint32_t seq, uint32_t ack, uint16_t win, uint8_t flags)
{
    tcp_hdr_set_seq(tcphdr, seq);
    tcp_hdr_set_ack(tcphdr, ack);
    tcp_hdr_set_hdr_len(tcphdr, sizeof(struct tcphdr));
    tcp_hdr_set_flags(tcphdr, flags);
    tcp_hdr_set_window(tcphdr, win);
    tcp_hdr_set_urg_ptr(tcphdr, 0);
    tcp_hdr_set_checksum(tcphdr, 0);
}

static void update_udp_hdr(struct udphdr *udphdr, int trim)
{
    uint16_t total = udp_hdr_get_total_len(udphdr);
    udp_hdr_set_total_len(udphdr, total - trim);
    udp_hdr_set_checksum(udphdr, 0);
}

static void update_ip4_hdr(struct ip *iphdr, uint16_t ipid, uint8_t ttl, int trim)
{
    int hdr_len = ipv4_hdr_get_hdr_len(iphdr);
    uint16_t total = ipv4_hdr_get_total_len(iphdr);
    ipv4_hdr_set_total_len(iphdr, total - trim);
    ipv4_hdr_set_ipid(iphdr, ipid);
    ipv4_hdr_set_ttl(iphdr, ttl);
    iphdr->ip_sum = 0;
    iphdr->ip_sum = checksum((char *)iphdr, hdr_len);
}

static void update_ip6_hdr(struct ip6_hdr *ip6hdr, int trim)
{
    uint16_t len = ipv6_hdr_get_payload_len(ip6hdr);
    ipv6_hdr_set_payload_len(ip6hdr, len - trim);
}

static inline void calc_tcp_seq_ack(const struct session *sess, enum flow_direction inject_dir, uint32_t *seq, uint32_t *ack)
{
    /*
     * +--------+  current packet  +---------+     C2S RST      +--------+
     * |        |----------------->|         |----------------->|        |
     * | Client |                  | Stellar |                  | Server |
     * |        |<-----------------|         |<-----------------|        |
     * +--------+    S2C RST       +---------+                  +--------+
     *
     * for example: current packet is C2S
     *
     * inject direction == current direction (inject C2S RST)
     *      seq = current_packet_seq
     *      ack = current_packet_ack
     *
     * inject direction != current direction (inject S2C RST)
     *      seq = current_packet_ack
     *      ack = current_packet_seq + current_packet_payload_len
     *  or if current packet is a SYN-ACK packet
     *      seq = current_packet_seq
     *      ack = current_packet_ack + current_packet_payload_len + 1
     */

    enum flow_direction curr_dir = session_get_flow_direction(sess);
    const struct tcp_half *tcp_curr_half = &sess->tcp_halfs[curr_dir];
    if (inject_dir == curr_dir)
    {
        *seq = tcp_curr_half->seq;
        *ack = tcp_curr_half->ack;
    }
    else
    {
        *seq = tcp_curr_half->ack;
        *ack = tcp_curr_half->seq + tcp_curr_half->len + (tcp_curr_half->flags & TH_SYN ? 1 : 0);
    }
}

static inline void calc_tcp_fingerprint(struct tcp_fingerprint *finger)
{
#define RANGE(rand, start, end) (start + rand % (end - start + 1)) // [start, end]

    struct timespec curtime;
    clock_gettime(CLOCK_THREAD_CPUTIME_ID, &curtime);
    uint64_t random = (0x013579ABCDEF ^ (uint64_t)curtime.tv_nsec);
    finger->ipid = (uint16_t)(RANGE(random, 32767, 65535));
    finger->ttl = (uint8_t)(RANGE(random, 48, 120));
    finger->win = (uint16_t)(RANGE(random, 1000, 1460));
}

// return packet length
int build_tcp_packet(const struct packet *first, uint16_t ip_id, uint8_t ip_ttl,
                     uint32_t tcp_seq, uint32_t tcp_ack, uint8_t tcp_flags, uint16_t tcp_win,
                     const char *tcp_pld, int pld_len, char *pkt_buff, int buff_size)
{
    int trim = 0;
    struct tcphdr *tcphdr;
    struct udphdr *udphdr;
    struct ip *iphdr;
    struct ip6_hdr *ip6hdr;
    struct packet_layer *curr;
    struct packet_layer *last;
    int len = packet_get_len(first);
    int8_t layers = packet_get_layers_number(first);

    if ((tcp_pld == NULL && pld_len > 0) || (tcp_pld != NULL && pld_len <= 0))
    {
        return -EINVAL;
    }

    if (len > buff_size)
    {
        return -ENOMEM;
    }

    memcpy(pkt_buff, packet_get_data(first), len);
    for (int8_t i = layers - 1; i >= 0; i--)
    {
        curr = (struct packet_layer *)packet_get_layer(first, i);
        switch (curr->type)
        {
        case LAYER_TYPE_TCP:
            trim = curr->hdr_len + curr->pld_len - sizeof(struct tcphdr) + pld_len;
            if (len - trim > buff_size)
            {
                return -ENOMEM;
            }
            tcphdr = (struct tcphdr *)(pkt_buff + curr->hdr_offset);
            update_tcp_hdr(tcphdr, tcp_seq, tcp_ack, tcp_win, tcp_flags);
            if (pld_len)
            {
                memcpy(pkt_buff + curr->hdr_offset + sizeof(struct tcphdr), tcp_pld, pld_len);
            }
            break;
        case LAYER_TYPE_UDP:
            udphdr = (struct udphdr *)(pkt_buff + curr->hdr_offset);
            update_udp_hdr(udphdr, trim);
            break;
        case LAYER_TYPE_IPV4:
            iphdr = (struct ip *)(pkt_buff + curr->hdr_offset);
            last = (struct packet_layer *)packet_get_layer(first, i + 1);
            if (last->type == LAYER_TYPE_TCP)
            {
                tcphdr = (struct tcphdr *)(pkt_buff + last->hdr_offset);
                tcphdr->th_sum = checksum_v4(tcphdr, len - trim - last->hdr_offset, IPPROTO_TCP, &iphdr->ip_src, &iphdr->ip_dst);
            }
            if (last->type == LAYER_TYPE_UDP)
            {
                udphdr = (struct udphdr *)(pkt_buff + last->hdr_offset);
                udphdr->uh_sum = checksum_v4(udphdr, len - trim - last->hdr_offset, IPPROTO_UDP, &iphdr->ip_src, &iphdr->ip_dst);
            }
            update_ip4_hdr(iphdr, ip_id, ip_ttl, trim);
            break;
        case LAYER_TYPE_IPV6:
            ip6hdr = (struct ip6_hdr *)(pkt_buff + curr->hdr_offset);
            last = (struct packet_layer *)packet_get_layer(first, i + 1);
            if (last->type == LAYER_TYPE_TCP)
            {
                tcphdr = (struct tcphdr *)(pkt_buff + last->hdr_offset);
                tcphdr->th_sum = checksum_v6(tcphdr, len - trim - last->hdr_offset, IPPROTO_TCP, &ip6hdr->ip6_src, &ip6hdr->ip6_dst);
            }
            if (last->type == LAYER_TYPE_UDP)
            {
                udphdr = (struct udphdr *)(pkt_buff + last->hdr_offset);
                udphdr->uh_sum = checksum_v6(udphdr, len - trim - last->hdr_offset, IPPROTO_UDP, &ip6hdr->ip6_src, &ip6hdr->ip6_dst);
            }
            update_ip6_hdr(ip6hdr, trim);
            break;
        case LAYER_TYPE_GRE:
            return -EPROTONOSUPPORT;
            // TODO
            break;
        default:
            break;
        }
    }

    return len - trim;
}

// return packet length
int build_udp_packet(const struct packet *first, const char *udp_pld, int pld_len, char *pkt_buff, int buff_size)
{
    int trim = 0;
    struct udphdr *udphdr;
    struct ip *iphdr;
    struct ip6_hdr *ip6hdr;
    struct packet_layer *curr;
    struct packet_layer *last;
    int len = packet_get_len(first);
    int8_t layers = packet_get_layers_number(first);

    if ((udp_pld == NULL && pld_len > 0) || (udp_pld != NULL && pld_len <= 0))
    {
        return -EINVAL;
    }

    if (len > buff_size)
    {
        return -ENOMEM;
    }

    memcpy(pkt_buff, packet_get_data(first), len);
    for (int8_t i = layers - 1; i >= 0; i--)
    {
        curr = (struct packet_layer *)packet_get_layer(first, i);
        switch (curr->type)
        {
        case LAYER_TYPE_UDP:
            trim = curr->hdr_len + curr->pld_len - sizeof(struct udphdr) + pld_len;
            if (len - trim > buff_size)
            {
                return -ENOMEM;
            }
            udphdr = (struct udphdr *)(pkt_buff + curr->hdr_offset);
            update_udp_hdr(udphdr, trim);
            if (pld_len)
            {
                memcpy(pkt_buff + curr->hdr_offset + sizeof(struct udphdr), udp_pld, pld_len);
            }
            break;
        case LAYER_TYPE_IPV4:
            iphdr = (struct ip *)(pkt_buff + curr->hdr_offset);
            last = (struct packet_layer *)packet_get_layer(first, i + 1);
            if (last->type == LAYER_TYPE_UDP)
            {
                udphdr = (struct udphdr *)(pkt_buff + last->hdr_offset);
                udphdr->uh_sum = checksum_v4(udphdr, len - trim - last->hdr_offset, IPPROTO_UDP, &iphdr->ip_src, &iphdr->ip_dst);
            }
            update_ip4_hdr(iphdr, 0, 0, trim);
            break;
        case LAYER_TYPE_IPV6:
            ip6hdr = (struct ip6_hdr *)(pkt_buff + curr->hdr_offset);
            last = (struct packet_layer *)packet_get_layer(first, i + 1);
            if (last->type == LAYER_TYPE_UDP)
            {
                udphdr = (struct udphdr *)(pkt_buff + last->hdr_offset);
                udphdr->uh_sum = checksum_v6(udphdr, len - trim - last->hdr_offset, IPPROTO_UDP, &ip6hdr->ip6_src, &ip6hdr->ip6_dst);
            }
            update_ip6_hdr(ip6hdr, trim);
            break;
        case LAYER_TYPE_GRE:
            return -EPROTONOSUPPORT;
            // TODO
            break;
        default:
            break;
        }
    }

    return len - trim;
}

int inject_tcp_packet(const struct session *sess, enum flow_direction inject_dir, uint8_t tcp_flags, const char *payload, uint16_t len)
{
#define TCP_FLAGS_LOG_FORMAT "URG:%d, ACK:%d, PSH:%d, RST:%d, SYN:%d, FIN:%d"
#define TCP_FLAGS_LOG_VALUE(flags)                                   \
    (((flags) & TH_URG) ? 1 : 0), (((flags) & TH_ACK) ? 1 : 0),      \
        (((flags) & TH_PUSH) ? 1 : 0), (((flags) & TH_RST) ? 1 : 0), \
        (((flags) & TH_SYN) ? 1 : 0), (((flags) & TH_FIN) ? 1 : 0)

    uint16_t thr_idx = stellar_get_current_thread_index();

    if (session_get_type(sess) != SESSION_TYPE_TCP)
    {
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        INJECT_PACKET_LOG_ERROR("session %ld is not a TCP session, cannot inject TCP packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d)",
                                session_get_id(sess), TCP_FLAGS_LOG_VALUE(tcp_flags), len);
        return 0;
    }

    const struct packet *pkt = session_get_first_packet(sess, inject_dir);
    if (pkt == NULL)
    {
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        INJECT_PACKET_LOG_ERROR("session %ld has no %s first packet, cannot inject TCP packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d)",
                                session_get_id(sess), flow_direction_to_str(inject_dir), TCP_FLAGS_LOG_VALUE(tcp_flags), len);
        return 0;
    }

    struct tcp_fingerprint finger = {0};
    uint32_t tcp_seq = 0;
    uint32_t tcp_ack = 0;
    char buff[4096] = {0};
    calc_tcp_seq_ack(sess, inject_dir, &tcp_seq, &tcp_ack);
    calc_tcp_fingerprint(&finger);

    int pkt_len = build_tcp_packet(pkt, finger.ipid, finger.ttl, tcp_seq, tcp_ack, tcp_flags, finger.win, payload, len, buff, sizeof(buff));
    if (pkt_len <= 0)
    {
        INJECT_PACKET_LOG_ERROR("session %ld build TCP %s packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d) failed, %s",
                                session_get_id(sess), flow_direction_to_str(inject_dir), TCP_FLAGS_LOG_VALUE(tcp_flags), len, strerror(len));
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        return 0;
    }

    struct inject_packet_meta meta = {0};
    meta.session_id = session_get_id(sess);
    session_get_route_ctx(sess, inject_dir, &meta.route);
    session_get_sid_list(sess, inject_dir, &meta.sids);

    struct packet inj_pkt;
    packet_parse(&inj_pkt, buff, pkt_len);
    packet_set_origin(&inj_pkt, PACKET_ORIGIN_USERSTACK);
    packet_set_origin_ctx(&inj_pkt, &meta);
    if (packet_io_inject(runtime->packet_io, thr_idx, &inj_pkt, 1) == 1)
    {
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_SUCCESS, 1);
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_BYTES_SUCCESS, pkt_len);
        INJECT_PACKE_LOG_DEBUG("session %ld inject TCP %s packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d) success",
                               session_get_id(sess), flow_direction_to_str(inject_dir), TCP_FLAGS_LOG_VALUE(tcp_flags), len);
        return pkt_len;
    }
    else
    {
        INJECT_PACKET_LOG_ERROR("session %ld inject TCP %s packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d) failed, packet I/O nospace",
                                session_get_id(sess), flow_direction_to_str(inject_dir), TCP_FLAGS_LOG_VALUE(tcp_flags), len);
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        return 0;
    }
}

int inject_udp_packet(const struct session *sess, enum flow_direction inject_dir, const char *payload, uint16_t len)
{
    uint16_t thr_idx = stellar_get_current_thread_index();

    if (session_get_type(sess) != SESSION_TYPE_UDP)
    {
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        INJECT_PACKET_LOG_ERROR("session %ld is not a UDP session, cannot inject UDP packet (payload len:%d)",
                                session_get_id(sess), len);
        return 0;
    }

    const struct packet *pkt = session_get_first_packet(sess, inject_dir);
    if (pkt == NULL)
    {
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        INJECT_PACKET_LOG_ERROR("session %ld has no %s first packet, cannot inject UDP packet (payload len:%d)",
                                session_get_id(sess), flow_direction_to_str(inject_dir), len);
        return 0;
    }

    char buff[4096] = {0};
    int pkt_len = build_udp_packet(pkt, payload, len, buff, sizeof(buff));
    if (pkt_len <= 0)
    {
        INJECT_PACKET_LOG_ERROR("session %ld build UDP %s packet (payload len:%d) failed, %s",
                                session_get_id(sess), flow_direction_to_str(inject_dir), len, strerror(len));
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        return 0;
    }

    struct inject_packet_meta meta = {0};
    meta.session_id = session_get_id(sess);
    session_get_route_ctx(sess, inject_dir, &meta.route);
    session_get_sid_list(sess, inject_dir, &meta.sids);

    struct packet inj_pkt;
    packet_parse(&inj_pkt, buff, pkt_len);
    packet_set_origin(&inj_pkt, PACKET_ORIGIN_USERSTACK);
    packet_set_origin_ctx(&inj_pkt, &meta);
    if (packet_io_inject(runtime->packet_io, thr_idx, &inj_pkt, 1) == 1)
    {
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_SUCCESS, 1);
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_BYTES_SUCCESS, pkt_len);
        INJECT_PACKE_LOG_DEBUG("session %ld inject UDP %s packet (payload len:%d) success",
                               session_get_id(sess), flow_direction_to_str(inject_dir), len);
        return pkt_len;
    }
    else
    {
        INJECT_PACKET_LOG_ERROR("session %ld inject UDP %s packet (payload len:%d) failed, packet I/O nospace",
                                session_get_id(sess), flow_direction_to_str(inject_dir), len);
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        return 0;
    }
}

/******************************************************************************
 * Public API
 ******************************************************************************/

int stellar_inject_tcp_rst(const struct session *sess, enum flow_direction inject_dir)
{
    return inject_tcp_packet(sess, inject_dir, TH_RST | TH_ACK, NULL, 0);
}

int stellar_inject_tcp_fin(const struct session *sess, enum flow_direction inject_dir)
{
    return inject_tcp_packet(sess, inject_dir, TH_FIN | TH_ACK, NULL, 0);
}

int stellar_inject_payload(const struct session *sess, enum flow_direction inject_dir, const char *payload, uint16_t len)
{
    switch (session_get_type(sess))
    {
    case SESSION_TYPE_TCP:
        return inject_tcp_packet(sess, inject_dir, TH_ACK, payload, len);
    case SESSION_TYPE_UDP:
        return inject_udp_packet(sess, inject_dir, payload, len);
    default:
        return 0;
    }
}

int stellar_inject_ctrl_msg(const struct session *sess, const struct sid_list *sids, const char *msg, uint16_t len)
{
    // TODO
    return 0;
}