#include <algorithm>
#include <arpa/inet.h>
#include <cassert>
#include <clocale>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <errno.h>
#include <fcntl.h>
#include <idn2.h>
#include <inttypes.h>
#include <netinet/in.h>
#include <signal.h>
#include <string_view>
#include <sys/socket.h>
#include <sys/uio.h>
#include <unistd.h>
#include <wchar.h>
extern "C" {
#include <libknot/codes.h>
#include <libknot/dname.h>
#include <libknot/packet/pkt.h>
#include <libknot/packet/wire.h>
#include <libknot/rrtype/rdname.h>
#include <libknot/rrtype/tsig.h>
#include <libknot/tsig-op.h>
#include <libknot/tsig.h>

#define KNOT_PREREQUISITE KNOT_ANSWER
#define KNOT_UPDATE KNOT_AUTHORITY
}


template <class F>
struct quickscope_wrapper {
	F func;

	~quickscope_wrapper() { func(); }
};

template <class F>
quickscope_wrapper(F) -> quickscope_wrapper<F>;


static const char * inet_ntop_generic(const struct sockaddr_storage * sa, char * dst, socklen_t size) {
	switch(sa->ss_family) {
		case AF_INET:
			return inet_ntop(AF_INET, &(((struct sockaddr_in *)sa)->sin_addr), dst, size);

		case AF_INET6:
			return inet_ntop(AF_INET6, &(((struct sockaddr_in6 *)sa)->sin6_addr), dst, size);

		default:
			return nullptr;
	}
}


int main(int, const char * const * argv) {
	std::setlocale(LC_ALL, "C.UTF-8");

	{
		struct sigaction act {};
		act.sa_handler = [](int) {};
		sigaction(SIGALRM, &act, nullptr);
	}
	setlinebuf(stderr);

	int knot = socket(AF_INET6, SOCK_DGRAM, 0);
	{
#if TARTA
		struct sockaddr_in knot_in6 = {.sin_family = AF_INET, .sin_port = htobe16(53)};
		inet_pton(AF_INET, "139.28.40.42", &knot_in6.sin_addr);
#else
		const struct sockaddr_in6 knot_in6 = {.sin6_family = AF_INET6, .sin6_port = htobe16(0x4849), .sin6_addr = IN6ADDR_LOOPBACK_INIT};  // 'HI', 18505
#endif
		connect(knot, reinterpret_cast<const struct sockaddr *>(&knot_in6), sizeof(knot_in6));
	}

	int validator = socket(AF_INET, SOCK_DGRAM, 0);
	{
		const struct sockaddr_in validator_in = {.sin_family = AF_INET, .sin_port = htobe16(5353), .sin_addr = {htobe32(INADDR_LOOPBACK)}};  // local unbound
		connect(validator, reinterpret_cast<const struct sockaddr *>(&validator_in), sizeof(validator_in));
	}


	knot_tsig_key_t key_public;  // public.zones.hinfo.network
	assert(!knot_tsig_key_init_str(&key_public, argv[2]));

	knot_tsig_key_t key_private;  // proxy.zones.hinfo.network
	if(auto priv = std::getenv("PRIVKEY"))
		assert(!knot_tsig_key_init_str(&key_private, priv));
	else
		assert(!knot_tsig_key_init_file(&key_private, "/dev/stdin"));

	auto zones_hinfo_network       = knot_dname_from_str_alloc("zones.hinfo.network.");
	auto hinfo_zones_hinfo_network = knot_dname_from_str_alloc("hinfo.zones.hinfo.network.");
	assert(zones_hinfo_network);
	assert(hinfo_zones_hinfo_network);


	sockaddr_in listening_in = {
	    .sin_family = AF_INET,
	    .sin_port   = htobe16(0x4849),
	};
	inet_pton(AF_INET, argv[1], &listening_in.sin_addr);
	int listening = socket(AF_INET, SOCK_DGRAM, 0);
	assert(!bind(listening, reinterpret_cast<const struct sockaddr *>(&listening_in), sizeof(listening_in)));

	uint8_t digest_buf[KNOT_WIRE_MAX_PKTSIZE];

	std::uint16_t forwarded_txid = -1, validation_txid = -1;
	knot_pkt * forward_pkt   = knot_pkt_new(nullptr, KNOT_WIRE_MAX_PKTSIZE, nullptr);
	knot_pkt * response_pkt  = knot_pkt_new(nullptr, KNOT_WIRE_MAX_PKTSIZE, nullptr);
	knot_pkt * validator_pkt = knot_pkt_new(nullptr, KNOT_WIRE_MAX_PKTSIZE, nullptr);
	assert(forward_pkt);
	assert(response_pkt);
	assert(validator_pkt);

	for(;;) {
		struct sockaddr_storage peer = {};
		char buf[KNOT_WIRE_MAX_PKTSIZE + 1];
		struct iovec iov  = {.iov_base = buf, .iov_len = sizeof(buf)};
		struct msghdr msg = {
		    .msg_name    = &peer,
		    .msg_namelen = sizeof(peer),

		    .msg_iov    = &iov,
		    .msg_iovlen = 1,
		};
		auto recvlen = recvmsg(listening, &msg, 0);
		if(recvlen == -1)
			continue;

		auto pkt = knot_pkt_new(buf, recvlen, nullptr);
		if(!pkt)
			continue;
		quickscope_wrapper pkt_deleter{[&] { knot_pkt_free(pkt); }};

		if(auto p = knot_pkt_parse(pkt, 0); p != KNOT_EOK && p != KNOT_ETRAIL)
			continue;


		{
			char peername[std::max(INET_ADDRSTRLEN, INET6_ADDRSTRLEN) + 1];
			struct timespec now;
			clock_gettime(CLOCK_REALTIME, &now);
			std::fprintf(stderr, "%llu.%09d\t%s", (unsigned long long)now.tv_sec, (int)now.tv_nsec,
			             inet_ntop_generic(&peer, peername, sizeof(peername)) ?: "<unknown>");
		}

		auto reply = [&](short rcode, short tsig_rcode, bool mac, auto && data) {
			knot_pkt_clear(response_pkt);
			if(knot_pkt_init_response(response_pkt, pkt) != KNOT_EOK) {
				std::fputs("\t_rerr\n", stderr);
				return;
			}
			{
				struct timespec now;
				clock_gettime(CLOCK_REALTIME, &now);
				auto rcode_l      = knot_lookup_by_id(knot_rcode_names, rcode);
				auto tsig_rcode_l = knot_lookup_by_id(knot_rcode_names, tsig_rcode);
				std::fprintf(stderr, "\t%s\t%s\t%llu.%09d\n", rcode_l ? rcode_l->name : "<unknown>", tsig_rcode_l ? tsig_rcode_l->name : "<unknown>",
				             (unsigned long long)now.tv_sec, (int)now.tv_nsec);
			}

			data(response_pkt);

			knot_wire_set_rcode(response_pkt->wire, rcode);

			if(mac)
				knot_tsig_sign(response_pkt->wire, &response_pkt->size, response_pkt->max_size,              //
				               knot_tsig_rdata_mac(pkt->tsig_rr), knot_tsig_rdata_mac_length(pkt->tsig_rr),  //
				               pkt->wire, &pkt->size,                                                        //
				               &key_public, tsig_rcode, knot_tsig_rdata_time_signed(pkt->tsig_rr));
			else
				knot_tsig_add(response_pkt->wire, &response_pkt->size, response_pkt->max_size, tsig_rcode, pkt->tsig_rr);


			iov.iov_base = response_pkt->wire;
			iov.iov_len  = response_pkt->size;
			(void)sendmsg(listening, &msg, 0);
		};
#define ERROR_REPLY(...)                     \
	__extension__({                            \
		reply(__VA_ARGS__, [](knot_pkt_t *) {}); \
		continue;                                \
	})
#define ERROR_REPLY_C2(...)                  \
	do {                                       \
		reply(__VA_ARGS__, [](knot_pkt_t *) {}); \
		goto continue2;                          \
	} while(false)


		// copied from src/knot/nameserver/process_query.c:process_query_verify()
		switch(knot_tsig_server_check(pkt->tsig_rr, pkt->wire, pkt->size, &key_public)) {
			case KNOT_EOK:
				break;
			case KNOT_TSIG_EBADKEY:
				ERROR_REPLY(KNOT_RCODE_NOTAUTH, KNOT_RCODE_BADKEY, false);
			case KNOT_TSIG_EBADSIG:
				ERROR_REPLY(KNOT_RCODE_NOTAUTH, KNOT_RCODE_BADSIG, false);
			case KNOT_TSIG_EBADTIME:
				ERROR_REPLY(KNOT_RCODE_NOTAUTH, KNOT_RCODE_BADTIME, false);
			case KNOT_EMALF:
				ERROR_REPLY(KNOT_RCODE_FORMERR, KNOT_RCODE_NOERROR, false);
			default:
				ERROR_REPLY(KNOT_RCODE_SERVFAIL, KNOT_RCODE_NOERROR, false);
		}

		if(!knot_dname_is_case_equal(knot_pkt_qname(pkt), zones_hinfo_network))
			ERROR_REPLY(KNOT_RCODE_NOTAUTH, KNOT_RCODE_NOERROR, true);

		if(pkt->sections[KNOT_PREREQUISITE].count)
			ERROR_REPLY(KNOT_RCODE_NOTIMPL, KNOT_RCODE_NOERROR, true);


		knot_rrset_t collected;
		knot_rrset_init(&collected, hinfo_zones_hinfo_network, KNOT_RRTYPE_PTR, KNOT_CLASS_IN, 3600);
		quickscope_wrapper collected_deleter{[&] {
			collected.owner = nullptr;
			knot_rrset_clear(&collected, nullptr);
		}};


		for(std::uint16_t i = 0; i < pkt->sections[KNOT_UPDATE].count; ++i) {
			auto rr = knot_pkt_rr(&pkt->sections[KNOT_UPDATE], i);
			//                                                               TTL=0 is delete
			if(!knot_dname_is_equal(rr->owner, hinfo_zones_hinfo_network) || !rr->ttl || rr->rclass != KNOT_CLASS_IN || rr->type != KNOT_RRTYPE_PTR)
				ERROR_REPLY_C2(KNOT_RCODE_REFUSED, KNOT_RCODE_NOERROR, true);

			for(std::uint16_t r = 0; r < rr->rrs.count; ++r) {
				char name[KNOT_DNAME_TXT_MAXLEN + 1];
				if(!knot_dname_to_str(name, knot_ptr_name(&rr->rrs.rdata[r]), sizeof(name)))  // artificial limit: 50-byte name
					ERROR_REPLY_C2(KNOT_RCODE_FORMERR, KNOT_RCODE_NOERROR, true);
				std::fputc('\t', stderr), std::fputs(name, stderr);
				if(std::strlen(name) > 50) {
					char * decoded{};
					if(idn2_to_unicode_lzlz(name, &decoded, 0) != IDNA_SUCCESS)
						ERROR_REPLY_C2(KNOT_RCODE_REFUSED, KNOT_RCODE_NOERROR, true);
					quickscope_wrapper decoded_deleter{[&] { std::free(decoded); }};

					std::string_view cur{decoded};
					std::mbstate_t ctx{};
					std::size_t width{};
					for(wchar_t c; !cur.empty();)
						switch(auto r = std::mbrtowc(&c, cur.data(), cur.size(), &ctx)) {
							case static_cast<std::size_t>(-1):  // EILSEQ: reset, try to go past
							case static_cast<std::size_t>(-2):  // incomplete: consumed, same as 0-width for output
							case 0:
								ERROR_REPLY_C2(KNOT_RCODE_REFUSED, KNOT_RCODE_NOERROR, true);
								break;
							default:
								width += std::max(wcwidth(c), 0);
								if(width > 50)
									ERROR_REPLY_C2(KNOT_RCODE_REFUSED, KNOT_RCODE_NOERROR, true);
								cur.remove_prefix(r);
								break;
						}
				}


				auto yxdomain = [&] {
					reply(KNOT_RCODE_YXDOMAIN, KNOT_RCODE_NOERROR, true, [&](knot_pkt_t * p) {
						if(knot_pkt_begin(response_pkt, KNOT_ADDITIONAL) == KNOT_EOK) {
							knot_rrset_t failed;
							knot_rrset_init(&failed, hinfo_zones_hinfo_network, KNOT_RRTYPE_PTR, KNOT_CLASS_IN, 0);
							if(knot_rrset_add_rdata(&failed, rr->rrs.rdata[r].data, rr->rrs.rdata[r].len, nullptr) == KNOT_EOK)
								knot_pkt_put(p, KNOT_COMPR_HINT_NONE, &failed, 0);
							failed.owner = nullptr;
							knot_rrset_clear(&failed, nullptr);
						}
					});
				};
#define YXDOMAIN()  \
	do {              \
		yxdomain();     \
		goto continue2; \
	} while(false)


				knot_pkt_clear(validator_pkt);
				if(knot_pkt_put_question(validator_pkt, knot_ptr_name(&rr->rrs.rdata[r]), KNOT_CLASS_IN, KNOT_RRTYPE_HINFO) != KNOT_EOK)
					YXDOMAIN();
				knot_wire_set_rd(validator_pkt->wire);
				knot_wire_set_id(validator_pkt->wire, ++validation_txid);

				if(send(validator, validator_pkt->wire, validator_pkt->size, 0) <= 0)
					YXDOMAIN();
				for(;;) {
					alarm(3);
					auto recvd = recv(validator, validator_pkt->wire, validator_pkt->max_size, 0);
					alarm(0);
					if(recvd <= 0)
						YXDOMAIN();
					if(knot_wire_get_id(validator_pkt->wire) != validation_txid)  // stale reply
						continue;
					if(knot_wire_get_rcode(validator_pkt->wire))
						YXDOMAIN();
					// loopback, and thus perfect, transport: don't retry
					// unbound is configured to validate: we just need to check for NOERROR and whether we only got HINFOs (so no SOA, no CNAME, &c.)

					auto repl_pkt = knot_pkt_new(validator_pkt->wire, recvd, nullptr);
					if(!repl_pkt)
						YXDOMAIN();
					quickscope_wrapper pkt_deleter{[&] { knot_pkt_free(repl_pkt); }};
					if(auto p = knot_pkt_parse(repl_pkt, 0); p != KNOT_EOK && p != KNOT_ETRAIL)
						YXDOMAIN();
					bool any = false;
					for(std::uint16_t a = 0; a < repl_pkt->sections[KNOT_ANSWER].count; ++a) {
						if(knot_pkt_rr(&repl_pkt->sections[KNOT_ANSWER], a)->type != KNOT_RRTYPE_HINFO)
							YXDOMAIN();
						any = true;
					}
					if(!any)
						YXDOMAIN();
					break;
				}


				if(knot_rrset_add_rdata(&collected, rr->rrs.rdata[r].data, rr->rrs.rdata[r].len, nullptr) != KNOT_EOK)
					ERROR_REPLY_C2(KNOT_RCODE_SERVFAIL, KNOT_RCODE_NOERROR, true);
			}
		}

		knot_pkt_clear(forward_pkt);
		if(knot_pkt_put_question(forward_pkt, zones_hinfo_network, KNOT_CLASS_IN, KNOT_RRTYPE_SOA) != KNOT_EOK ||  //
		   knot_pkt_begin(forward_pkt, KNOT_UPDATE) != KNOT_EOK ||                                                 //
		   knot_pkt_put(forward_pkt, KNOT_COMPR_HINT_NONE, &collected, 0) != KNOT_EOK)
			ERROR_REPLY(KNOT_RCODE_SERVFAIL, KNOT_RCODE_NOERROR, true);

		knot_wire_set_opcode(forward_pkt->wire, KNOT_OPCODE_UPDATE);
		knot_wire_set_id(forward_pkt->wire, ++forwarded_txid);

		std::size_t digest_len;
		if(knot_tsig_sign(forward_pkt->wire, &forward_pkt->size, forward_pkt->max_size,  //
		                  nullptr, 0,                                                    //
		                  digest_buf, &(digest_len = sizeof(digest_buf)),                //
		                  &key_private, KNOT_RCODE_NOERROR, 0 != KNOT_EOK))
			ERROR_REPLY(KNOT_RCODE_SERVFAIL, KNOT_RCODE_NOERROR, true);

		// loopback, and thus perfect, transport: don't retry
		if(send(knot, forward_pkt->wire, forward_pkt->size, 0) <= 0)
			ERROR_REPLY(KNOT_RCODE_SERVFAIL, KNOT_RCODE_NOERROR, true);
		for(;;) {
			alarm(3);
			auto recvd = recv(knot, forward_pkt->wire, forward_pkt->max_size, 0);
			alarm(0);
			if(recvd <= 0)
				ERROR_REPLY(KNOT_RCODE_SERVFAIL, KNOT_RCODE_NOERROR, true);
			if(knot_wire_get_id(forward_pkt->wire) != forwarded_txid)  // stale reply
				continue;

			auto repl_pkt = knot_pkt_new(forward_pkt->wire, recvd, nullptr);
			if(!repl_pkt)
				ERROR_REPLY(KNOT_RCODE_SERVFAIL, KNOT_RCODE_NOERROR, true);
			quickscope_wrapper pkt_deleter{[&] { knot_pkt_free(repl_pkt); }};
			if(auto p = knot_pkt_parse(repl_pkt, 0); p != KNOT_EOK && p != KNOT_ETRAIL)
				ERROR_REPLY(KNOT_RCODE_SERVFAIL, KNOT_RCODE_NOERROR, true);

			reply(knot_wire_get_rcode(forward_pkt->wire), knot_pkt_has_tsig(repl_pkt) ? (knot_rcode_t)knot_tsig_rdata_error(repl_pkt->tsig_rr) : KNOT_RCODE_NOERROR,
			      true, [&](knot_pkt_t * p) {
				      for(int s = 0; s < KNOT_PKT_SECTIONS; ++s)
					      if(repl_pkt->sections[s].count)
						      if(knot_pkt_begin(response_pkt, static_cast<knot_section_t>(s)) == KNOT_EOK)
							      for(std::uint16_t r = 0; r < repl_pkt->sections[r].count; ++r)
								      knot_pkt_put(p, KNOT_COMPR_HINT_NONE, knot_pkt_rr(&repl_pkt->sections[s], r), 0);
			      });
			break;
		}

	continue2:;
	}
}