#include "teadns/parser.h"

#include <iostream>
#include <codecvt>
#include <algorithm>

#ifdef WIN32
    #include <WS2tcpip.h>
    #include <in6addr.h>
    #include <ip2string.h>
    #include <inaddr.h>
    #include <WinDNS.h>
#else
	#include <arpa/inet.h>
	#include <netinet/in.h>
#include <cstring>

#endif

using namespace tc::dns;
using namespace tc::dns::response;
using namespace tc::dns::response::rrparser;

thread_local std::vector<size_t> visited_links;
std::string DNSResponseData::parse_dns_dn(std::string &error, size_t &index, bool allow_compression) {
	if(allow_compression) {
		visited_links.clear();
		visited_links.reserve(8);

		if(std::find(visited_links.begin(), visited_links.end(), index) != visited_links.end()) {
			error = "circular link detected";
			return "";
		}
		visited_links.push_back(index);
	}

	error.clear();

	std::string result;
	result.reserve(256); //Max length is 253

	while(true) {
		if(index + 1 > this->length) {
			error = "truncated data (missing code)";
			goto exit;
		}

		auto code = this->buffer[index++];
		if(code == 0) break;

		if((code >> 6U) == 3) {
			if(!allow_compression) {
				error = "found link, but links are not allowed";
				goto exit;
			}

			auto lower_addr = this->buffer[index++];
			if(index + 1 > this->length) {
				error = "truncated data (missing lower link address)";
				goto exit;
			}

			size_t addr = ((code & 0x3FU) << 8U) | lower_addr;
			if(addr >= this->length) {
				error = "invalid link address";
				goto exit;
			}
			auto tail = this->parse_dns_dn(error, addr, true);
			if(!error.empty())
				goto exit;

			if(!result.empty())
				result += "." + tail;
			else
				result = tail;
			break;
		} else {
			if(code > 63) {
				error = "max domain label length is 63 characters";
				goto exit;
			}

			if(!result.empty())
				result += ".";

			if(index + code >= this->length) {
				error = "truncated data (domain label)";
				goto exit;
			}

			result.append((const char*) (this->buffer + index), code);
			index += code;
		}
	}

	exit:
	if(allow_compression) visited_links.pop_back();
	return result;
}

DNSResponseData::~DNSResponseData() {
	::free(this->buffer);
}


DNSResponse::DNSResponse(uint8_t secure_state, const char* bogus, void *packet, size_t size) {
	this->bogus = bogus ? std::string{bogus} : std::string{"packet is secure"};
	this->secure_state = secure_state;

	this->data = std::make_shared<DNSResponseData>();
	this->data->buffer = (uint8_t*) malloc(size);
	this->data->length = size;

	memcpy(this->data->buffer, packet, size);
}

response::DNSHeader DNSResponse::header() const {
	return response::DNSHeader{this};
}

bool DNSResponse::parse(std::string &error) {
	if(this->is_parsed) {
		error = this->parse_error;
		return error.empty();
	}
	error.clear();
	this->is_parsed = true;

	auto header = this->header();
	size_t index = 12; /* 12 bits for the header */

	{
		auto count = header.query_count();
		this->parsed_queries.reserve(count);

		for(size_t idx = 0; idx < count; idx++) {
			auto dn = this->data->parse_dns_dn(error, index, true);
			if(!error.empty()) {
				error = "failed to parse query " + std::to_string(idx) + " dn: " + error; // NOLINT(performance-inefficient-string-concatenation)
				goto error_exit;
			}

			if(index + 4 > this->packet_length()) {
				error = "truncated data for query " + std::to_string(index);
				goto error_exit;
			}

			auto type = (rrtype::value) ntohs(*(uint16_t*) (this->data->buffer + index));
			index += 2;

			auto klass = (rrclass::value) ntohs(*(uint16_t*) (this->data->buffer + index));
			index += 2;

			this->parsed_queries.emplace_back(new response::DNSQuery{dn, type, klass});
		}
	}

	{
		auto count = header.answer_count();
		this->parsed_answers.reserve(count);

		for(size_t idx = 0; idx < count; idx++) {
			this->parsed_answers.push_back(this->parse_rr(error, index, true));
			if(!error.empty()) {
				error = "failed to parse answer " + std::to_string(idx) + ": " + error; // NOLINT(performance-inefficient-string-concatenation)
				goto error_exit;
			}
		}
	}

	{
		auto count = header.authority_count();
		this->parsed_authorities.reserve(count);

		for(size_t idx = 0; idx < count; idx++) {
			this->parsed_authorities.push_back(this->parse_rr(error, index, true));
			if(!error.empty()) {
				error = "failed to parse authority " + std::to_string(idx) + ": " + error; // NOLINT(performance-inefficient-string-concatenation)
				goto error_exit;
			}
		}
	}

	{
		auto count = header.additional_count();
		this->parsed_additionals.reserve(count);

		for(size_t idx = 0; idx < count; idx++) {
			this->parsed_additionals.push_back(this->parse_rr(error, index, true));
			if(!error.empty()) {
				error = "failed to parse additional " + std::to_string(idx) + ": " + error; // NOLINT(performance-inefficient-string-concatenation)
				goto error_exit;
			}
		}
	}

	return true;

	error_exit:
	this->parsed_queries.clear();
	this->parsed_answers.clear();
	this->parsed_authorities.clear();
	this->parsed_additionals.clear();
	return false;
}

std::shared_ptr<response::DNSResourceRecords> DNSResponse::parse_rr(std::string &error, size_t &index, bool allow_compressed) {
	auto dn = this->data->parse_dns_dn(error, index, allow_compressed);
	if(!error.empty()) {
		error = "failed to parse rr dn: " + error; // NOLINT(performance-inefficient-string-concatenation)
		return nullptr;
	}

	if(index + 10 > this->packet_length()) {
		error = "truncated header";
		return nullptr;
	}

	auto type = (rrtype::value) ntohs(*(uint16_t*) (this->data->buffer + index));
	index += 2;

	auto klass = (rrclass::value) ntohs(*(uint16_t*) (this->data->buffer + index));
	index += 2;

	auto ttl = ntohl(*(uint32_t*) (this->data->buffer + index));
	index += 4;

	auto payload_length = ntohs(*(uint16_t*) (this->data->buffer + index));
	index += 2;

	if(index + payload_length > this->packet_length()) {
		error = "truncated body";
		return nullptr;
	}

	auto response = std::shared_ptr<response::DNSResourceRecords>(new response::DNSResourceRecords{this->data, index, payload_length, ttl, dn, type, klass});
	index += payload_length;
	return response;
}


#ifndef WIN32
uint16_t DNSHeader::field(int index) const {
	return ((uint16_t*) this->response->packet_data())[index];
}

DNSResourceRecords::DNSResourceRecords(std::shared_ptr<DNSResponseData> packet, size_t payload_offset, size_t length, uint32_t ttl, std::string name, rrtype::value type, rrclass::value klass)
		: offset{payload_offset}, length{length}, ttl{ttl}, name{std::move(name)}, type{type}, klass{klass} {
	this->data = std::move(packet);
}

const uint8_t* DNSResourceRecords::payload_data() const  {
	return this->data->buffer + this->offset;
}
#else
DNSResourceRecords::DNSResourceRecords(std::shared_ptr<tc::dns::DNSResponseData> data, PDNS_RECORDA rdata) : nrecord{rdata}, data{std::move(data)} { }
bool DNSResourceRecords::is_wide_string() const {
    return this->data->wide_string;
}
#endif

bool A::is_valid() {
#ifdef WIN32
    return true;
#else
    return this->handle->payload_length() == 4;
#endif
}

in_addr A::address() {
#ifdef WIN32
    in_addr result{};
    result.S_un.S_addr = this->handle->native_record()->Data.A.IpAddress;
    return result;
#else
    //TODO: Attention: Unaligned access
    return {*(uint32_t*) this->handle->payload_data()};
#endif
}

std::string A::address_string() {
#ifdef WIN32
    struct in_addr address = this->address();
    char buffer[17];
    RtlIpv4AddressToStringA(&address, buffer);
    return std::string{buffer};
#else
    auto _1 = this->handle->payload_data()[0],
            _2 = this->handle->payload_data()[1],
            _3 = this->handle->payload_data()[2],
            _4 = this->handle->payload_data()[3];
    return std::to_string(_1) + "." + std::to_string(_2) + "." + std::to_string(_3) + "." + std::to_string(_4);
#endif
}

//---------------- AAAA
#ifdef WIN32
bool AAAA::is_valid() {
    return true;
}

std::string AAAA::address_string() {
    struct in6_addr address = this->address();
    char buffer[47];
    RtlIpv6AddressToStringA(&address, buffer); //Supported for Win7 as well and not only above 8.1 like inet_ntop
    return std::string{buffer};
}

in6_addr AAAA::address() {
    in6_addr result{};
    memcpy(result.u.Byte, this->handle->native_record()->Data.AAAA.Ip6Address.IP6Byte, 16);
    return result;
}
#else
bool AAAA::is_valid() {
    return this->handle->payload_length() == 16;
}

std::string AAAA::address_string() {
    auto address = this->address();

    char buffer[INET6_ADDRSTRLEN];
	if(!inet_ntop(AF_INET6, (void*) &address, buffer, INET6_ADDRSTRLEN)) return "";
	return std::string(buffer);
}

in6_addr AAAA::address() {
	return {
			.__in6_u = {
				.__u6_addr32 = {
						//TODO: Attention unaligned memory access
						((uint32_t*) this->handle->payload_data())[0],
						((uint32_t*) this->handle->payload_data())[1],
						((uint32_t*) this->handle->payload_data())[2],
						((uint32_t*) this->handle->payload_data())[3]
				}
			}
	};
}
#endif

//---------------- SRV
#ifdef WIN32
bool SRV::is_valid() { return true; }
std::string SRV::target_hostname() {
    if(this->handle->is_wide_string()) {
        auto result = std::wstring{ ((PDNS_RECORDW) this->handle->native_record())->Data.Srv.pNameTarget };
        return std::string{result.begin(), result.end()};
    } else {
        return std::string{ this->handle->native_record()->Data.Srv.pNameTarget };
    }
}
uint16_t SRV::priority() { return this->handle->native_record()->Data.SRV.wPriority; }
uint16_t SRV::weight() { return this->handle->native_record()->Data.SRV.wWeight; }
uint16_t SRV::target_port() { return this->handle->native_record()->Data.SRV.wPort; }
#else
bool SRV::is_valid() {
	if(this->handle->payload_length() < 7)
		return false;
	size_t index = this->handle->payload_offset() + 6;
	std::string error{};
	this->handle->dns_data()->parse_dns_dn(error, index, true);
	return error.empty();
}

std::string SRV::target_hostname() {
	size_t index = this->handle->payload_offset() + 6;
	std::string error{};
	return this->handle->dns_data()->parse_dns_dn(error, index, true);
}


uint16_t SRV::priority() { return ntohs(((uint16_t*) this->handle->payload_data())[0]); }
uint16_t SRV::weight() { return ntohs(((uint16_t*) this->handle->payload_data())[1]); }
uint16_t SRV::target_port() { return ntohs(((uint16_t*) this->handle->payload_data())[2]); }
#endif

//---------------- All types with a name
bool named_base::is_valid() {
#ifdef WIN32
    return true;
#else
	size_t index = this->handle->payload_offset();
	std::string error{};
	this->handle->dns_data()->parse_dns_dn(error, index, true);
	return error.empty();
#endif
}

std::string named_base::name() {
#ifdef WIN32
    if(this->handle->is_wide_string()) {
        auto result = std::wstring{ ((PDNS_RECORDW) this->handle->native_record())->Data.Cname.pNameHost };
        return std::string{result.begin(), result.end()};
    } else {
        return std::string{ this->handle->native_record()->Data.Cname.pNameHost };
    }
#else
    size_t index = this->handle->payload_offset();
	std::string error{};
	return this->handle->dns_data()->parse_dns_dn(error, index, true);
#endif
}