#include <iostream>
#include <rte_eal.h>
#include <rte_ethdev.h>
#include <rte_vhost.h>
#include <cstring>
#include <thread>
#include <unistd.h>
#include <algorithm>
#include <arpa/inet.h>
#include <sys/epoll.h>

#define NUM_MBUFS 8192
#define NUM_RXTX_QUEUES 37

static int get_portid_by_name(const char* if_name) {
  uint16_t portid;
  RTE_ETH_FOREACH_DEV(portid) {
    char name[RTE_ETH_NAME_MAX_LEN];
    rte_eth_dev_get_name_by_port(portid, name);
    if (strcmp(name, if_name) == 0)
      return portid;
  }
  return -1; // Interface não encontrada
}

static rte_spinlock_t locks[RTE_MAX_ETHPORTS];

static void turn_on_intr(const int epfd, const uint16_t port_id, const uint16_t queue) {
  rte_spinlock_lock(&locks[port_id]);

  if (rte_eth_dev_rx_intr_ctl_q(port_id, queue, epfd, RTE_INTR_EVENT_ADD, nullptr) != 0)
    std::cerr << "Erro ao configurar a interrupção RX para a porta " << port_id << " fila " << queue << std::endl;

  rte_spinlock_unlock(&locks[port_id]);
}

static void configure_port(rte_mempool* mempool, const uint16_t port_id, bool has_interrupt = true) {
  rte_eth_dev_info dev_info {};
  if (rte_eth_dev_info_get(port_id, &dev_info) != 0) {
    std::cerr << "Erro ao obter informações da porta: " << port_id << std::endl;
    exit(0);
  }

  rte_eth_conf port_conf {
    .rxmode = {
      .mq_mode = RTE_ETH_MQ_RX_RSS,
      .offloads = RTE_ETH_RX_OFFLOAD_CHECKSUM
    },
    .intr_conf = {
      .rxq = has_interrupt
    }
  };

  uint16_t max_queues = NUM_RXTX_QUEUES;
  uint16_t nb_rx_queues = std::min(dev_info.max_rx_queues, max_queues);
  uint16_t nb_tx_queues = std::min(dev_info.max_tx_queues, max_queues);

  if (dev_info.tx_offload_capa & RTE_ETH_TX_OFFLOAD_MBUF_FAST_FREE)
    port_conf.txmode.offloads |=
            RTE_ETH_TX_OFFLOAD_MBUF_FAST_FREE;

  if (rte_eth_dev_configure(port_id, nb_rx_queues, nb_tx_queues, &port_conf) != 0) {
    std::cerr << "Erro ao configurar a porta: " << port_id << std::endl;
    exit(0);
  }

  uint16_t max_descs = NUM_MBUFS;
  uint16_t nb_rx_desc = std::min(dev_info.rx_desc_lim.nb_max, max_descs);
  uint16_t nb_tx_desc = std::min(dev_info.tx_desc_lim.nb_max, max_descs);
  if (rte_eth_dev_adjust_nb_rx_tx_desc(port_id, &nb_rx_desc, &nb_tx_desc) != 0) {
    std::cerr << "Erro ao ajustar descritores RX/TX: " << port_id << std::endl;
    exit(0);
  }

  for (uint16_t q = 0; q < nb_rx_queues; q++) {
    auto rxq_conf = dev_info.default_rxconf;
    rxq_conf.rx_thresh.pthresh = 8;
    rxq_conf.rx_thresh.hthresh = 8;
    rxq_conf.rx_thresh.wthresh = 4;
    rxq_conf.rx_drop_en = 0;

    if (rte_eth_rx_queue_setup(port_id, q, nb_rx_desc,
                               rte_eth_dev_socket_id(port_id), &rxq_conf, mempool) != 0) {
      std::cerr << "Erro ao configurar fila RX: " << q << " na porta: " << port_id << std::endl;
      exit(0);
    }
  }

  for (uint16_t q = 0; q < nb_tx_queues; q++) {
    auto txq_conf = dev_info.default_txconf;
    txq_conf.tx_thresh.pthresh = 8;
    txq_conf.tx_thresh.hthresh = 8;
    txq_conf.tx_thresh.wthresh = 4;

    if (rte_eth_tx_queue_setup(port_id, q, nb_tx_desc,
                               rte_eth_dev_socket_id(port_id), &txq_conf) != 0) {
      std::cerr << "Erro ao configurar fila TX: " << q << " na porta: " << port_id << std::endl;
      exit(0);
    }
  }

  if (rte_eth_dev_start(port_id) != 0) {
    std::cerr << "Erro ao iniciar a porta: " << port_id << std::endl;
    exit(0);
  }
}

static void net_to_virtio_queue(const int port_id, const int queue) {
  std::string virtio_net = std::string("virtio_user") + std::to_string(port_id);
  auto virtio_port_id = get_portid_by_name(virtio_net.c_str());

  if (virtio_port_id == -1) {
    std::cerr << "Erro: Interface " << virtio_net << " não encontrada" << std::endl;
    return;
  }

  rte_epoll_event event {};

  turn_on_intr(RTE_EPOLL_PER_THREAD, port_id, queue);

  while (true) {
    assert(rte_eth_dev_rx_intr_enable(port_id, queue) == 0);
    int nfds = rte_epoll_wait(RTE_EPOLL_PER_THREAD, &event, 1, -1);
    if (nfds == 0) continue;

    rte_mbuf* rx_packets[2048];
    uint16_t num_rx = rte_eth_rx_burst(port_id, queue, rx_packets, 2048);
    if (num_rx == 0) continue;

    rte_eth_tx_burst(virtio_port_id, 0, rx_packets, num_rx);

    rte_pktmbuf_free_bulk(rx_packets, num_rx);
  }
}

static void virtio_to_net_queue(const int port_id) {
  std::string virtio_net = std::string("virtio_user") + std::to_string(port_id);
  auto virtio_port_id = get_portid_by_name(virtio_net.c_str());

  if (virtio_port_id == -1) {
    std::cerr << "Erro: Interface " << virtio_net << " não encontrada" << std::endl;
    return;
  }

  while (true) {
    rte_mbuf* rx_packets[2048];
    uint16_t num_rx = rte_eth_rx_burst(virtio_port_id, 0, rx_packets, 2048);
    if (num_rx == 0) {
      usleep(100);
      continue;
    }

    rte_eth_tx_burst(port_id, NUM_RXTX_QUEUES-1, rx_packets, num_rx);

    rte_pktmbuf_free_bulk(rx_packets, num_rx);
  }
}

static void redirect_to_virtio(const int port_id) {
  rte_eth_dev_info dev_info {};
  rte_eth_dev_info_get(port_id, &dev_info);

  uint16_t nb_rx_queues = dev_info.nb_rx_queues;

  for (uint16_t q = 0; q < nb_rx_queues; q++) {
    std::thread thread(net_to_virtio_queue, port_id, q);

    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(q, &cpuset);
    pthread_setaffinity_np(thread.native_handle(),
                          sizeof(cpu_set_t), &cpuset);

    thread.detach();
  }

  std::thread (virtio_to_net_queue, port_id).detach();
}

int main(int argc, char **argv) {
  int ret = rte_eal_init(argc, argv);
  if (ret < 0) {
    std::cerr << "Erro na inicialização do EAL" << std::endl;
    return -1;
  }

  uint16_t nb_ports = rte_eth_dev_count_avail();
  uint16_t port_id;

  int port_count = 0;

  rte_mempool *mempool = rte_pktmbuf_pool_create("MBUF_POOL", NUM_MBUFS * (NUM_RXTX_QUEUES * 2) * nb_ports,
      256, 0, RTE_MBUF_DEFAULT_BUF_SIZE, rte_socket_id());

  if (mempool == nullptr) {
    std::cerr << "Erro ao alocar mempool: " << rte_strerror(rte_errno) << std::endl;
    return 0;
  }

  RTE_ETH_FOREACH_DEV(port_id) {
    char portname[32];
    char portargs[256];
    rte_ether_addr addr {};

    if (++port_count > nb_ports)
      break;

    rte_eth_macaddr_get(port_id, &addr);

    snprintf(portname, sizeof(portname), "virtio_user%u", port_id);
    snprintf(portargs, sizeof(portargs),
     "path=/dev/vhost-net,queues=1,queue_size=%u,iface=%s,mac=" RTE_ETHER_ADDR_PRT_FMT,
     1024, portname, RTE_ETHER_ADDR_BYTES(&addr));

    if (rte_eal_hotplug_add("vdev", portname, portargs) < 0) {
      rte_exit(EXIT_FAILURE, "Cannot create paired port for port %u\n", port_id);
    }

    // Atribuir IP e ativar a interface
    system("ifconfig virtio_user0 67.159.40.18/24 up");
    system("ip route add default via 67.159.40.17 dev virtio_user0");

    rte_spinlock_init(&locks[port_id]);

    configure_port(mempool, port_id, true);

    std::string virtio_net = std::string("virtio_user") + std::to_string(port_id);
    auto virtio_port_id = get_portid_by_name(virtio_net.c_str());

    if (virtio_port_id == -1) {
      std::cerr << "Erro: Interface " << virtio_net << " não encontrada" << std::endl;
      return 0;
    }

    configure_port(mempool, virtio_port_id, false);

    redirect_to_virtio(port_id);
  }

  while (true) {
    sleep(1);
  }

  return 0;
}