#include <string.h>
#include <pcap/pcap.h>
#include <glib.h>
#include <glib/gstdio.h>
#include <gio/gio.h>
#include <gst/gst.h>

#include "config.h"
#include "omvs_gst.h"

#define _OMVS_IPPROTO_UDP 17

static gint _omvs_net_dev_idx;
static gboolean _omvs_list_devs;
static gchar *_omvs_outdir = "omvs_out";
static gint _omvs_jobs = 1;
static gint _omvs_sleep = 1000;
static gint _omvs_timeout = 10000;
static gboolean _omvs_quiet;

static gchar **_omvs_net_dev_names;
static gchar **_omvs_net_dev_descs;
static gint _omvs_num_net_devs;

static GOptionEntry _omvs_entries[] =
{
  { "interface", 'i', 0, G_OPTION_ARG_INT, &_omvs_net_dev_idx,
    "Network device interface index", "dev_idx" },
  { "list-inteface", 'l', 0, G_OPTION_ARG_NONE, &_omvs_list_devs,
    "List network device interfaces", NULL },
  { "jobs", 'j', 0, G_OPTION_ARG_INT, &_omvs_jobs,
    "Number of jobs", "jobs" },
  { "output", 'o', 0, G_OPTION_ARG_FILENAME, &_omvs_outdir,
    "Output directory name", "dirname" },
  { "sleep", 's', 0, G_OPTION_ARG_INT, &_omvs_sleep,
    "Sleep time(milliseconds) between scans", "ms" },
  { "timeout", 't', 0, G_OPTION_ARG_INT, &_omvs_timeout,
    "Timeout time(milliseconds) in each scan", "ms" },
  { "quiet", 'q', 0, G_OPTION_ARG_NONE, &_omvs_quiet,
    "Print no log except for errors", NULL },
  { NULL }
};

typedef struct _OMVSEthHdr {
  guint8 ether_dhost[6];
  guint8 ether_shost[6];
  guint16 ether_type;
} __attribute__((packed)) OMVSEthHdr;

typedef struct _OMVSIPHdr {
  guint8 vhl;
  guint8 tos;
  guint16 tot_len;
  guint16 id;
  guint16 frag_off;
  guint8 ttl;
  guint8 protocol;
  guint16 check;
  guint32 saddr;
  guint32 daddr;
} OMVSIPHdr;

typedef struct _OMVSUDPHdr {
  guint16 uh_sport;
  guint16 uh_dport;
  guint16 uh_ulen;
  guint16 uh_sum;
} OMVSUDPHdr;

typedef struct _OMVSIPAddr {
  guint32 addr;
  OMVSGst *gsts;
  guint16 *ports;
  gint num_ports;
} OMVSIPAddr;

typedef struct _OMVSScanner {
  OMVSIPAddr *ipaddrs;
  gint num_ipaddrs;
} OMVSScanner;

static gint _omvs_init_net_devs_info(void);
static gint _omvs_print_net_devs_info(void);
static gint _omvs_deinit_net_devs_info(void);
static gint _omvs_start_scan(OMVSScanner *scanner);
static gpointer _omvs_start_scan_job(gpointer data);

static gint _omvs_init_net_devs_info(void) {
  gchar errbuf[PCAP_ERRBUF_SIZE];
  pcap_if_t *alldevs;
  pcap_if_t *dev;
  gint i;

  if (pcap_findalldevs(&alldevs, errbuf) != 0) {
    g_print("pcap_findalldevs() fail: %s", errbuf);
    return -1;
  }
  dev = alldevs;
  while (dev) {
    _omvs_num_net_devs++;
    dev = dev->next;
  }
  _omvs_net_dev_names = g_malloc(sizeof(gchar *) * _omvs_num_net_devs);
  _omvs_net_dev_descs = g_malloc(sizeof(gchar *) * _omvs_num_net_devs);

  for (i = 0, dev = alldevs; dev; i++, dev = dev->next) {
    _omvs_net_dev_names[i] = g_strdup(dev->name);
    _omvs_net_dev_descs[i] = g_strdup(dev->description);
  }
  pcap_freealldevs(alldevs);

  return 0;
}

static gint _omvs_print_net_devs_info(void) {
  gint i;

  for (i = 0; i < _omvs_num_net_devs; i++) {
    g_print("dev_idx(%d) dev_name(%s) dev_desc(%s)\n",
      i, _omvs_net_dev_names[i], _omvs_net_dev_descs[i]);
  }

  return 0;
}

static gint _omvs_deinit_net_devs_info(void) {
  gint i;

  for (i = 0; i < _omvs_num_net_devs; i++) {
    g_free(_omvs_net_dev_names[i]);
    g_free(_omvs_net_dev_descs[i]);
  }
  g_free(_omvs_net_dev_names);
  g_free(_omvs_net_dev_descs);

  return 0;
}

static gpointer _omvs_start_scan_job(gpointer data) {
  OMVSScanner *scanner;
  GError *error;
  GSocket *socket;
  GSocketAddress *socket_address;
  GInetAddress *inet_address;
  gint i;
  pcap_t *pcap;
  gchar errbuf[PCAP_ERRBUF_SIZE];
  const guint8 *packet;
  struct pcap_pkthdr *header;

  scanner = (OMVSScanner *)data;

  socket = g_socket_new(G_SOCKET_FAMILY_IPV4, G_SOCKET_TYPE_DATAGRAM,
    G_SOCKET_PROTOCOL_DEFAULT, &error);
  if (!socket) {
    g_printerr("[%p] can't create socket: %s\n", (void *)g_thread_self(),
      error->message);
    goto finish_return1;
  }

  inet_address = g_inet_address_new_any(G_SOCKET_FAMILY_IPV4);
  socket_address = g_inet_socket_address_new(inet_address, 0xffff);
  g_object_unref(inet_address);
  if (!g_socket_bind(socket, socket_address, TRUE, &error)) {
    g_printerr("[%p] can't bind socket: %s\n", (void *)g_thread_self(),
      error->message);
    goto finish_return2;
  }

  pcap = pcap_open_live(_omvs_net_dev_names[_omvs_net_dev_idx],
    sizeof(OMVSEthHdr) + sizeof(OMVSIPHdr) + sizeof(OMVSUDPHdr),
    TRUE, 500, errbuf);
  if (pcap == NULL) {
    g_printerr("[%p] can't open pcap: %s\n", (void *)g_thread_self(), errbuf);
    goto finish_return2;
  }
  if (pcap_datalink(pcap) != DLT_EN10MB) {
    g_printerr("[%p] %s is not ethernet\n", (void *)g_thread_self(),
      _omvs_net_dev_names[_omvs_net_dev_idx]);
    goto finish_return3;
  }

  for (i = 0; i < scanner->num_ipaddrs; i++) {
    OMVSIPAddr *o_ipaddr;
    guint8 addr[4];
    guint32 naddr;
    gchar *addr_str;
    GInetAddress *group_address;
    gint64 time1, time2;
    gint j;

    o_ipaddr = &scanner->ipaddrs[i];

    naddr = g_htonl(o_ipaddr->addr);
    addr[0] = (guint8)((o_ipaddr->addr >> 24) & 0xff);
    addr[1] = (guint8)((o_ipaddr->addr >> 16) & 0xff);
    addr[2] = (guint8)((o_ipaddr->addr >> 8) & 0xff);
    addr[3] = (guint8)((o_ipaddr->addr >> 0) & 0xff);

    addr_str =
      g_strdup_printf("%d.%d.%d.%d", addr[0], addr[1], addr[2], addr[3]);

    if (!_omvs_quiet) {
      g_print("[%p] start scanning %s\n", (void *)g_thread_self(), addr_str);
    }
    group_address = g_inet_address_new_from_string(addr_str);
    if (!g_socket_join_multicast_group(socket, group_address, FALSE,
      _omvs_net_dev_names[_omvs_net_dev_idx], &error)) {
      g_printerr("[%p] can't join %s: %s\n", (void *)g_thread_self(),
        addr_str, error->message);
    }

    time1 = g_get_monotonic_time();
    for (;;) {
      gint ret;
      guint caplen;
      OMVSIPHdr *ip;
      OMVSUDPHdr *udp;
      guint ip_hdr_len;
      gint udp_port;
      gboolean is_udp_port_found;
      gboolean is_gst_finished;
      gboolean is_all_gst_finished;

      ret = pcap_next_ex(pcap, &header, &packet);

      if (ret < 0) {
        break;
      }

      time2 = g_get_monotonic_time();
      if (time2 - time1 > (gint64)((gint64)_omvs_timeout * 1000)) {
        break;
      }

      if (ret == 0) {
        continue;
      }

      caplen = header->caplen;

      if (caplen < sizeof(OMVSEthHdr)) {
        continue;
      }

      packet += sizeof(OMVSEthHdr);
      caplen -= sizeof(OMVSEthHdr);
      if (caplen < sizeof(OMVSIPHdr)) {
        continue;
      }

      ip = (OMVSIPHdr *)packet;
      ip_hdr_len = (ip->vhl & 0x0f) * 4;
      if (caplen < ip_hdr_len) {
        continue;
      }

      if (ip->protocol != _OMVS_IPPROTO_UDP) {
        continue;
      }

      if (naddr != ip->daddr) {
        continue;
      }

      packet += ip_hdr_len;
      caplen -= ip_hdr_len;

      if (caplen < sizeof(OMVSUDPHdr)) {
        continue;
      }

      udp = (OMVSUDPHdr *)packet;
      udp_port = g_ntohs(udp->uh_dport);

      is_udp_port_found = FALSE;
      for (j = 0; j < o_ipaddr->num_ports; j++) {
        if (o_ipaddr->ports[j] == udp_port) {
          is_udp_port_found = TRUE;
          break;
        }
      }

      if (!is_udp_port_found) {
        gchar *uri;
        gchar *filename;
        uri = g_strdup_printf("udp://%s:%d", addr_str, udp_port);
        filename =
          g_strdup_printf("%s/%s-%d.png", _omvs_outdir, addr_str, udp_port);
        if (!_omvs_quiet) {
          g_print("[%p] trying to save %s to %s\n", (void *)g_thread_self(),
            uri, filename);
        }
        o_ipaddr->gsts = g_realloc(o_ipaddr->gsts,
          (o_ipaddr->num_ports + 1) * sizeof(OMVSGst));
        o_ipaddr->gsts[o_ipaddr->num_ports] = omvs_gst_open(uri, filename);
        g_free(filename);
        g_free(uri);
        o_ipaddr->ports = g_realloc(o_ipaddr->ports,
          (o_ipaddr->num_ports + 1) * sizeof(guint16));
        o_ipaddr->ports[o_ipaddr->num_ports] = udp_port;
        o_ipaddr->num_ports++;
      }

      is_gst_finished = FALSE;
      for (j = 0; j < o_ipaddr->num_ports; j++) {
        if (o_ipaddr->gsts[j]) {
          if (omvs_gst_is_finished(o_ipaddr->gsts[j])) {
            omvs_gst_close(o_ipaddr->gsts[j]);
            o_ipaddr->gsts[j] = NULL;
            is_gst_finished = TRUE;
          }
        }
      }

      is_all_gst_finished = FALSE;
      if (is_gst_finished) {
        is_all_gst_finished = TRUE;
        for (j = 0; j < o_ipaddr->num_ports; j++) {
          if (o_ipaddr->gsts[j]) {
            is_all_gst_finished = FALSE;
            break;
          }
        }
      }

      if (is_all_gst_finished) {
        break;
      }
    }

    for (j = 0; j < o_ipaddr->num_ports; j++) {
      omvs_gst_close(o_ipaddr->gsts[j]);
    }
    g_free(o_ipaddr->gsts);

    if (!_omvs_quiet) {
      g_print("[%p] finish scanning %s\n", (void *)g_thread_self(), addr_str);
    }
    if (!g_socket_leave_multicast_group(socket, group_address, FALSE,
      _omvs_net_dev_names[_omvs_net_dev_idx], &error)) {
      g_printerr("[%p] can't leave %s: %s\n", (void *)g_thread_self(),
        addr_str, error->message);
    }
    g_object_unref(group_address);
    g_free(addr_str);

    g_usleep(_omvs_sleep * 1000);
  }

finish_return3:
  pcap_close(pcap);
finish_return2:
  g_object_unref(socket_address);
  g_object_unref(socket);
finish_return1:
  return NULL;
}

static gint _omvs_start_scan(OMVSScanner *scanner) {
  gint i;
  gint num_ipaddrs_for_each_thread;
  OMVSScanner *scanners;
  GThread **jobs;

  num_ipaddrs_for_each_thread = scanner->num_ipaddrs / _omvs_jobs;
  if (num_ipaddrs_for_each_thread == 0) {
    num_ipaddrs_for_each_thread = 1;
    _omvs_jobs = scanner->num_ipaddrs;
  }

  scanners = g_malloc0(_omvs_jobs * sizeof(OMVSScanner));
  jobs = g_malloc0(_omvs_jobs * sizeof(GThread *));

  for (i = 0; i < _omvs_jobs; i++) {
    scanners[i].ipaddrs = &scanner->ipaddrs[i * num_ipaddrs_for_each_thread];
    scanners[i].num_ipaddrs = num_ipaddrs_for_each_thread;
  }
  scanners[_omvs_jobs - 1].num_ipaddrs += (scanner->num_ipaddrs % _omvs_jobs);

  for (i = 0; i < _omvs_jobs; i++) {
    jobs[i] = g_thread_new(NULL, _omvs_start_scan_job, &scanners[i]);
  }

  for (i = 0; i < _omvs_jobs; i++) {
    g_thread_join(jobs[i]);
  }

  g_free(jobs);
  g_free(scanners);

  return 0;
}

int main(int argc, char *argv[]) {
  GError *error = NULL;
  GOptionContext *context;
  int ret = 0;
  gint i;
  gint num_ipaddrs;
  gint idx_ipaddr;
  OMVSScanner omvs_scanner;

  if (_omvs_init_net_devs_info() != 0) {
    ret = -1;
    goto finish_return2;
  }

  context = g_option_context_new("<multicast ip address or net ... [#N]>");
  g_option_context_set_summary(context,
    "omvs scans to find free video streaming multicast ip addresses."
  );
  g_option_context_set_description(context,
    PACKAGE_STRING " is Written by Taeho Oh <ohhara@postech.edu>.\n"
    PACKAGE_URL "\n");
  g_option_context_add_main_entries(context, _omvs_entries, NULL);
  g_option_context_add_group(context, gst_init_get_option_group());
  if (!g_option_context_parse(context, &argc, &argv, &error)) {
    g_print("option parsing failed: %s\n", error->message);
    g_error_free(error);
    ret = -2;
    goto finish_return;
  }

  if (_omvs_list_devs) {
    _omvs_print_net_devs_info();
    ret = 0;
    goto finish_return;
  }

  if (argc < 2 || _omvs_net_dev_idx >= _omvs_num_net_devs) {
    gchar *help_msg;
    help_msg = g_option_context_get_help(context, TRUE, NULL);
    g_print("%s", help_msg);
    g_free(help_msg);
    ret = -2;
    goto finish_return;
  }

  g_mkdir(_omvs_outdir, 0755);

  num_ipaddrs = 0;
  for (i = 1; i < argc; i++) {
    GInetAddressMask *mask;
    GInetAddress *addr;
    guint len;
    mask = g_inet_address_mask_new_from_string(argv[i], &error);
    if (!mask) {
      g_print("%s parsing failed: %s\n", argv[i], error->message);
      g_error_free(error);
      ret = -3;
      goto finish_return;
    }
    if (g_inet_address_mask_get_family(mask) != G_SOCKET_FAMILY_IPV4) {
      g_print("%s is not ipv4 ip address\n", argv[i]);
      g_object_unref(mask);
      ret = -4;
      goto finish_return;
    }
    addr = g_inet_address_mask_get_address(mask);
    if (!g_inet_address_get_is_multicast(addr)) {
      g_print("%s is not multicast ip address\n", argv[i]);
      g_object_unref(mask);
      ret = -5;
      goto finish_return;
    }
    len = g_inet_address_mask_get_length(mask);
    if (len) {
      num_ipaddrs += (1 << (32 - len));
    } else {
      num_ipaddrs++;
    }
    g_object_unref(mask);
  }

  memset(&omvs_scanner, 0, sizeof(OMVSScanner));
  omvs_scanner.ipaddrs = g_malloc0(num_ipaddrs * sizeof(OMVSIPAddr));
  omvs_scanner.num_ipaddrs = num_ipaddrs;

  idx_ipaddr = 0;
  for (i = 1; i < argc; i++) {
    GInetAddressMask *mask;
    GInetAddress *addr;
    guint len;
    guint32 start_ipaddr;
    const guint8 *addr_bytes;
    mask = g_inet_address_mask_new_from_string(argv[i], &error);
    addr = g_inet_address_mask_get_address(mask);
    len = g_inet_address_mask_get_length(mask);
    addr_bytes = g_inet_address_to_bytes(addr);
    start_ipaddr =
      (((guint32)addr_bytes[0]) << 24) |
      (((guint32)addr_bytes[1]) << 16) |
      (((guint32)addr_bytes[2]) << 8) |
      (((guint32)addr_bytes[3]) << 0);
    if (len) {
      guint n;
      guint j;
      n = (1 << (32 - len));
      for (j = 0; j < n; j++) {
        omvs_scanner.ipaddrs[idx_ipaddr].addr = start_ipaddr;
        idx_ipaddr++;
        start_ipaddr++;
      }
    } else {
      omvs_scanner.ipaddrs[idx_ipaddr].addr = start_ipaddr;
      idx_ipaddr++;
    }
  }

  _omvs_start_scan(&omvs_scanner);
  g_free(omvs_scanner.ipaddrs);

finish_return:
  g_option_context_free(context);
  _omvs_deinit_net_devs_info();
finish_return2:

  return ret;
}