#include "./handler.h"
#include "./server.h"
#include "./net.h"
#include "logger.h"

#include <algorithm>
#include <iostream>
#include <teadns/parser.h>
#include <teadns/builder.h>
#include <netinet/in.h>

using namespace ts::dns;
using namespace ts::dns::builder;

std::vector<std::string> le_token;

struct ClientAddress {
	std::string str;
	const sockaddr_storage &addr;
};

void create_answer(DNSBuilder& response, const ClientAddress &client_address, const parser::DNSQuery& query);
void WebDNSHandler::handle_message(const std::shared_ptr<DNSServerBinding>& binding, const sockaddr_storage &address, void *buffer, size_t size) {
	const ClientAddress addr_data{net::to_string(address), address};
	auto logger = log_dns();
	const auto begin = std::chrono::system_clock::now();

	DNSParser parser{0, nullptr, buffer, size};

	std::string error;
	if(!parser.parse(error)) {
		logger->warn("[{}] Received invalid DNS request: {}", addr_data.str, error);
		return;
	}
	logger->info("[{}] Received DNS request with {} queries", addr_data.str, parser.queries().size());
	logger->trace("[{}] Query type: {}", addr_data.str, (uint32_t) parser.header().query_type());
	logger->trace("[{}] Queries ({}):", addr_data.str, (uint32_t) parser.queries().size());
	for(auto& query : parser.queries())
		logger->trace("[{}]   {} ({}::{})", addr_data.str, query->qname(), rrclass::name(query->qclass()), rrtype::name(query->qtype()));

	{
		DNSBuilder response{};

		response.header().id(*(uint16_t*) buffer);
		response.header().set_answer(true);
		response.header().set_query_type(parser.header().query_type());

		for(auto& query : parser.queries()) {
			auto& q = response.push_query();
			q.set_qclass(query->qclass());
			q.set_qname(query->qname());
			q.set_qtype(query->qtype());

			try {
				create_answer(response, addr_data, *query);
			} catch(std::exception& ex) {
				logger->error("[{}] Failed to generate answer for query. Caught an exception: {}", addr_data.str, ex.what());
			}
		}

		if(response.answer_count() == 0)
			response.header().set_response_code(rcode::NXDOMAIN);
		else
			response.header().set_response_code(rcode::NOERROR);

		char rbuffer[1024];
		auto len = response.build(rbuffer, 1024, error);
		if(!len) {
			//TODO: Send prebuild server fail packet
			logger->error("[{}] Failed to build response: {}", error);
		} else {
			binding->send(address, rbuffer, len);
		}
	}

	const auto end = std::chrono::system_clock::now();
	logger->info("[{}] Query completed in {}us", addr_data.str, std::chrono::floor<std::chrono::microseconds>(end - begin).count());
}

#define MAX_IPV6_ADDRESS_STR_LEN  39

//Address at least one char long!
bool parse_v6(uint8_t(result)[16], const std::string_view& address) {
	uint16_t accumulator{0};
	uint8_t colon_count{0}, pos{0};

	memset(result, 0, 16);

	// Step 1: look for position of ::, and count colons after it
	for(uint8_t i = 1; i <= MAX_IPV6_ADDRESS_STR_LEN && address.length() > i; i++) {
		if (address[i] == ':') {
			if (address[i - 1] == ':') {
				colon_count = 14; // Double colon!
			} else if (colon_count) {
				colon_count -= 2; // Count backwards the number of colons after the ::
			}
		}
	}

	// Step 2: convert from ascii to binary
	for(uint8_t i=0; i <= MAX_IPV6_ADDRESS_STR_LEN && pos < 16; i++) {
		if (address[i] == ':' || address.length() == i) {
			result[pos] = accumulator >> 8U;
			result[pos + 1] = accumulator;
			accumulator = 0;

			if (colon_count && i && address[i - 1] == ':') {
				pos = colon_count;
			} else {
				pos += 2;
			}
		} else {
			(uint8_t&) address[i] |= 0x20U;
			accumulator <<= 4U;

			if (address[i] >= '0' && address[i] <= '9') {
				accumulator |= (uint8_t) (address[i] - '0');
			} else if (address[i] >= 'a' && address[i] <= 'f') {
				accumulator |= (uint8_t) ((address[i] - 'a') + 10);
			} else {
				return false; // Not hex or colon: fail
			}
		}
	}

	return true;
}


void create_answer(DNSBuilder& response, const ClientAddress & client_address, const parser::DNSQuery& query) {
	if(query.qclass() != rrclass::IN) return;

	if(query.qtype() == rrtype::A) {
		auto dn = query.qname();
		uint8_t resp[4];
		{
			size_t index = 0;
			size_t aindex = 0;
			do {
				auto found = dn.find('-', index);
				auto length = index == -1 ? dn.length() - index : found - index;

				try {
					resp[aindex] = std::stoul(dn.substr(index, length));
				} catch(std::exception& ex) {
					break;
				}
				aindex++;
				index = found;
			} while(index++ && aindex < 4);
			if(aindex != 4)
				return;
		}

		log_dns()->info("[{}] Sending requested IPv4 ({}): {}.{}.{}.{}", client_address.str, dn, resp[0], resp[1], resp[2], resp[3]);
		auto& a = response.push_answer(query.qname());
		a.set_class(query.qclass());
		a.set_type(rrtype::A);
		a.set_ttl(120);
		a.builder<rrbuilder::A>().set_address(resp);
	} else if(query.qtype() == rrtype::AAAA) {
		auto dn = query.qname();
		const auto parts = parse_dn(dn);

		if(parts.size() != 3) return;

		in6_addr result{};
		if(!parse_v6(result.__in6_u.__u6_addr8, parts[0])) return;

		log_dns()->info("[{}] Sending requested IPv6 ({}): {}", client_address.str, dn, net::to_string(result));

		auto& a = response.push_answer(query.qname());
		a.set_class(query.qclass());
		a.set_type(rrtype::AAAA);
		a.set_ttl(120);
		a.builder<rrbuilder::AAAA>().set_address(result.__in6_u.__u6_addr32);
	} else if(query.qtype() == rrtype::TXT) {
		auto dn = query.qname();
		std::transform(dn.begin(), dn.end(), dn.begin(), tolower);

		if(dn == "_acme-challenge.con-gate.work") {
			log_dns()->info("[{}] Received LetsEncrypt auth request. Sending {} predefined keys.", le_token.size());

			for(auto& key : le_token) {
				auto& a = response.push_answer(query.qname());
				a.set_class(query.qclass());
				a.set_type(query.qtype());
				a.set_ttl(120);
				a.builder<rrbuilder::TXT>().set_text(key);
			}
		}
	}
}