# SPDX-FileCopyrightText: 2024 Łukasz Niemier <#@hauleth.dev>
#
# SPDX-License-Identifier: MPL-2.0

defmodule AwwTest.Dns do
  use GenServer

  @callback handle_query(
              String.t(),
              :inet_req.dns_class(),
              :inet_req.dns_rr_type()
            ) :: map()

  def start_link(module: mod) do
    with {:ok, pid} <- GenServer.start_link(__MODULE__, mod) do
      {:ok, pid, ns(pid)}
    end
  end

  def ns(pid), do: GenServer.call(pid, :name)

  def init(mod) do
    {:ok, sock} = :gen_udp.open(0, mode: :binary)

    {:ok, %{sock: sock, mod: mod}}
  end

  def handle_call(:name, _ref, state) do
    {:ok, name} = :inet.sockname(state.sock)

    {:reply, name, state}
  end

  def handle_info({:udp, sock, from_ip, from_port, data}, %{sock: sock} = state) do
    with {:ok, req} <- :inet_dns.decode(data) do
      req = :inet_dns.msg(req)
      header = :inet_dns.header(req[:header])
      questions = req[:qdlist]

      answers =
        for q <- questions,
            query = :inet_dns.dns_query(q),
            answer <- handle(state, query) do
          :inet_dns.make_rr(answer)
        end

      resp_header =
        :inet_dns.make_header(
          id: header[:id],
          aa: false,
          qr: true,
          rd: true,
          ra: true,
          opcode: :query,
          rcode: 0
        )

      resp =
        :inet_dns.make_msg(
          header: resp_header,
          qdlist: questions,
          anlist: answers
        )

      :gen_udp.send(sock, from_ip, from_port, :inet_dns.encode(resp))
    end

    {:noreply, state}
  end

  defp handle(%{mod: mod}, query) do
    domain = query[:domain]
    class = query[:class]
    type = query[:type]

    resp =
      try do
        mod.handle_query(List.to_string(domain), class, type)
      catch
        _, _ -> []
      end

    for %{ttl: ttl, data: data} <- List.wrap(resp) do
      [domain: domain, type: type, class: class, ttl: ttl, data: data]
    end
  end
end