#include <stdio.h>
#include <netdb.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <arpa/nameser.h>
#include <resolv.h>
#include <errno.h>
#include <getdns/getdns.h>
#include <string.h>
#include "ip.h"
#include "ipalloc.h"
#include "fmt.h"
#include "alloc.h"
#include "str.h"
#include "stralloc.h"
#include "dns.h"
#include "case.h"

static getdns_context *context = NULL;
static stralloc glue = {0};

static int getdns_resolve_query(const char *name, uint16_t type,
    getdns_dict **response)
{
  if (getdns_general_sync(context, name, type, NULL, response)
      != GETDNS_RETURN_GOOD)
    return DNS_HARD;

  uint32_t error;
  getdns_dict_get_int(*response, "status", &error);
  if (error != GETDNS_RESPSTATUS_GOOD)
    return DNS_HARD;

  return 0;
}

static int getdns_resolve_loop(getdns_dict *response, uint16_t type,
    int (*cb_rr)(getdns_dict *response, getdns_dict *rdata, void *userarg),
    void *userarg)
{
  getdns_list *answers;
  if (getdns_dict_get_list(response, "replies_tree", &answers)
      != GETDNS_RETURN_GOOD)
    return DNS_HARD;

  size_t num_answers, rec_ix, num_cb = 0;
  getdns_list_get_length(answers, &num_answers);
  for (rec_ix = 0; rec_ix < num_answers; ++rec_ix) {
    getdns_dict *record;
    getdns_list_get_dict(answers, rec_ix, &record);

    getdns_list *answer;
    getdns_dict_get_list(record, "answer", &answer);

    size_t num_rr, rr_ix;
    getdns_list_get_length(answer, &num_rr);
    for (rr_ix = 0; rr_ix < num_rr; ++rr_ix) {
      getdns_dict *rr, *rr_rdata;
      uint32_t rr_type;

      getdns_list_get_dict(answer, rr_ix, &rr);
      getdns_dict_get_int(rr, "type", &rr_type);
      if (rr_type != type)
        continue;

      if (cb_rr != NULL)
      {
        getdns_dict_get_dict(rr, "rdata", &rr_rdata);
        int cb = cb_rr(response, rr_rdata, userarg);
        if (cb != 0)
          return cb;
      }
      ++num_cb;
    }
  }

  return num_cb;
}

static int getdns_resolve(const char *name, uint16_t type,
    int (*cb_rr)(getdns_dict *response, getdns_dict *rdata, void *userarg),
    void *userarg)
{
  getdns_dict *response = NULL;
  int r = getdns_resolve_query(name, type, &response);
  r |= getdns_resolve_loop(response, type, cb_rr, userarg);
  getdns_dict_destroy(response);
  return r;
}

static int getdns_resolve_num(getdns_dict *response, uint16_t type, size_t *num)
{
  int r = getdns_resolve_loop(response, type, NULL, NULL);
  if (r < 0)
    return r;
  *num = r;
  return 0;
}

void dns_init(int flagsearch)
{
  getdns_context_create(&context, 1);
  getdns_context_set_resolution_type(context, GETDNS_RESOLUTION_STUB);
  if (!flagsearch)
    getdns_context_set_append_name(context, GETDNS_APPEND_NAME_NEVER);
  else
    fprintf(stderr, "not supported!\n"); //FIXME
}

static int dns_cname_cb(getdns_dict *response, getdns_dict *rdata,
    void *userarg)
{
  stralloc *sa = userarg;

  getdns_bindata *bindata = NULL;
  if (getdns_dict_get_bindata(rdata, "cname", &bindata)
      != GETDNS_RETURN_GOOD)
    return DNS_SOFT;

  char *dname;
  getdns_convert_dns_name_to_fqdn(bindata, &dname);
  dname[strlen(dname) - 1] = '\0';
  if (!stralloc_copys(sa, dname)) {
    free(dname);
    return DNS_MEM;
  }
  free(dname);
  return 0;
}

int dns_cname(stralloc *sa)
{
  int loop;
  for (loop = 0; loop < 10; ++loop)
  {
    if (!sa->len)
      return loop;
    if (sa->s[sa->len - 1] == ']')
      return loop;
    if (sa->s[sa->len - 1] == '.') {
      --sa->len;
      continue;
    }

    if (!stralloc_0(sa))
      return DNS_MEM;
    switch(getdns_resolve(sa->s, GETDNS_RRTYPE_CNAME, dns_cname_cb, sa))
    {
      case DNS_MEM:
        return DNS_MEM;
      case DNS_SOFT:
        return DNS_SOFT;
      case DNS_HARD:
      case 0:
        return loop;
    }
  }
  return DNS_HARD; /* alias loop */
}

static int iaafmt(char *s, struct ip_address *ip)
{
  unsigned int i, len = 0;
  i = fmt_ulong(s, (unsigned long) ip->d[3]); len += i; if (s) s += i;
  i = fmt_str(s, "."); len += i; if (s) s += i;
  i = fmt_ulong(s, (unsigned long) ip->d[2]); len += i; if (s) s += i;
  i = fmt_str(s, "."); len += i; if (s) s += i;
  i = fmt_ulong(s, (unsigned long) ip->d[1]); len += i; if (s) s += i;
  i = fmt_str(s, "."); len += i; if (s) s += i;
  i = fmt_ulong(s, (unsigned long) ip->d[0]); len += i; if (s) s += i;
  i = fmt_str(s, ".in-addr.arpa."); len += i; if (s) s += i;
  return len;
}

static int dns_ptr_cb(getdns_dict *response, getdns_dict *rdata,
    void *userarg)
{
  stralloc *sa = userarg;

  getdns_bindata *bindata = NULL;
  if (getdns_dict_get_bindata(rdata, "rdata_raw", &bindata)
      != GETDNS_RETURN_GOOD)
    return DNS_SOFT;

  char *dname;
  getdns_convert_dns_name_to_fqdn(bindata, &dname);
  dname[strlen(dname) - 1] = '\0';
  if (!stralloc_copys(sa, dname)) {
    free(dname);
    return DNS_MEM;
  }
  free(dname);
  return 0;
}

int dns_ptr(stralloc *sa, struct ip_address *ip)
{
  if (!stralloc_ready(sa, iaafmt((char *)0, ip)))
    return DNS_MEM;
  sa->len = iaafmt(sa->s, ip);
  if (!stralloc_0(sa))
    return DNS_MEM;

  int num = getdns_resolve(sa->s, GETDNS_RRTYPE_PTR, dns_ptr_cb, sa);
  return (num < 0) ? num : 0;
}

struct dns_ipplus_user {
  ipalloc *ia;
#ifdef IX_FQDN
  char *fqdn;
#endif
  int pref;
};

static int dns_ipplus_cb(getdns_dict *response, getdns_dict *rdata,
    void *userarg)
{
  struct dns_ipplus_user *u = userarg;
  struct ip_mx ix = {0};

  getdns_bindata *bindata = NULL;
  if (getdns_dict_get_bindata(rdata, "ipv4_address", &bindata)
      != GETDNS_RETURN_GOOD)
    return DNS_SOFT;

  memcpy(ix.ip.d, bindata->data, 4);
  ix.pref = u->pref;
#ifdef IX_FQDN
  ix.fqdn = u->fqdn;
#endif
  if (!ipalloc_append(u->ia, &ix))
    return DNS_MEM;

  return 0;
}

static int dns_ipplus(ipalloc *ia, stralloc *sa, int pref)
{
  struct ip_mx ix = {0};

  if (!stralloc_copy(&glue, sa))
    return DNS_MEM;
  if (!stralloc_0(&glue))
    return DNS_MEM;
  if (glue.s[0]) {
    if (!glue.s[ip_scan(glue.s, &ix.ip)] || !glue.s[ip_scanbracket(glue.s, &ix.ip)])
    {
      if (!ipalloc_append(ia,&ix))
        return DNS_MEM;
      return 0;
    }
  }

  struct dns_ipplus_user u = {
    .ia   = ia,
#ifdef IX_FQDN
    .fqdn = glue.s,
#endif
    .pref = pref,
  };
  int num = getdns_resolve(glue.s, GETDNS_RRTYPE_A, dns_ipplus_cb, &u);
#ifdef IX_FQDN
  glue.s = 0;
#endif
  if (num == 0)
    num = DNS_HARD;
  return (num < 0) ? num : 0;
}

int dns_ip(ipalloc *ia, stralloc *sa)
{
  if (!ipalloc_readyplus(ia, 0))
    return DNS_MEM;
  ia->len = 0;
  return dns_ipplus(ia, sa, 0);
}

struct dns_mxip_user {
  unsigned num;
  struct mx_user {
    stralloc sa;
    unsigned short p;
  } *mx;
};

static int dns_mxip_cb(getdns_dict *response, getdns_dict *rdata,
    void *userarg)
{
  struct dns_mxip_user *mx = userarg;

  getdns_bindata *bindata = NULL;
  if (getdns_dict_get_bindata(rdata, "exchange", &bindata)
      != GETDNS_RETURN_GOOD)
    return DNS_SOFT;

  char *dname;
  getdns_convert_dns_name_to_fqdn(bindata, &dname);
  dname[strlen(dname) - 1] = '\0';

  uint32_t pref;
  getdns_dict_get_int(rdata, "preference", &pref);

  mx->mx[mx->num].p = pref;
  mx->mx[mx->num].sa.s = 0;
  if (!stralloc_copys(&mx->mx[mx->num].sa, dname))
  {
    free(dname);
    return DNS_MEM;
  }
  mx->num++;
  free(dname);
  return 0;
}

int dns_mxip(ipalloc *ia, stralloc *sa, unsigned long random)
{
  struct ip_mx ix = {0};

  if (!ipalloc_readyplus(ia, 0))
    return DNS_MEM;
  ia->len = 0;

  if (!stralloc_copy(&glue, sa))
    return DNS_MEM;
  if (!stralloc_0(&glue))
    return DNS_MEM;
  if (glue.s[0]) {
    if (!glue.s[ip_scan(glue.s, &ix.ip)] || !glue.s[ip_scanbracket(glue.s, &ix.ip)])
    {
      if (!ipalloc_append(ia,&ix))
        return DNS_MEM;
      return 0;
    }
  }

  getdns_dict *response = NULL;
  size_t num_rr = 0;
  int ret = getdns_resolve_query(glue.s, GETDNS_RRTYPE_MX, &response);
  ret |= getdns_resolve_num(response, GETDNS_RRTYPE_MX, &num_rr);
  if (ret != 0 || num_rr == 0) {
    getdns_dict_destroy(response);
    return (ret != 0) ? ret : dns_ip(ia, sa); /* e.g., CNAME -> A */
  }

  struct dns_mxip_user mx = {
    .num = 0,
    .mx  = (struct mx_user*)alloc(num_rr * sizeof(struct mx_user)),
  };
  if (!mx.mx)
    return DNS_MEM;
  ret = getdns_resolve_loop(response, GETDNS_RRTYPE_MX, dns_mxip_cb, &mx);
  if (ret < 0)
  {
    while (mx.num > 0)
      alloc_free(mx.mx[--mx.num].sa.s);
    alloc_free(mx.mx);
    getdns_dict_destroy(response);
    return ret;
  }
  getdns_dict_destroy(response);

  int flagsoft = 0;
  while (mx.num > 0)
  {
    int i, j;
    unsigned long numsame = 1;

    i = 0;
    for (j = 1; j < mx.num; ++j) {
      if (mx.mx[j].p < mx.mx[i].p)
      {
        i = j;
        numsame = 1;
      }
      else if (mx.mx[j].p == mx.mx[i].p)
      {
        ++numsame;
        random = random * 69069 + 1;
        if ((random / 2) < (2147483647 / numsame))
          i = j;
      }
    }

    switch(dns_ipplus(ia, &mx.mx[i].sa, mx.mx[i].p))
    {
      case DNS_MEM:
      case DNS_SOFT:
        flagsoft = 1;
        break;
    }

    alloc_free(mx.mx[i].sa.s);
    mx.mx[i] = mx.mx[--mx.num];
  }

  alloc_free(mx.mx);
  return flagsoft;
}