(defpackage :cl-mdns/packet
  (:use :cl :arrow-macros)
  (:local-nicknames (:nb :nibbles))
  (:export :build-query-packet :build-response :make-ip-from-string :encode-name :parse-query))

(in-package :cl-mdns/packet)

(defun encode-name (name-string)
  "Converts local domain name to encoded byte array.
  E.g. converts 'name.local' to #(4 110 97 109 101 5 108 111 99 97 108 0)"
  (let ((parts (uiop:split-string name-string :separator "."))
        (output (make-array 0 :element-type '(unsigned-byte 8) :adjustable t :fill-pointer 0)))
    (dolist (part parts)
      (vector-push-extend (length part) output)
      (loop for char across part do
        (vector-push-extend (char-code char) output)))
    (vector-push-extend 0 output) ; Terminate
    output))

(defun make-question-header ()
  (let ((buffer (make-array 12 :element-type '(unsigned-byte 8) :initial-element 0)))
    ;; Transaction ID (2 bytes): 0 for mDNS
    ;; Flags (2 bytes): 0x0000 basic query
    ;; QDCOUNT (2 bytes): 1
    (setf (nb:ub16ref/be buffer 4) 1)
    buffer))

(defun make-question-tail ()
  (let ((buffer (make-array 4 :element-type '(unsigned-byte 8))))
    (setf (nibbles:ub16ref/be buffer 0) #x01) ; Type: AA
    (setf (nibbles:ub16ref/be buffer 2) 1)  ; Class: IN
    buffer))

(defun build-query-packet (service-name)
  "Creates 'question' packet for given SERVICE-NAME."
  (let ((header (make-question-header))
        (name (encode-name service-name))
        (tail (make-question-tail)))
    (concatenate '(vector (unsigned-byte 8)) header name tail)))

(defun make-dns-header ()
  (let ((buffer (make-array 12 :element-type '(unsigned-byte 8) :initial-element 0)))
    ;; Header
    (setf (nb:ub16ref/be buffer 0) 0)
    (setf (nb:ub16ref/be buffer 2) #x8400)   ; Flags: Standard response, AA
    (setf (nb:ub16ref/be buffer 4) 0)        ; Questions
    (setf (nb:ub16ref/be buffer 6) 1)        ; Answer RRs
    (setf (nb:ub16ref/be buffer 8) 0)        ; Authority RRs
    (setf (nb:ub16ref/be buffer 10) 0)       ; Additional RRs

    buffer))

(defun ip-addr-list-p (x)
  (and (listp x)
       (= (length x) 4)
       (every #'integerp x)))

(deftype ip-addr-list ()
  '(satisfies ip-addr-list-p))

(defun make-ip-from-string (ip-str)
  (-<>> (uiop:split-string ip-str :separator ".")
    (mapcar #'parse-integer)
    (lambda (x) (if (typep x 'ip-addr-list)
		    x
		    (error "~S is not in valid ipv4 format." ip-str)))))

(defun make-ip-body (ip)
  (destructuring-bind (ip-a ip-b ip-c ip-d) ip
    (let ((buffer (make-array 14 :element-type '(unsigned-byte 8) :initial-element 0))
	  (pos 0))
      ;; Type A (1), Class IN (1)
      (setf (nb:ub16ref/be buffer pos) 1) (incf pos 2)
      (setf (nb:ub16ref/be buffer pos) 1) (incf pos 2)
      ;; TTL (120 seconds)
      (setf (nb:ub32ref/be buffer pos) 120) (incf pos 4)
      ;; Data Length (4 bytes for IPv4)
      (setf (nb:ub16ref/be buffer pos) 4) (incf pos 2)
      (setf (aref buffer pos) ip-a) (incf pos)
      (setf (aref buffer pos) ip-b) (incf pos)
      (setf (aref buffer pos) ip-c) (incf pos)
      (setf (aref buffer pos) ip-d)
      buffer)))

(defun build-response (domain-name ip)
  (check-type ip ip-addr-list)
  (let ((header (make-dns-header))
	(name (encode-name domain-name))
	(ip-part (make-ip-body ip)))
    (concatenate '(vector (unsigned-byte 8)) header name ip-part)))

(defun parse-query-body-message (lst)
  "Parses the body of message.
The rest of the packet (meaning the tail) is ignored."
  (labels ((parse-rec (acc rem-lst)
	     (destructuring-bind (cnt &rest msg) rem-lst
	       (if (or (null cnt) (zerop cnt))
		   acc
		   (parse-rec
		    (-<>> (subseq msg 0 cnt)
		      (mapcar #'code-char)
		      (coerce <> 'string)
		      list
		      (concatenate 'list acc))
		    (nthcdr cnt msg))))))
    (parse-rec nil lst)))

(defun parse-query (query)
  "Parses question/query and returns domain name.
Completely ignores header and tail of packet for now."
  (let ((message (coerce (subseq query 12) 'list)))
    (format nil "~{~A~^.~}"
     (parse-query-body-message message))))