;;;; -*- mode: lisp -*-
(defpackage :geoip
  (:use :common-lisp :cffi :split-sequence)
  (:export
   :get-in :parse-ip :query-ip
   :make-mmdb :close-mmdb :with-mmdb))

(in-package :geoip)

(defun get-in (data &rest keys)
  "Get nested value from alist/map structure"
  (reduce (lambda (acc key) (cdr (assoc key acc :test #'equal))) keys :initial-value data))

(defun parse-ipv4 (ip-address)
  (reduce (lambda (acc n)
            (assert (<= 0 n 255))
            (+ (ash acc 8) n))
          (split-sequence #\. ip-address)
          :key #'parse-integer
          :initial-value 0))

(defun parse-ipv6 (ip-address)
  (reduce (lambda (acc n)
            (assert (<= 0 n #xFFFF))
            (+ (ash acc 16) n))
          (split-sequence #\: ip-address)
          :key (lambda (part)(if (string= part "") 0 (parse-integer part :radix 16)) )
          :initial-value 0))

(defun parse-ip (ip-address)
  (if (find #\. ip-address)
      (cons 4 (parse-ipv4 ip-address))
      (cons 6 (parse-ipv6 ip-address))))

(alexandria:define-constant
    +metadata-marker+
    (make-array 14 :element-type '(unsigned-byte 8)
                   :initial-contents
                (list* #xab #xcd #xef (map 'list #'char-code "MaxMind.com")))
  :test #'equalp)

(defstruct maxmind-database
  (filename nil :read-only t)
  (ptr nil :read-only t)
  (fd nil :read-only t)
  (size 0 :read-only t)
  (node-count 0 :read-only t :type fixnum)
  (record-size 0 :read-only t :type fixnum)
  (ip-version 0 :read-only t :type fixnum)
  (data-offset 0 :read-only t :type fixnum)
  (metadata nil))

(defstruct db-reader
  (db-ptr)
  (head-ptr)
  (data-offset 0 :read-only t :type fixnum))

(defstruct maxmind-database-metadata
  (node-count 0 :read-only t :type fixnum)
  (record-size 0 :read-only t :type fixnum)
  (ip-version 0 :read-only t :type fixnum)
  (database-type nil :read-only t :type string)
  (languages nil :read-only t :type list)
  (binary-format-major-version 0 :read-only t :type fixnum)
  (binary-format-minor-version 0 :read-only t :type fixnum)
  (build-epoch 0 :read-only t :type fixnum)
  (description nil :read-only t :type list))


(defun make-metadata (data)
  (make-maxmind-database-metadata
   :binary-format-major-version (get-in data :binary-format-major-version)
   :binary-format-minor-version (get-in data :binary-format-minor-version)
   :build-epoch (get-in data :build-epoch)
   :database-type (get-in data :database-type)
   :description (get-in data :description)
   :ip-version (get-in data :ip-version)
   :languages (get-in data :languages)
   :node-count (get-in data :node-count)
   :record-size (get-in data :record-size)))

(defun metadata-marker-p (ptr offset)
  (loop for i below #.(length +metadata-marker+)
        always (eql (mem-ref ptr :uchar (+ i offset))
                    (aref +metadata-marker+ i))))

(defun find-metadata-start (ptr size)
  (loop for offset from (- size 1 #.(length +metadata-marker+))
          downto (- size 1 (* 128 1024))
        when (metadata-marker-p ptr offset)
          return (+ offset #.(length +metadata-marker+))))

(defun read-db-char (reader)
  (with-slots (db-ptr) reader
    (prog1 (mem-ref db-ptr :uchar)
      (incf-pointer db-ptr))))

(defun bytes-from-foreign (reader data-length)
  (let ((bytes (make-array data-length :element-type '(unsigned-byte 8))))
    (with-slots (db-ptr) reader
      (foreign-array-to-lisp db-ptr
                             `(:array :uchar ,data-length)
                             :element-type '(unsigned-byte 8)
                             :displaced-to bytes)
      (incf-pointer db-ptr data-length))
    bytes))

(defun mread-uft8 (reader length)
  (with-slots (db-ptr) reader
    (prog1 (foreign-string-to-lisp db-ptr :count length :encoding :utf-8)
      (incf-pointer db-ptr length))))

(defmacro mem-uint (ptr size)
  `(let ((r 0))
     (dotimes (i ,size r)
       (setf r (+ (ash r 8) (mem-ref ,ptr :uint8 i))))))

(defun mread-unsigned (reader size)
  (with-slots (db-ptr) reader
    (prog1 (mem-uint db-ptr size)
      (incf-pointer db-ptr size))))

(defun mread-int32 (db-ptr length)
  (let ((uval (mread-unsigned db-ptr length)))
    (if (and (= length 4) (logbitp 31 uval))
        (- uval #.(expt 2 32))
        uval)))

(defun mread-datafield-metadata (db-ptr)
  (let* ((control-byte (read-db-char db-ptr))
         (type (ldb (byte 3 5) control-byte))
         (length (ldb (byte 5 0) control-byte)))
    (when (zerop type)
      (setf type (+ 7 (read-db-char db-ptr))))
    (list type
          (cond
            ((or (= type 1) (< length 29))
             length)
            ((= length 29)
             (+ 29 (read-db-char db-ptr)))
            ((= length 30)
             (+ 285 (mread-unsigned db-ptr 2)))
            ((= length 31)
             (+ 65821 (mread-unsigned db-ptr 3)))))))


(defun mread-map (db-ptr length)
  (loop repeat length
        collect
        (cons
         (intern (string-upcase (substitute #\- #\_ (mread-data db-ptr))) :keyword)
         (mread-data db-ptr))))

(defun mread-list (db-ptr length)
  (loop repeat length collect (mread-data db-ptr)))

(defun read-pointer (reader length)
  (let ((size-bits (ldb (byte 2 3) length))
        (value-bits (ldb (byte 3 0) length)))
    (ecase size-bits
      (0 (+ (ash value-bits 8) (read-db-char reader)))
      (1 (+ (ash value-bits 16) (mread-unsigned reader 2) 2048))
      (2 (+ (ash value-bits 24) (mread-unsigned reader 3) 526336))
      (3 (mread-unsigned reader 4)))))

(defun mread-pointer (reader length)
  (with-slots (db-ptr head-ptr data-offset) reader
    (let ((target (+ (read-pointer reader length) data-offset))
          (curr-ptr db-ptr))
      (setf db-ptr (inc-pointer head-ptr target))
      (prog1 (mread-data reader)
        (setf db-ptr curr-ptr)))))

(defun mread-data (db-ptr)
  (destructuring-bind (type length) (mread-datafield-metadata db-ptr)
    (ecase type
      (1 (mread-pointer db-ptr length))
      (2 (mread-uft8 db-ptr length))
      (3 (ieee-floats:decode-float64 (mread-unsigned db-ptr 8)))
      (4 (bytes-from-foreign db-ptr length))
      ((5 6 9 10) (mread-unsigned db-ptr length))
      (7 (mread-map db-ptr length))
      (8 (mread-int32 db-ptr length))
      (11 (mread-list db-ptr length))
      (14 (< 0 length)) ; bool
      (15 (ieee-floats:decode-float32 (mread-unsigned db-ptr 4)))
      )))

(defun read-node-record (mmdb node-number bit)
  "Read left (bit=0) or right (bit=1) record from node"
  (with-slots (ptr record-size) mmdb
    (let* ((node-bytes (/ record-size 4))
           (record-bytes (floor record-size 8))
           (node-offset (* node-bytes node-number)))
      (if (evenp node-bytes)
          ;; Simple case: records align to byte boundaries
          (mem-uint (inc-pointer ptr (+ node-offset (* bit record-bytes))) record-bytes)
          ;; Complex case: records share a byte
          (let* ((base-value (mem-uint (inc-pointer ptr (+ node-offset (* bit (1+ record-bytes)))) record-bytes))
                 (shared-byte (mem-ref (inc-pointer ptr (+ node-offset record-bytes)) :uchar))
                 (high-bits (if (zerop bit) (ash shared-byte -4) (logand shared-byte #xf))))
            (+ (ash high-bits (* record-bytes 8)) base-value))))))

(defun record-value (mmdb ip-bits)
  (with-slots (node-count) mmdb
    (loop for start = 0 then node
          for bit in ip-bits
          for node = (read-node-record mmdb start bit)
          when (= node node-count)
            return nil
          when (> node node-count)
            return node)))

(defun lookup-ip-bits (mmdb ip-bits)
  (with-slots (ptr data-offset node-count) mmdb
    (alexandria:when-let* ((record (record-value mmdb ip-bits))
                           (offset (+ data-offset (- record node-count 16))))
      (mread-data
       (make-db-reader :db-ptr (inc-pointer ptr offset)
                       :head-ptr ptr :data-offset data-offset)))))

(defun integer-to-bits (n bit-count)
  "Convert integer to list of bits (MSB first)"
  (loop for i from (1- bit-count) downto 0
        collect (if (logbitp i n) 1 0)))

(defun query-ip (mmdb ip-string)
  (destructuring-bind (ip-v . ip-val) (parse-ip ip-string)
    (with-slots (ip-version) mmdb
      (when (and (= ip-version 4) (= ip-v 6))
        (error "Can't query IPv6 address in IPv4 database"))
      (lookup-ip-bits mmdb (integer-to-bits ip-val (if (= ip-version 4) 32 128))))))

(defun mmap->mmdb (file ptr fd size)
  (let* ((metadata-start (find-metadata-start ptr size))
         (db-reader (make-db-reader :db-ptr (inc-pointer ptr metadata-start)))
         (metadata (make-metadata (mread-data db-reader))))
    (with-slots (record-size node-count ip-version) metadata
      (make-maxmind-database
       :filename file
       :ptr ptr :fd fd :size size
       :metadata metadata
       :node-count node-count
       :record-size record-size
       :ip-version ip-version
       ;; :data-offset (+ 16 (* (/ (* 2 record-size) 8) node-count))
       :data-offset (+ 16 (/ (* record-size node-count) 4))))))

(defun make-mmdb (file)
  (multiple-value-bind (ptr fd size) (mmap:mmap file)
    (mmap->mmdb file ptr fd size)))

(defun close-mmdb (mmdb)
  (with-slots (ptr fd size) mmdb
    (mmap:munmap ptr fd size)))

(defmacro with-mmdb ((mmdb file) &body body)
  `(mmap:with-mmap (ptr fd size ,file)
     (let ((,mmdb (mmap->mmdb ,file ptr fd size)))
       ,@body)))