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

#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>

#include <sys/poll.h>

#include <fcntl.h>

#include <unistd.h>

#include <errno.h>

#include <sys/time.h>

#include "dns2.h"
#include "utils.h"
#include "packet.h"
#include "constants.h"
#include "glist.h"
#include "rr.h"
#include "question.h"

#include "dns2-private.h"

typedef struct {
	char name[512];
	
	struct sockaddr_storage server_addr;
	socklen_t addr_size;
	
	int socket;
} NameServer;


/* Prototipos de función */
int _dns2_internal_tcp_start_connect (DNS2 *obj);
int _dns2_internal_udp_create (DNS2 *obj);

int _dns2_internal_tcp_connect_continue (DNS2 *obj);
int _dns2_internal_udp_send (DNS2 *obj);
int _dns2_internal_tcp_send (DNS2 *obj);
int _dns2_internal_udp_recv (DNS2 *obj);
int _dns2_internal_tcp_recv (DNS2 *obj);
int _dns2_internal_tcp_next_or_error (DNS2 *obj);
void dns2_clean_network (DNS2 *obj);

double _dns2_diff_timespec (const struct timespec *time1, const struct timespec *time0) {
  return (time1->tv_sec - time0->tv_sec)
      + (time1->tv_nsec - time0->tv_nsec) / 1000000000.0;
}

int dns2_async (DNS2 *obj) {
	int length;
	// checkServers
	DNSRR *signature;
	DNSQuestion *question;
	int max_udp_size;
	
	if (obj->state != DNS_STATE_FREE) {
		return DNS_COMMAND_ERROR;
	}
	
	if (obj->packet == NULL) {
		return DNS_COMMAND_ERROR;
	}
	
	/* Resetear el estado del paquete */
	dns_packet_reset (obj->packet);
	
	obj->current_ns = g_list_first (obj->nameservers);
	obj->retry_count = 0;
	if (obj->current_ns == NULL) {
		return DNS_COMMAND_ERROR;
	}
	
	/* Actualizar el contador de paquetes */
	dns_packet_recount_items (obj->packet);
	
	if (obj->type == DNS_UPDATER && (obj->packet->header.qdcount == 0 || obj->packet->header.nscount == 0)) {
		
		/* Error, nada que mandar en el update */
		return DNS_COMMAND_ERROR;
	}
	
	max_udp_size = DNS_MAX_UDP_SIZE;
	
	/* Si tiene una firma, generar la firma */
	if (obj->auth_signature != NULL) {
		if (obj->auth_signature->type == RR_TYPE_TSIG /*||
		    obj->auth_signature->tipo == RR_TYPE_SIG*/) {
		    
		    signature = dns2_rr_dup (obj->auth_signature);
		    dns_rr_pack_generate_signature (signature, obj->packet);
			obj->packet->additional = g_list_append (obj->packet->additional, signature);
			obj->packet->header.arcount = g_list_length (obj->packet->additional);
		}
		max_udp_size = 4000;
	}
	
	dns_packet_pack_data (obj->packet);
	
	/* ------ Para depuración ------
	unsigned char *data;
	
	length = dns_packet_get_data (obj->packet, &data);
	for (int g = 0; g<length; g++) {
		printf ("%02x", (unsigned int) data[g]);
	}
	printf ("\nEnd data.\n");
	 ------ FIN de la depuración ------ */
	
	/* Revisar el tipo de pregunta, algunos tipos de pregunta, como los ANY requieren TCP */
	question = (DNSQuestion *) obj->packet->questions->data;
	
	if (question->qtype == RR_TYPE_ANY || question->qtype == RR_TYPE_AXFR) {
		obj->force_use_tcp = 1;
	}
	
	length = dns_packet_get_data (obj->packet, NULL);
	if (length > max_udp_size || obj->force_use_tcp == 1) {
		obj->state = DNS_STATE_TCP_CONNECT;
		
		return _dns2_internal_tcp_start_connect (obj);
	} else {
		obj->state = DNS_STATE_UDP;
		
		return _dns2_internal_udp_create (obj);
	}
}

int dns2_async_continue (DNS2 *obj, int happen) {
	if (obj->state == DNS_STATE_FREE) {
		return DNS_COMMAND_ERROR;
	}
	
	if (obj->direction == DNS_POLL_WANT_READ) {
		if ((happen & DNS_POLL_WANT_READ) == 0) {
			return DNS_COMMAND_POLL;
		}
		
		switch (obj->state) {
			case DNS_STATE_UDP:
				return _dns2_internal_udp_recv (obj);
			case DNS_STATE_TCP:
				return _dns2_internal_tcp_recv (obj);
		}
	} else if (obj->direction == DNS_POLL_WANT_WRITE) {
		if ((happen & DNS_POLL_WANT_WRITE) == 0) {
			return DNS_COMMAND_POLL;
		}
		
		switch (obj->state) {
			case DNS_STATE_TCP_CONNECT:
				return _dns2_internal_tcp_connect_continue (obj);
			case DNS_STATE_UDP:
				return _dns2_internal_udp_send (obj);
			case DNS_STATE_TCP:
				return _dns2_internal_tcp_send (obj);
		}
	}
	
	return DNS_COMMAND_ERROR;
}

int dns2_sync (DNS2 *obj) {
	int res, res_poll;
	struct pollfd poller;
	int wanted_events;
	int events;
	
	res = dns2_async (obj);
	
	while (res == DNS_COMMAND_POLL) {
		poller.fd = dns2_get_socket (obj);
		poller.revents = 0;
		poller.events = 0;
		/* FIXME AQUI */
		wanted_events = dns2_get_socket_poll_flags (obj);
		if (wanted_events == DNS_POLL_WANT_READ) {
			poller.events = POLLIN;
		} else if (wanted_events == DNS_POLL_WANT_WRITE) {
			poller.events = POLLOUT;
		}
		
		res_poll = poll (&poller, 1, 30);
		
		if (res_poll == 0) {
			/* Si no ocurrió evento, mandar a llamar por timeouts: */
			res = dns2_check_timeout (obj);
			continue;
		} else if (res_poll < 0) {
			if (errno == EINTR) continue;
			/* Caso contrario, salir */
			res = DNS_COMMAND_ERROR;
			break;
		}
		
		/* Ocurrió un evento, revisar cuál se quería y retornarlo para el continue */
		if (poller.revents & POLLOUT) {
			events = DNS_POLL_WANT_WRITE;
		} else if (poller.revents & POLLIN) {
			events = DNS_POLL_WANT_READ;
		}
		res = dns2_async_continue (obj, events);
	}
	
	/* Leer el paquete respuesta y regresarlo */
	return res;
}

int _dns2_internal_tcp_start_connect (DNS2 *obj) {
	NameServer *ns;
	int flags;
	int res;
	
	ns = (NameServer *) obj->current_ns->data;
	ns->socket = socket (ns->server_addr.ss_family, SOCK_STREAM, 0);
	
	/* Si no pude crear un socket, ni siquiera intentar con el siguiente nameserver */
	if (ns->socket < 0) {
		obj->state = DNS_STATE_FREE;
		obj->direction = DNS_POLL_WANT_NONE;
		
		return DNS_COMMAND_ERROR;
	}
	
	/* Aplicar el no-bloqueante */
	flags = fcntl (ns->socket, F_GETFL, 0);
	flags = flags | O_NONBLOCK;
	fcntl (ns->socket, F_SETFL, flags);
	
	/* Ejecutar el connect */
	res = connect (ns->socket, (struct sockaddr *) &ns->server_addr, ns->addr_size);
	
	if (res < 0) {
		if (errno == EINPROGRESS) {
			obj->direction = DNS_POLL_WANT_WRITE;
			return DNS_COMMAND_POLL;
		}
		
		return _dns2_internal_tcp_next_or_error (obj);
	}
	
	/* Tengo conexión, pasar de estado */
	obj->state = DNS_STATE_TCP;
	obj->direction = DNS_POLL_WANT_WRITE;
	
	return dns2_async_continue (obj, DNS_POLL_WANT_WRITE);
}

int _dns2_internal_tcp_connect_continue (DNS2 *obj) {
	int optval;
	socklen_t optlen;
	int res;
	NameServer *ns;
	
	ns = (NameServer *) obj->current_ns->data;
	optlen = sizeof (optval);
	
	res = getsockopt (ns->socket, SOL_SOCKET, SO_ERROR, &optval, &optlen);
	
	if (res < 0 || optval != 0) {
		return _dns2_internal_tcp_next_or_error (obj);
	}
	
	/* Pasar de estado */
	obj->state = DNS_STATE_TCP;
	obj->direction = DNS_POLL_WANT_WRITE;
	
	return dns2_async_continue (obj, DNS_POLL_WANT_WRITE);
}

int _dns2_internal_tcp_next_or_error (DNS2 *obj) {
	NameServer *ns;
	
	ns = (NameServer *) obj->current_ns->data;
	close (ns->socket);
	ns->socket = -1;
	
	obj->current_ns = g_list_next (obj->current_ns);
	if (obj->current_ns != NULL) {
		obj->state = DNS_STATE_TCP_CONNECT;
		
		return _dns2_internal_tcp_start_connect (obj);
	} else {
		obj->state = DNS_STATE_FREE;
		obj->direction = DNS_POLL_WANT_NONE;
		
		/* Como ya no hay más nameservers en la lista, entregar error */
		return DNS_COMMAND_ERROR;
	}
}

int _dns2_internal_udp_next_or_error (DNS2 *obj) {
	/* Los sockets UDP no se cierran, se dejan para usar después */
	
	obj->current_ns = g_list_next (obj->current_ns);
	if (obj->current_ns != NULL) {
		/* Aún hay siguiente servidor, intentarlo con éste */
		obj->state = DNS_STATE_UDP;
		
		return _dns2_internal_udp_create (obj);
	} else {
		/* Como ya no hay mas en la lista, esperar al timeout */
		obj->direction = DNS_POLL_WANT_READ;
		
		return DNS_COMMAND_POLL;
	}
}

int _dns2_internal_udp_send (DNS2 *obj) {
	int res;
	void *data;
	int data_length;
	NameServer *ns;
	
	ns = (NameServer *) obj->current_ns->data;
	
	data_length = dns_packet_get_data (obj->packet, (unsigned char **) &data);
	
	clock_gettime (CLOCK_MONOTONIC, &obj->last_time_sent);
	
	res = sendto (ns->socket, data, data_length, 0, (struct sockaddr *) &ns->server_addr, ns->addr_size);
	
	if (res < 0) {
		if (errno == EWOULDBLOCK || errno == EAGAIN || errno == EINTR) {
			obj->direction = DNS_POLL_WANT_WRITE;
			
			return DNS_COMMAND_POLL;
		}
		
		return _dns2_internal_udp_next_or_error (obj);
	} else if (res < data_length) {
		/* Envié la petición incompleta, manejar este error */
		
		return _dns2_internal_udp_next_or_error (obj);
	}
	
	/* Cambio de estado, esperar la lectura */
	obj->direction = DNS_POLL_WANT_READ;
	
	return dns2_async_continue (obj, DNS_POLL_WANT_READ);
}

int _dns2_internal_tcp_send (DNS2 *obj) {
	int res;
	void *data;
	int data_length;
	unsigned char data_plus[8192 + 2];
	uint16_t t16;
	NameServer *ns;
	
	ns = (NameServer *) obj->current_ns->data;
	data_length = dns_packet_get_data (obj->packet, (unsigned char **) &data);
	
	/* En TCP, agregar la longitud previa a los datos */
	t16 = htons (data_length);
	memcpy (data_plus, &t16, 2);
	
	memcpy (&data_plus[2], data, data_length);
	
	res = write (ns->socket, data_plus, data_length + 2);
	
	if (res < 0) {
		if (errno == EWOULDBLOCK || errno == EAGAIN || errno == EINTR) {
			obj->direction = DNS_POLL_WANT_WRITE;
			
			return DNS_COMMAND_POLL;
		}
		
		return _dns2_internal_tcp_next_or_error (obj);
	} else if (res < data_length) {
		/* Envié la petición incompleta, manejar este error */
		
		return _dns2_internal_tcp_next_or_error (obj);
	}
	
	/* Cambio de estado, esperar la lectura */
	obj->direction = DNS_POLL_WANT_READ;
	
	return dns2_async_continue (obj, DNS_POLL_WANT_READ);
}

int _dns2_internal_udp_recv (DNS2 *obj) {
	int res;
	struct sockaddr_storage origen;
	socklen_t origen_size;
	unsigned char buffer[8192];
	NameServer *ns;
	
	ns = (NameServer *) obj->current_ns->data;
	
	origen_size = sizeof (origen);
	res = recvfrom (ns->socket, buffer, sizeof (buffer), 0, (struct sockaddr *) &origen, &origen_size);
	
	if (res < 0) {
		if (errno == EWOULDBLOCK || errno == EAGAIN || errno == EINTR) {
			obj->direction = DNS_POLL_WANT_READ;
			
			return DNS_COMMAND_POLL;
		}
		
		/* ¿Otro error en un socket UDP? -> Mejor seguir esperando una lectura */
		
		obj->direction = DNS_POLL_WANT_READ;
			
		return DNS_COMMAND_POLL;
	} /* Revisar que el paquete venga del mismo origen a donde envié la petición original */
	else if (sockaddr_cmp ((struct sockaddr *) &origen, (struct sockaddr *) &ns->server_addr) == 0) {
		/* Por el momento, dumpear la información
		int g;
		printf ("Respuesta recibida:\n");
		for (g = 0; g < res; g++) {
			printf ("%02x", ((unsigned int) buffer[g]));
		}
		
		printf ("\nEnd response\n");*/
		
		if (obj->packet_response != NULL) {
			dns_packet_free_full (obj->packet_response);
			obj->packet_response = NULL;
		}
		
		obj->packet_response = dns_packet_create_response (SOCK_DGRAM, buffer, res);
		
		if (obj->packet_response == NULL) {
			/* Como no es válida la respuesta, pasar al siguiente servidor */
			return _dns2_internal_udp_next_or_error (obj);
		} else {
			if (obj->packet_response->header.id != obj->packet->header.id) {
				/* Esta no es la respuesta que esperaba */
				obj->direction = DNS_POLL_WANT_READ;
			
				return DNS_COMMAND_POLL;
			}
			/* Finalizar la petición */
			dns2_clean_network (obj);
			
			obj->state = DNS_STATE_FREE;
			obj->direction = DNS_POLL_WANT_NONE;
			
			return DNS_COMMAND_SUCCESS;
		}
	}
	obj->direction = DNS_POLL_WANT_READ;
			
	return DNS_COMMAND_POLL;
}

int _dns2_internal_tcp_recv (DNS2 *obj) {
	int res;
	uint16_t t16;
	NameServer *ns;
	
	ns = (NameServer *) obj->current_ns->data;
	
	res = read (ns->socket, &obj->buffer_read[obj->buffer_read_pos], sizeof (obj->buffer_read) - obj->buffer_read_pos);
	
	if (res < 0) {
		if (errno == EWOULDBLOCK || errno == EAGAIN || errno == EINTR) {
			obj->direction = DNS_POLL_WANT_READ;
			
			return DNS_COMMAND_POLL;
		}
		
		/* Si obtuvimos un error, pasar al siguiente servidor */
		return _dns2_internal_tcp_next_or_error (obj);
	} else if (res == 0) {
		/* ¿Socket cerrado? -> Pasar al siguiente servidor */
		return _dns2_internal_tcp_next_or_error (obj);
	} else {
		/* Por el momento, dumpear la información
		int g;
		printf ("Bytes recibidos por TCP:\n");
		for (g = obj->buffer_read_pos; g < obj->buffer_read_pos + res; g++) {
			printf ("%02x", ((unsigned int) obj->buffer_read[g]));
		}
		
		printf ("\nFin bytes\n");*/
		
		/* Incrementar el buffer */
		obj->buffer_read_pos = obj->buffer_read_pos + res;
		
		/* Validar la longitud de los datos TCP aquí */
		memcpy (&t16, obj->buffer_read, 2);
		t16 = ntohs (t16);
		
		if (t16 + 2 < obj->buffer_read_pos) {
			/* Datos incompletos TCP, seguir haciendo POLL */
			obj->direction = DNS_POLL_WANT_READ;
			
			return DNS_COMMAND_POLL;
		}
		
		if (obj->packet_response != NULL) {
			dns_packet_free_full (obj->packet_response);
			obj->packet_response = NULL;
		}
		
		obj->packet_response = dns_packet_create_response (SOCK_STREAM, &obj->buffer_read[2], res - 2);
		
		if (obj->packet_response == NULL) {
			/* Error de parseo */
			return _dns2_internal_tcp_next_or_error (obj);
		} else {
			/* If packet == good */
			//dns2_updater_packet_has_valid_tsig (obj, packet);
			
			/* Pasar de estado */
			close (ns->socket);
			ns->socket = -1;
			
			obj->state = DNS_STATE_FREE;
			obj->direction = DNS_POLL_WANT_NONE;
			
			return DNS_COMMAND_SUCCESS;
		}
	}
}

int _dns2_internal_udp_create (DNS2 *obj) {
	NameServer *ns;
	int flags;
	
	obj->direction = DNS_POLL_WANT_WRITE;
	
	ns = (NameServer *) obj->current_ns->data;
	
	/* Crear el socket UDP de una buena vez */
	if (ns->socket < 0) {
		ns->socket = socket (ns->server_addr.ss_family, SOCK_DGRAM, 0);
		
		/* Aplicar el no-bloqueante */
		flags = fcntl (ns->socket, F_GETFL, 0);
		flags = flags | O_NONBLOCK;
		fcntl (ns->socket, F_SETFL, flags);
	}
	
	/* Si no pude crear un socket, tenemos un serio problema */
	if (ns->socket < 0) {
		obj->state = DNS_STATE_FREE;
		obj->direction = DNS_POLL_WANT_NONE;
		
		return DNS_COMMAND_ERROR;
	}
	
	return dns2_async_continue (obj, DNS_POLL_WANT_WRITE);
}

void dns2_sign_tsig (DNS2 *obj, const char *keyname, const char *signature, const char *algorithm) {
	DNSRR *rr;
	
	/* TODO: Revisar si deberíamos liberar el auth_signature anterior */
	rr = dns2_rr_create_tsig (keyname, signature, algorithm);
	
	if (rr != NULL) {
		obj->auth_signature = rr;
	}
}

/* Crear un objeto */
DNS2 *dns2_new (void) {
	DNS2 *obj;
	
	obj = (DNS2 *) malloc (sizeof (DNS2));
	
	if (obj == NULL) {
		return NULL;
	}
	
	obj->packet = NULL;
	
	obj->state = DNS_STATE_FREE;
	obj->direction = DNS_POLL_WANT_NONE;
	obj->nameservers = NULL;
	obj->timeout = 5;
	
	obj->current_ns = NULL;
	obj->zone = NULL;
	obj->force_use_tcp = 0;
	obj->want_recursion = 0;
	obj->buffer_read_pos = 0;
	obj->auth_signature = NULL;
	obj->packet_response = NULL;
	
	return obj;
}

void dns2_add_nserver (DNS2 *obj, const char *ip) {
	int res;
	struct sockaddr_in v4;
	struct sockaddr_in6 v6;
	NameServer *ns;
	
	memset (&v4, 0, sizeof (v4));
	res = inet_pton (AF_INET, ip, &v4.sin_addr);
	if (res > 0) {
		/* Es una IP de 4 */
		ns = (NameServer *) malloc (sizeof (NameServer));
		
		if (ns == NULL) return;
		
		v4.sin_family = AF_INET;
		v4.sin_port = htons (53);
		
		memcpy (&ns->server_addr, &v4, sizeof (v4));
		ns->addr_size = sizeof (v4);
		strncpy (ns->name, ip, sizeof (ns->name));
		ns->socket = -1;
		
		obj->nameservers = g_list_append (obj->nameservers, ns);
		
		return;
	}
	
	memset (&v6, 0, sizeof (v6));
	res = inet_pton (AF_INET6, ip, &v6.sin6_addr);
	if (res > 0) {
		ns = (NameServer *) malloc (sizeof (NameServer));
		
		if (ns == NULL) return;
		
		v6.sin6_family = AF_INET6;
		v6.sin6_port = htons (53);
		
		memcpy (&ns->server_addr, &v6, sizeof (v6));
		ns->addr_size = sizeof (v6);
		strncpy (ns->name, ip, sizeof (ns->name));
		
		obj->nameservers = g_list_append (obj->nameservers, ns);
		
		return;
	}
}

void dns2_clean_network (DNS2 *obj) {
	GList *g;
	NameServer *ns;
	
	for (g = obj->nameservers; g != NULL; g = g->next) {
		ns = (NameServer *) g->data;
		
		if (ns->socket >= 0) {
			close (ns->socket);
			ns->socket = -1;
		}
	}
}

void dns2_free (DNS2 *obj) {
	
	/* Cerrar el socket, si es que hay alguno abierto */
	dns2_clean_network (obj);
	
	/* Liberar la lista de NS */
	g_list_free_full (obj->nameservers, (GDestroyNotify) free);
	
	if (obj->zone != NULL) {
		free (obj->zone);
	}
	
	/* Ahora, buscar si el auth_signature ya está en la lista de additionals. Si lo está, no hay necesidad de liberarlo, se liberará cuando se libere el paquete */
	if (obj->auth_signature != NULL) {
		/* Liberar la firma, porque está fuera del paquete */
		dns2_rr_free (obj->auth_signature);
		obj->auth_signature = NULL;
	}
	
	if (obj->packet != NULL) {
		dns_packet_free_full (obj->packet);
	}
	
	if (obj->packet_response != NULL) {
		dns_packet_free_full (obj->packet_response);
	}
	
	free (obj);
}

int dns2_check_timeout (DNS2 *obj) {
	struct timespec now;
	double d;
	
	if (obj->state != DNS_STATE_UDP) return DNS_COMMAND_POLL;
	
	clock_gettime (CLOCK_MONOTONIC, &now);
	
	d = _dns2_diff_timespec (&now, &obj->last_time_sent);
	if (obj->current_ns->next == NULL) {
		/* Si estamos al borde del último servidor, necesitamos el timeout largo */
		if (d > (double) obj->timeout) {
			/* Como cayó el timeout largo, incrementar los retries y regresar al primer servidor */
			obj->retry_count++;
			if (obj->retry_count >= 3) {
				/* Ya agoté todos los reintentos y servidores, bye */
				dns2_clean_network (obj);
				obj->current_ns = NULL;
				
				return DNS_COMMAND_ERROR;
			}
			
			obj->current_ns = obj->nameservers;
			return _dns2_internal_udp_create (obj);
		}
	} else if (d > 1.0) {
		/* 1 segundo entre servidor y servidor */
		return _dns2_internal_udp_next_or_error (obj);
	}
	
	return DNS_COMMAND_POLL;
}

int dns2_get_socket (DNS2 *obj) {
	NameServer *ns;
	
	if (obj == NULL) return -1;
	
	if (obj->state == DNS_STATE_FREE) return -1;
	
	ns = (NameServer *) obj->current_ns->data;
	
	return ns->socket;
}

int dns2_get_socket_poll_flags (DNS2 *obj) {
	if (obj == NULL) return -1;
	
	if (obj->direction == DNS_POLL_WANT_READ) {
		return DNS_POLL_WANT_READ;
	} else if (obj->direction == DNS_POLL_WANT_WRITE) {
		return DNS_POLL_WANT_WRITE;
	}
	
	return 0;
}

void dns2_set_use_tcp (DNS2 *obj) {
	if (obj == NULL) return;
	
	obj->force_use_tcp = 1;
}

DNSRR **dns2_get_answers_from_response (DNS2 *obj) {
	if (obj == NULL) return NULL;
	if (obj->packet_response == NULL) return NULL;
	
	return dns_packet_get_answers (obj->packet_response);
}

DNSRR **dns2_get_authority_from_response (DNS2 *obj) {
	if (obj == NULL) return NULL;
	if (obj->packet_response == NULL) return NULL;
	
	return dns_packet_get_authority (obj->packet_response);
}

DNSRR **dns2_get_additional_from_response (DNS2 *obj) {
	if (obj == NULL) return NULL;
	if (obj->packet_response == NULL) return NULL;
	
	return dns_packet_get_additional (obj->packet_response);
}

int dns2_get_rcode (DNS2 *obj) {
	if (obj == NULL) return -1;
	if (obj->packet_response == NULL) return -1;
	
	return obj->packet_response->header.rcode;
}