#include <time.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"

// OK
#define CHECKSUM_CARRY(x) (x = (x >> 16) + (x & 0xffff), (~(x + (x >> 16)) & 0xffff))
static inline int checksum(uint16_t *data, int len)
{
    int sum = 0;
    int nleft = len;
    uint16_t ans = 0;
    uint16_t *w = data;

    while (nleft > 1)
    {
        sum += *w++;
        nleft -= 2;
    }

    if (nleft == 1)
    {
        *(char *)(&ans) = *(char *)w;
        sum += ans;
    }

    return sum;
}

// OK
static inline void calc_tcp_seq_and_ack(const struct session *sess, enum session_direction inj_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 = s2c_direction_last_seq + s2c_direction_last_len
     *      ack = current_packet_seq
     */

    enum session_direction curr_dir = session_get_current_direction(sess);
    enum session_direction peer_dir = (curr_dir == SESSION_DIRECTION_C2S) ? SESSION_DIRECTION_S2C : SESSION_DIRECTION_C2S;
    const struct tcp_half *tcp_curr_half = &sess->tcp_halfs[curr_dir];
    const struct tcp_half *tcp_peer_half = &sess->tcp_halfs[peer_dir];
    if (inj_dir == curr_dir)
    {
        *seq = tcp_curr_half->seq;
        *ack = tcp_curr_half->ack;
    }
    else
    {
        *seq = tcp_peer_half->seq + tcp_peer_half->len;
        *ack = tcp_curr_half->seq;
    }
}

// OK
#define RANGE(rand, start, end) (start + rand % (end - start + 1)) // [start, end]
static inline void calc_ipid_ttl_win(uint16_t *ipid, uint8_t *ttl, uint16_t *win)
{
    struct timespec curtime;
    clock_gettime(CLOCK_THREAD_CPUTIME_ID, &curtime);
    uint64_t random = (0x013579ABCDEF ^ (uint64_t)curtime.tv_nsec);
    *ipid = (uint16_t)(RANGE(random, 32767, 65535));
    *ttl = (uint8_t)(RANGE(random, 48, 120));
    *win = (uint16_t)(RANGE(random, 1000, 1460));
}

// OK
static inline 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);
}

// OK
static inline 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);
    ipv4_hdr_set_checksum(iphdr, 0);
    uint16_t sum = checksum((uint16_t *)iphdr, hdr_len);
    sum = CHECKSUM_CARRY(sum);
    ipv4_hdr_set_checksum(iphdr, ntohs(sum));
}

// OK
static inline 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_checksum(tcphdr, 0);
    tcp_hdr_set_urg_ptr(tcphdr, 0);
    uint16_t sum = checksum((uint16_t *)tcphdr, sizeof(struct tcphdr));
    sum = CHECKSUM_CARRY(sum);
    tcp_hdr_set_checksum(tcphdr, ntohs(sum));
}

// OK
static inline void update_udp_hdr(struct udphdr *udphdr, uint16_t trim)
{
    uint16_t total = udp_hdr_get_total_len(udphdr);
    udp_hdr_set_total_len(udphdr, total - trim);
    udp_hdr_set_checksum(udphdr, 0);
    uint16_t sum = checksum((uint16_t *)udphdr, total - trim);
    sum = CHECKSUM_CARRY(sum);
    udp_hdr_set_checksum(udphdr, ntohs(sum));
}

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

// return 0: success, -1: failed

int stellar_inject_icmp_unreach(const struct session *sess, enum session_direction inj_dir, uint16_t thr_idx)
{
    // TODO
    return -1;
}

// OK
int stellar_inject_tcp_rst(const struct session *sess, enum session_direction inj_dir, uint16_t thr_idx)
{
    if (session_get_type(sess) != SESSION_TYPE_TCP)
    {
        session_inc_stat((struct session *)sess, inj_dir, STAT_INJ_PKTS_FAIL, 1);
        return -1;
    }

    const struct packet *pkt = session_get_1st_packet(sess, inj_dir);
    if (pkt == NULL)
    {
        session_inc_stat((struct session *)sess, inj_dir, STAT_INJ_PKTS_FAIL, 1);
        return -1;
    }

    uint8_t ttl = 0;
    uint16_t win = 0;
    uint16_t ipid = 0;
    uint32_t seq = 0;
    uint32_t ack = 0;

    int trim = 0;
    char buff[4096] = {0};
    struct ip *iphdr;
    struct ip6_hdr *ip6hdr;
    struct tcphdr *tcphdr;
    struct packet_layer *layer;

    int len = packet_get_len(pkt);
    int8_t layers = packet_get_layers(pkt);
    memcpy(buff, packet_get_data(pkt), len);

    calc_tcp_seq_and_ack(sess, inj_dir, &seq, &ack);
    calc_ipid_ttl_win(&ipid, &ttl, &win);

    for (int8_t i = layers - 1; i >= 0; i--)
    {
        layer = (struct packet_layer *)packet_get_layer(pkt, i);
        switch (layer->type)
        {
        case LAYER_TYPE_TCP:
            trim = layer->hdr_len + layer->pld_len - sizeof(struct tcphdr);
            tcphdr = (struct tcphdr *)(buff + layer->hdr_offset);
            update_tcp_hdr(tcphdr, seq, ack, win, TH_RST | TH_ACK);
            break;
        case LAYER_TYPE_IPV4:
            iphdr = (struct ip *)(buff + layer->hdr_offset);
            update_ip4_hdr(iphdr, ipid, ttl, trim);
            break;
        case LAYER_TYPE_IPV6:
            ip6hdr = (struct ip6_hdr *)(buff + layer->hdr_offset);
            update_ip6_hdr(ip6hdr, trim);
            break;
        default:
            break;
        }
    }

    struct packet inj_pkt;
    packet_parse(&inj_pkt, buff, len - trim);
    if (packet_io_inject(runtime->packet_io, thr_idx, &inj_pkt, 1) == 1)
    {
        session_inc_stat((struct session *)sess, inj_dir, STAT_INJ_PKTS_SUSS, 1);
        session_inc_stat((struct session *)sess, inj_dir, STAT_INJ_BYTES_SUSS, len - trim);
        return 0;
    }
    else
    {
        session_inc_stat((struct session *)sess, inj_dir, STAT_INJ_PKTS_FAIL, 1);
        return -1;
    }
}

int stellar_inject_tcp_payload(const struct session *sess, enum session_direction inj_dir, uint16_t thr_idx, const char *payload, uint16_t len)
{
    // TODO
    return -1;
}

int stellar_inject_udp_payload(const struct session *sess, enum session_direction inj_dir, uint16_t thr_idx, const char *payload, uint16_t len)
{
    // TODO
    return -1;
}

int stellar_inject_ctrl_msg(const struct session *sess, uint16_t *sids, int nr_sid, uint16_t thr_idx, const char *msg, uint16_t len)
{
    // TODO
    return -1;
}

int stellar_inject_stateless_pkt(uint16_t link_id, uint16_t thr_idx, const char *packet, uint16_t len)
{
    // TODO
    return -1;
}