#include "retransmit.h"

#include <libwifi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "transmit.h"

static pthread_mutex_t retransmitter_mutex = PTHREAD_MUTEX_INITIALIZER;

void ack_map_init(struct ack_map *map)
{
  map->top = NULL;
}

void ack_map_add(struct ack_map *map, const unsigned char dst_addr[6], const unsigned char *packet, int packet_size)
{
  struct ack_map_node *cur = map->top;
  struct ack_map_node *prev = NULL;

  while (cur != NULL)
  {
    prev = cur;
    cur = cur->next;
  }

  struct ack_map_node *next = (struct ack_map_node *)malloc(sizeof(struct ack_map_node));
  next->next = NULL;
  memcpy(next->dst_addr, dst_addr, 6);
  next->packet = (unsigned char *)malloc(packet_size);
  memcpy(next->packet, packet, packet_size);
  next->packet_size = packet_size;
  next->retries_left = 7;

  if (prev == NULL)
  {
    map->top = next;
  }
  else
  {
    prev->next = next;
  }
}

void ack_map_pop(struct ack_map *map)
{
  if (map->top != NULL)
  {
    struct ack_map_node *cur = map->top;
    map->top = cur->next;

    free(cur->packet);
    free(cur);
  }
}

int ack_map_length(struct ack_map *map)
{
  int result = 0;

  struct ack_map_node *cur = map->top;
  while (cur != NULL)
  {
    result++;

    cur = cur->next;
  }

  return result;
}

void ack_map_clear(struct ack_map *map)
{
  while (map->top != NULL)
  {
    ack_map_pop(map);
  }
}

void *retransmit_thread(void *arg)
{
  struct retransmitter *retransmitter = (struct retransmitter *)arg;

  for (;;)
  {
    pthread_mutex_lock(&retransmitter_mutex);

    if (retransmitter->quit)
    {
      pthread_mutex_unlock(&retransmitter_mutex);

      break;
    }

    struct ack_map_node *cur = retransmitter->messages.top;
    if (cur != NULL)
    {
      send_packet_with_radiotap(retransmitter->sock, retransmitter->device_index, cur->dst_addr, cur->packet, cur->packet_size);

      // Set retransmit flag for future.
      cur->packet[1] |= (1 << 3);

      cur->retries_left--;

      if (cur->retries_left <= 0)
      {
        ack_map_pop(&retransmitter->messages);
      }
    }

    pthread_mutex_unlock(&retransmitter_mutex);

    usleep(5000);
  }

  return NULL;
}

struct retransmitter *start_retransmit_thread(int sock, int device_index)
{
  struct retransmitter *retransmitter = (struct retransmitter *)malloc(sizeof(struct retransmitter));
  retransmitter->sock = sock;
  retransmitter->device_index = device_index;
  ack_map_init(&retransmitter->messages);
  retransmitter->quit = false;

  pthread_create(&retransmitter->thread_id, NULL, &retransmit_thread, retransmitter);

  return retransmitter;
}

void send_until_acked(struct retransmitter *retransmitter, const unsigned char dst_addr[6], const unsigned char *packet, int packet_size)
{
  pthread_mutex_lock(&retransmitter_mutex);

  ack_map_add(&retransmitter->messages, dst_addr, packet, packet_size);

  printf("Adding message for " MACSTR ". Current length: %d\n", MAC2STR(dst_addr), ack_map_length(&retransmitter->messages));

  pthread_mutex_unlock(&retransmitter_mutex);
}

void handle_ack(struct retransmitter *retransmitter)
{
  pthread_mutex_lock(&retransmitter_mutex);

  if (retransmitter->messages.top != NULL)
  {
    retransmitter->messages.top->retries_left = 0;
  }

  printf("Received ACK. Current length: %d\n", ack_map_length(&retransmitter->messages));

  pthread_mutex_unlock(&retransmitter_mutex);
}

void free_retransmitter(struct retransmitter *retransmitter)
{
  pthread_mutex_lock(&retransmitter_mutex);
  retransmitter->quit = true;
  pthread_mutex_unlock(&retransmitter_mutex);

  pthread_join(retransmitter->thread_id, NULL);

  ack_map_clear(&retransmitter->messages);
  free(retransmitter);
}