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

#include "times.h"
#include "packet_build.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__)

static inline void calc_tcp_seq_ack(const struct session *sess, enum flow_direction inject_dir, uint32_t *seq, uint32_t *ack, uint8_t flags, uint16_t len)
{
    /*
     * +--------+  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_current_flow_direction(sess);
    struct tcp_half *tcp_curr_half = (struct tcp_half *)&sess->tcp_halfs[curr_dir];
    if (inject_dir == curr_dir)
    {
        *seq = uint32_add(tcp_curr_half->seq, tcp_curr_half->inject_inc_seq_offset);
        *ack = tcp_curr_half->ack;

        tcp_curr_half->inject_inc_seq_offset += len;
        // inject RST packer after FIN packer, seq should be increased by 1
        tcp_curr_half->inject_inc_seq_offset += (flags & TH_FIN) ? 1 : 0;
    }
    else
    {
        *seq = uint32_add(tcp_curr_half->ack, tcp_curr_half->inject_inc_ack_offset);
        *ack = uint32_add(tcp_curr_half->seq, tcp_curr_half->len + (tcp_curr_half->flags & TH_SYN ? 1 : 0));

        tcp_curr_half->inject_inc_ack_offset += len;
        // inject RST packer after FIN packer, ack should be increased by 1
        tcp_curr_half->inject_inc_ack_offset += (flags & TH_FIN) ? 1 : 0;
    }
}

// return 1 if success
// return 0 if failed
int stellar_send_packet(struct stellar *st, struct packet *pkt)
{
    uint64_t time_ms = stellar_get_monotonic_time_msec();
    uint16_t thr_idx = stellar_get_current_thread_index();
    struct packet_io *packet_io = stellar_get_packet_io(st);
    struct session_manager *sess_mgr = stellar_get_session_manager(st);
    session_manager_record_duplicated_packet(sess_mgr, pkt, time_ms);

    if (packet_get_origin_ctx(pkt))
    {
        // TODO
        assert(0);
        packet_io_egress(packet_io, thr_idx, pkt, 1);
        return 1;
    }
    else
    {
        return packet_io_inject(packet_io, thr_idx, pkt, 1);
    }
}

static int inject_tcp_packet(struct stellar *st, const struct session *sess, enum flow_direction inject_dir, uint8_t tcp_flags, const char *tcp_payload, uint16_t tcp_payload_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)

    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 %s is not a TCP session, cannot inject TCP packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d)",
                                session_get_id(sess), session_get0_readable_addr(sess),
                                TCP_FLAGS_LOG_VALUE(tcp_flags), tcp_payload_len);
        return 0;
    }

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

    uint32_t tcp_seq = 0;
    uint32_t tcp_ack = 0;
    calc_tcp_seq_ack(sess, inject_dir, &tcp_seq, &tcp_ack, tcp_flags, tcp_payload_len);

    struct packet *new_pkt = imitate_tcp_packet(origin_pkt, tcp_seq, tcp_ack, tcp_flags, tcp_payload, tcp_payload_len);
    if (new_pkt == NULL)
    {
        INJECT_PACKET_LOG_ERROR("session %ld %s build TCP %s packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d) failed",
                                session_get_id(sess), session_get0_readable_addr(sess),
                                flow_direction_to_str(inject_dir), TCP_FLAGS_LOG_VALUE(tcp_flags), tcp_payload_len);
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        return 0;
    }

    int pkt_len = packet_get_raw_len(new_pkt);

    if (stellar_send_packet(st, new_pkt) == 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 %s inject TCP %s packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d) success",
                               session_get_id(sess), session_get0_readable_addr(sess),
                               flow_direction_to_str(inject_dir), TCP_FLAGS_LOG_VALUE(tcp_flags), tcp_payload_len);
        return pkt_len;
    }
    else
    {
        INJECT_PACKET_LOG_ERROR("session %ld %s inject TCP %s packet (" TCP_FLAGS_LOG_FORMAT ", payload len:%d) failed, packet I/O nospace",
                                session_get_id(sess), session_get0_readable_addr(sess),
                                flow_direction_to_str(inject_dir), TCP_FLAGS_LOG_VALUE(tcp_flags), tcp_payload_len);
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        return 0;
    }
}

static int inject_udp_packet(struct stellar *st, const struct session *sess, enum flow_direction inject_dir, const char *payload, uint16_t len)
{
    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 %s is not a UDP session, cannot inject UDP packet (payload len:%d)",
                                session_get_id(sess), session_get0_readable_addr(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 %s has no %s first packet, cannot inject UDP packet (payload len:%d)",
                                session_get_id(sess), session_get0_readable_addr(sess),
                                flow_direction_to_str(inject_dir), len);
        return 0;
    }

    struct packet *new_pkt = imitate_udp_packet(pkt, payload, len);
    if (new_pkt == 0)
    {
        INJECT_PACKET_LOG_ERROR("session %ld %s build UDP %s packet (payload len:%d) failed",
                                session_get_id(sess), session_get0_readable_addr(sess),
                                flow_direction_to_str(inject_dir), len);
        session_inc_stat((struct session *)sess, inject_dir, STAT_INJECTED_PACKETS_FAILED, 1);
        return 0;
    }

    int pkt_len = packet_get_raw_len(new_pkt);
    if (stellar_send_packet(st, new_pkt) == 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 %s inject UDP %s packet (payload len:%d) success",
                               session_get_id(sess), session_get0_readable_addr(sess),
                               flow_direction_to_str(inject_dir), len);
        return pkt_len;
    }
    else
    {
        INJECT_PACKET_LOG_ERROR("session %ld %s inject UDP %s packet (payload len:%d) failed, packet I/O nospace",
                                session_get_id(sess), session_get0_readable_addr(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(struct stellar *st, const struct session *sess, enum flow_direction inject_dir)
{
    return inject_tcp_packet(st, sess, inject_dir, TH_RST | TH_ACK, NULL, 0);
}

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

int stellar_inject_tcp_payload(struct stellar *st, const struct session *sess, enum flow_direction inject_dir, const char *payload, uint16_t len)
{
    return inject_tcp_packet(st, sess, inject_dir, TH_ACK, payload, len);
}

int stellar_inject_udp_payload(struct stellar *st, const struct session *sess, enum flow_direction inject_dir, const char *payload, uint16_t len)
{
    return inject_udp_packet(st, sess, inject_dir, payload, len);
}

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