package net.azib.ipscan.util;

import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.DatagramPacket;
import java.net.InetAddress;
import java.net.MulticastSocket;

public class MDNSResolver implements Closeable {
	InetAddress mdnsIP = InetAddress.getByName("224.0.0.251");
	private int mdnsPort = 5353;
	MulticastSocket socket = new MulticastSocket();

	public MDNSResolver(int timeout) throws IOException {
		socket.setSoTimeout(timeout);
		socket.setTimeToLive(1);
		// TODO: iPhones respond only back to the multicast address and port 5353 provided requestId == 0
		// socket.joinGroup(mdnsIP);
	}

	void writeName(DataOutputStream out, String name) throws IOException {
		int s = 0, e;
		while ((e = name.indexOf('.', s)) != -1) {
			out.writeByte(e - s);
			out.write(name.substring(s, e).getBytes());
			s = e + 1;
		}
		out.write(name.length() - s);
		out.write(name.substring(s).getBytes());
		out.writeByte(0);
	}

	String decodeName(byte[] data, int offset, int length) {
		var s = new StringBuilder(length);
		for (var i = offset; i < offset + length; i++) {
			var len = data[i];
			if (len == 0) break;
			s.append(new String(data, i + 1, len)).append('.');
			i += len;
		}
		s.setLength(s.length() - 1);
		return s.toString();
	}

	byte[] dnsRequest(int id, String name) throws IOException {
		var baos = new ByteArrayOutputStream();
		var out = new DataOutputStream(baos);
		out.writeShort(id);
		out.write(new byte[] {0, 0, 0, 1, 0, 0, 0, 0, 0, 0});
		writeName(out, name);
		out.write(new byte[] {0, 0xc, 0, 1});
		return baos.toByteArray();
	}

	String reverseName(byte[] addr) {
		// note: only IPv4 is supported here
		return (addr[3]&0xFF) + "." + (addr[2]&0xFF) + "." + (addr[1]&0xFF) + "." + (addr[0]&0xFF) + ".in-addr.arpa";
	}

	public String resolve(InetAddress ip) throws IOException {
		var addr = ip.getAddress();
		var requestId = addr[2] * 0xFF + addr[3];
		var request = dnsRequest(requestId, reverseName(addr));
		socket.send(new DatagramPacket(request, request.length, mdnsIP, mdnsPort));

		var respPacket = new DatagramPacket(new byte[512], 512);
		socket.receive(respPacket);
		var response = respPacket.getData();
		if (response[0] != request[0] && response[1] != request[1]) return null;
		int numQueries = response[5];
		var offset = (numQueries == 0 ? 12 : request.length) + 2 + 2 + 2 + 4 + 2;
		return decodeName(response, offset, respPacket.getLength() - offset);
	}

	public void close() {
		socket.close();
	}

	public static void main(String[] args) throws IOException {
		System.out.println(new MDNSResolver(2000).resolve(InetAddress.getByName("192.168.0.10")));
	}
}
