(*---------------------------------------------------------------------------
   Copyright (c) 2025 Thomas Gazagnaire. All rights reserved.
   SPDX-License-Identifier: MIT
  ---------------------------------------------------------------------------*)

(** mdns-query - mDNS service discovery CLI tool *)

let pp_domain ppf name = Fmt.pf ppf "%s" (Domain_name.to_string name)

let pp_response ppf (r : Mdns.response) =
  List.iter
    (fun (service, instance) ->
      Fmt.pf ppf "@[<h>PTR  %a -> %a@]@." pp_domain service pp_domain instance)
    r.ptrs;
  List.iter
    (fun (instance, port, target) ->
      Fmt.pf ppf "@[<h>SRV  %a -> %a:%d@]@." pp_domain instance pp_domain target
        port)
    r.srvs;
  List.iter
    (fun (instance, txts) ->
      Fmt.pf ppf "@[<h>TXT  %a -> %a@]@." pp_domain instance
        Fmt.(list ~sep:(any " ") string)
        txts)
    r.txts;
  List.iter
    (fun (name, ip) ->
      Fmt.pf ppf "@[<h>A    %a -> %a@]@." pp_domain name Ipaddr.V4.pp ip)
    r.addrs;
  List.iter
    (fun (name, ip) ->
      Fmt.pf ppf "@[<h>AAAA %a -> %a@]@." pp_domain name Ipaddr.V6.pp ip)
    r.addrs6

let is_empty (r : Mdns.response) =
  r.ptrs = [] && r.srvs = [] && r.txts = [] && r.addrs = [] && r.addrs6 = []

let run service timeout quiet =
  Eio_main.run @@ fun env ->
  Eio.Switch.run @@ fun sw ->
  let net = Eio.Stdenv.net env in
  let clock = Eio.Stdenv.clock env in
  let name =
    try Domain_name.of_string_exn service
    with Invalid_argument msg ->
      Fmt.epr "Error: invalid service name %S: %s@." service msg;
      exit 1
  in
  if not quiet then Fmt.pr "Querying %s (timeout: %.1fs)...@.@." service timeout;
  let responses = Mdns.query ~sw ~net ~clock ~timeout name in
  let merged = Mdns.merge responses in
  if is_empty merged then begin
    if not quiet then Fmt.pr "No responses received.@.";
    exit 1
  end
  else begin
    if not quiet then Fmt.pr "Results:@.";
    pp_response Fmt.stdout merged
  end

(* Cmdliner interface *)
open Cmdliner

let service =
  let doc =
    "Service type to query. Use DNS-SD format: $(b,_service._proto.local) \
     (e.g., $(b,_http._tcp.local), $(b,_hap._tcp.local))."
  in
  Arg.(required & pos 0 (some string) None & info [] ~docv:"SERVICE" ~doc)

let timeout =
  let doc = "Timeout in seconds to wait for responses." in
  Arg.(value & opt float 3.0 & info [ "t"; "timeout" ] ~docv:"SECONDS" ~doc)

let quiet =
  let doc = "Suppress informational messages; only print results." in
  Arg.(value & flag & info [ "q"; "quiet" ] ~doc)

let cmd =
  let doc = "Query mDNS services on the local network" in
  let man =
    [
      `S Manpage.s_description;
      `P
        "$(tname) sends an mDNS query for the specified service type and \
         displays all responses received within the timeout period.";
      `P "Exit status is 0 if at least one response was received, 1 otherwise.";
      `S Manpage.s_examples;
      `P "Query for HomeKit accessories:";
      `Pre "  $(tname) _hap._tcp.local";
      `P "Query for Matter devices:";
      `Pre "  $(tname) _matter._tcp.local";
      `P "Query for HTTP servers with longer timeout:";
      `Pre "  $(tname) -t 5 _http._tcp.local";
      `P "Quiet mode for scripting:";
      `Pre "  $(tname) -q _hap._tcp.local";
      `S Manpage.s_see_also;
      `P "$(b,dns-sd)(1), RFC 6762 (mDNS), RFC 6763 (DNS-SD)";
    ]
  in
  let info = Cmd.info "mdns-query" ~version:"%%VERSION%%" ~doc ~man in
  Cmd.v info Term.(const run $ service $ timeout $ quiet)

let () = exit (Cmd.eval cmd)