#include <arpa/inet.h>
#include <assert.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#include "dns.h"
#include "port_utils.h"
#include "utils.h"

#define BUFFER_MAX_SIZE 4096

const char *HOSTNAME = "127.0.0.1";

ssize_t
query_dns (unsigned char *query, size_t size, unsigned char **response)
{

  // hardcoded for dukens server
  const char *port_str = get_port (); // same port dukens listens on

  int port = strtol (port_str, NULL, 10);
  if (port <= 0 || port > 65535)
    {
      fprintf (stderr, "Invalid port: %s\n", port_str);
      exit (EXIT_FAILURE);
    }

  struct sockaddr_in addr = { 0 };
  addr.sin_family = AF_INET;
  addr.sin_port = htons (port);
  if (inet_aton (HOSTNAME, &addr.sin_addr) != 1)
    {
      fprintf (stderr, "Invalid IPv4 address: %s\n", HOSTNAME);
      exit (EXIT_FAILURE);
    }

  // create socket
  int sockfd = socket (AF_INET, SOCK_DGRAM, 0);
  if (sockfd < 0)
    {
      perror ("socket");
      exit (EXIT_FAILURE);
    }

  ssize_t sent = sendto (sockfd, query, size, 0, (struct sockaddr *)&addr,
                         sizeof (addr));
  if (sent < 0)
    {
      perror ("sendto");
      close (sockfd);
      exit (EXIT_FAILURE);
    }

  // receive and print response
  *response = calloc (BUFFER_MAX_SIZE, sizeof (unsigned char));
  ssize_t received
      = recvfrom (sockfd, *response, BUFFER_MAX_SIZE, 0, NULL, NULL);
  if (received < 0)
    {
      perror ("recvfrom");
      close (sockfd);
      exit (EXIT_FAILURE);
    }

  // cleanup
  close (sockfd);

  return received;
}

void
phase_one ()
{
  // hard coded DNS request
  unsigned char query[17]
      = { 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
          0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01 };

  unsigned char *res_bytes;
  ssize_t received = query_dns (query, sizeof (query), &res_bytes);

  printf ("Sending %zu bytes to %s:\n", sizeof (query), HOSTNAME);
  print_byte_block (query, sizeof (query));

  printf ("\nReceived %zd bytes from %s:\n", received, HOSTNAME);
  print_byte_block (res_bytes, received);
  printf ("\n");

  dns_response_t res = { 0 };
  parse_dns_response (res_bytes, sizeof (query), &res);
  dns_answer_t answer = res.answers[0];

  printf ("Domain name: %s\n", answer.domain);
  printf ("Record type: %s %s\n", qclass_to_str (answer.qclass),
          record_to_str (answer.qtype));
  printf ("TTL: %d\n", answer.ttl);
  printf ("IPv4 address: %s\n", answer.rdata);

  free (res_bytes);
}

void
not_phase_one (int argc, char *argv[])
{
  // building query
  uint16_t xid = strtol (argv[2], NULL, 10);

  uint16_t record = get_record_from_str (argv[4]);

  char domain_name[BUFFER_LENGTH];

  if (record == PTR)
    {
      ptr_ip (domain_name, BUFFER_LENGTH, argv[3]);
    }
  else
    {
      strncpy (domain_name, argv[3], sizeof (domain_name) - 1);
      domain_name[sizeof (domain_name) - 1] = '\0';
    }

  unsigned char query_bytes[512];
  size_t offset = 0;

  // DNS Header
  query_bytes[offset++] = (xid >> 8) & 0xFF;
  query_bytes[offset++] = xid & 0xFF;
  query_bytes[offset++] = 0x01; // flags
  query_bytes[offset++] = 0x00;
  query_bytes[offset++] = 0x00; // QDCOUNT high
  query_bytes[offset++] = 0x01; // QDCOUNT low
  // rest zeros (ANCOUNT, NSCOUNT, ARCOUNT)
  for (int i = 0; i < 6; i++)
    query_bytes[offset++] = 0x00;

  // Question section: domain name
  char *token = strtok (domain_name, ".");
  while (token)
    {
      size_t len = strlen (token);
      query_bytes[offset++] = (unsigned char)len;
      memcpy (&query_bytes[offset], token, len);
      offset += len;
      token = strtok (NULL, ".");
    }
  query_bytes[offset++] = 0x00; // end of domain

  query_bytes[offset++] = record >> 8;
  query_bytes[offset++] = record & 0xff;
  query_bytes[offset++] = 0x00; // QCLASS = IN
  query_bytes[offset++] = 0x01;

  unsigned char *res_bytes;
  query_dns (query_bytes, offset, &res_bytes);

  dns_query_t query = { 0 };
  parse_dns_query (&query_bytes[0], &query);
  dns_response_t response = { 0 };
  parse_dns_response (&res_bytes[0], offset, &response);

  // print header
  print_dns_header (&response.query.header);

  // print question section
  print_question_section (&query);

  uint8_t status = (response.query.header.flags >> 0) & 0xF;

  if (status == 0 && response.query.header.ancount > 0)
    {
      printf ("\n;; ANSWER SECTION:\n");

      for (size_t i = 0; i < response.query.header.ancount; i++)
        {
          dns_answer_t answer = response.answers[i];
          print_answer_section (&answer);
        }
    }

  if (response.query.header.nscount > 0)
    {
      printf ("\n;; AUTHORITY SECTION:\n");

      for (size_t i = 0; i < response.query.header.nscount; i++)
        {
          dns_answer_t authority = response.authorities[i];
          print_answer_section (&authority);
        }
    }

  if (response.query.header.arcount > 0)
    {
      printf ("\n;; ADDITIONAL SECTION:\n");

      for (size_t i = 0; i < response.query.header.arcount; i++)
        {
          dns_answer_t additional = response.additionals[i];
          print_answer_section (&additional);
        }
    }

  free (res_bytes);
}

static bool
get_args (int argc, char **argv, bool *opendns, bool *xflag)
{
  int ch = 0;
  while ((ch = getopt (argc, argv, "ox")) != -1)
    {
      switch (ch)
        {
        case 'o':
          *opendns = true;
          break;
        case 'x':
          *xflag = true;
          break;
        default:
          return false;
        }
    }

  return true;
}

int
main (int argc, char *argv[])
{
  bool opendns = false;
  bool xflag = false;

  if (!get_args (argc, argv, &opendns, &xflag))
    {
      fprintf (stderr, "ERROR: Show usage\n");
      return EXIT_FAILURE;
    }

  if (!xflag)
    {
      phase_one ();
    }
  else
    {
      not_phase_one (argc, argv);
    }

  return EXIT_SUCCESS;
}