b(in-package :cl-user)
(defpackage #:ocm.osm
  (:use :cl :alexandria))
(in-package :ocm.osm)

(defun next-power-of-two (number)
  (ash 1 (ceiling (log number 2))))

(defvar *read-buffer* (proto:make-byte-vector 4096))
(defun ensure-buffer (size)
  (when (> size (array-total-size *read-buffer*))
    (setf *read-buffer* (proto:make-byte-vector (next-power-of-two size)))))

(defmacro read-big-endian (stream &optional (sizeof 4))
  (with-gensyms (i unsigned-value)
    `(let ((,unsigned-value 0))
       (dotimes (,i ,sizeof ,unsigned-value)
         (setf ,unsigned-value (+ (* ,unsigned-value #x100)
                                  (read-byte ,stream)))))))

(defmacro read-little-endian (stream &optional (sizeof 4))
  (with-gensyms (i unsigned-value)
    `(let ((,unsigned-value 0))
       (dotimes (,i ,sizeof ,unsigned-value)
         (setf ,unsigned-value (+ ,unsigned-value
                                  (ash (read-byte ,stream )
                                       (* 8 ,i))))))))

(defvar *stream* nil "PBF file stream")
(defvar *blob-header-size*  nil "Current header size during file parsing")
(defvar *blob-header* nil "Current header")
(defvar *blob-size*  nil "Current blob size during file parsing")
(defvar *blob-type*  nil "Current blob type")
(defvar *blob* nil "Current blob")
(defvar *block-size* nil "Current block size")
(defvar *header-block* nil "File header block")
(defvar *primitive-block* nil "Current primitive block")
(defparameter +blob-type-osm-header+ "OSMHeader")
(defparameter +blob-type-osm-data+ "OSMData")

(defun read-proto (type size &optional no-read)
  (unless no-read
    (ensure-buffer size)
    (read-sequence *read-buffer* *stream* :end size))
  (proto:deserialize-object type *read-buffer* 0 size))

(defmacro iterate-osm ((filespec) &body body)
  `(with-open-file (*stream* ,filespec :element-type 'unsigned-byte)
     (handler-case
         (loop
            (let* ((*blob-header* (read-proto 'osmpbf:blob-header (read-big-endian *stream* 4)))
                   (*blob-type* (osmpbf:type *blob-header*))
                   (*blob-size* (osmpbf:datasize *blob-header*)))
              (if (or (equal *blob-type* +blob-type-osm-data+)
                      (equal *blob-type* +blob-type-osm-header+))
                  (let* ((*blob* (read-proto 'osmpbf:blob *blob-size*))
                         (*block-size* (unpack-block *blob*)))
                    (if (equal *blob-type* +blob-type-osm-header+)
                        (ensure-features (read-proto 'osmpbf:header-block *block-size* t))
                        (let ((*primitive-block* (read-proto 'osmpbf:primitive-block *block-size* t)))
                          ,@body)))
                  (file-position *stream* (+ (file-position *stream*) *blob-size*)))))
       (end-of-file (e)
         (declare (ignore e))))))

(defun degran-lat (lat block)
  (+ (osmpbf:lat-offset block)
     (* lat (osmpbf:granularity block))))
(defun degran-lon (lon block)
  (+ (osmpbf:lon-offset block)
     (* lon (osmpbf:granularity block))))

(defun sid (indexes string-table)
  (loop for i in indexes
     collect (trivial-utf-8:utf-8-bytes-to-string (elt string-table i))))

(defmacro string-table (block)
  `(osmpbf:s (osmpbf:stringtable ,block)))

(defmacro unsid (object accessor block)
  `(setf (,accessor ,object)
         (sid (,accessor ,object)
              (string-table ,block))))

(defun unsid2 (sids string-table)
  (loop for i from 0
     for sid in sids
     do (setf (elt sids i)
              (trivial-utf-8:utf-8-bytes-to-string (elt string-table sid)))))

(defmacro iterate-nodes ((node block) &body body)
  (with-gensyms (group info)
    `(loop for ,group in (osmpbf:primitivegroup ,block)
        do (loop for ,node in (osmpbf:nodes ,group)
              do (progn
                   (unsid ,node osmpbf:keys ,block)
                   (unsid ,node osmpbf:vals ,block)
                   (when-let (,info (osmpbf:info )))
                   (setf (osmpbf:lat ,node) (degran-lat (osmpbf:lat ,node) ,block)
                         (osmpbf:lon ,node) (degran-lon (osmpbf:lon ,node) ,block))
                   ,@body)))))

(defmacro iterate-dense-nodes ((node block) &body body)
  (with-gensyms (group string-table dense dense-info id lat lon tag-start tag-end info ith key val)
    `(loop for ,group in (osmpbf:primitivegroup ,block)
        with ,string-table = (osmpbf:s (osmpbf:stringtable ,block))
        do (when-let (,dense (osmpbf:dense ,group))
             (undelta (osmpbf:id ,dense))
             (undelta (osmpbf:lat ,dense))
             (undelta (osmpbf:lon ,dense))
             (when-let (,dense-info (osmpbf:denseinfo ,dense))
               (undelta (osmpbf:timestamp ,dense-info))
               (undelta (osmpbf:changeset ,dense-info))
               (undelta (osmpbf:uid ,dense-info))
               (undelta (osmpbf:user-sid ,dense-info)))
             (loop
                for ,id in (osmpbf:id ,dense)
                for ,lat in (osmpbf:lat ,dense)
                for ,lon in (osmpbf:lon ,dense)
                for ,tag-start = 0 then ,tag-end
                for ,tag-end = (position 0 (osmpbf:keys-vals ,dense) :start ,tag-start)
                for ,ith from 0
                do (let* ((,info (when-let (,dense-info (osmpbf:denseinfo ,dense))
                                   (make-instance 'osmpbf:info
                                                  :version (elt (osmpbf:version ,dense-info) ,ith)
                                                  :timestamp (elt (osmpbf:timestamp ,dense-info) ,ith)
                                                  :changeset (elt (osmpbf:changeset ,dense-info) ,ith)
                                                  :uid (elt (osmpbf:uid ,dense-info) ,ith)
                                                  :user-sid (car (sid (list (elt (osmpbf:user-sid ,dense-info) ,ith))
                                                                      ,string-table)))))
                          (,node (make-instance 'osmpbf:node
                                                :id ,id
                                                :keys (sid (loop for ,key
                                                              in (subseq (osmpbf:keys-vals ,dense)
                                                                         ,tag-start ,tag-end)
                                                              by #'cddr
                                                              collect ,key)
                                                           ,string-table)
                                                :vals (sid (loop for ,val
                                                              in (subseq (osmpbf:keys-vals ,dense)
                                                                         (1+ ,tag-start) ,tag-end)
                                                              by #'cddr
                                                              collect ,val)
                                                           ,string-table)
                                                :lat (degran-lat ,lat ,block)
                                                :lon (degran-lon ,lon ,block)
                                                :info ,info)))
                     ,@body))))))

(defmacro iterate-dense ((id lat lon keys vals block) &body body)
  (with-gensyms (group string-table dense tag-start tag-end key val id-cur lat-cur lon-cur id-prev lat-prev lon-prev lat-acc lon-acc)
    `(loop for ,group in (osmpbf:primitivegroup ,block)
        with ,string-table = (string-table ,block)
        do (when-let (,dense (osmpbf:dense ,group))
             (loop
                for ,id-cur in (osmpbf:id ,dense)
                for ,lat-cur in (osmpbf:lat ,dense)
                for ,lon-cur in (osmpbf:lon ,dense)
                for ,id-prev = 0 then ,id
                for ,lat-prev = 0 then ,lat-acc
                for ,lon-prev = 0 then ,lon-acc
                for ,tag-start = 0 then ,tag-end
                for ,tag-end = (position 0 (osmpbf:keys-vals ,dense) :start ,tag-start)
                for ,id = (+ ,id-cur ,id-prev)
                for ,lat-acc = (+ ,lat-cur ,lat-prev)
                for ,lon-acc = (+ ,lon-cur ,lon-prev)
                for ,lat = (degran-lat ,lat-acc ,block)
                for ,lon = (degran-lon ,lon-acc ,block)
                do (multiple-value-bind (,keys ,vals)
                       (loop for (,key ,val) on (subseq (osmpbf:keys-vals ,dense) ,tag-start ,tag-end) by #'cddr
                          collecting ,key into ,keys
                          collecting ,val into ,vals
                          finally (return (values ,keys ,vals)))
                     (let ((,keys (sid ,keys ,string-table))
                           (,vals (sid ,vals ,string-table)))
                       ,@body)))))))

(defmacro iterate-ways ((way block) &body body)
  (with-gensyms (group)
    `(loop for ,group in (osmpbf:primitivegroup ,block)
        do (loop for ,way in (osmpbf:ways ,group)
              do (progn
                   (unsid ,way osmpbf:keys ,block)
                   (unsid ,way osmpbf:vals ,block)
                   (undelta (osmpbf:refs ,way))
                   ,@body)))))

(defmacro iterate-rels ((rel block) &body body)
  (with-gensyms (group string-table)
    `(loop for ,group in (osmpbf:primitivegroup ,block)
        with ,string-table = (string-table ,block)
        do (loop for ,rel in (osmpbf:relations ,group)
              do (progn
                   (unsid2 (osmpbf:keys ,rel) ,string-table)
                   (unsid2 (osmpbf:vals ,rel) ,string-table)
                   (unsid2 (osmpbf:roles-sid ,rel) ,string-table)
                   (undelta (osmpbf:memids ,rel))
                   ,@body)))))

(defun get-tags (object string-table)
  (loop for key in (osmpbf:keys object)
     for val in (osmpbf:vals object)
     collect  (cons (trivial-utf-8:utf-8-bytes-to-string (elt string-table key))
                    (trivial-utf-8:utf-8-bytes-to-string (elt string-table val)))))

(defun delta-decode (seq)
  (loop for elt in seq
     for prev = 0 then cur
     for cur = (+ elt prev)
     collect cur))

(defmacro iterate-dense2 ((id tags lat lon block) &body body)
  (with-gensyms (group string-table dense tag-start tag-end key val id-cur lat-cur lon-cur id-prev lat-prev lon-prev lat-acc lon-acc key-vals clo-start clo-end)
    `(loop for ,group in (osmpbf:primitivegroup ,block)
        with ,string-table = (string-table ,block)
        do (when-let (,dense (osmpbf:dense ,group))
             (loop
                with ,key-vals = (osmpbf:keys-vals ,dense)
                for ,id-cur in (osmpbf:id ,dense)
                for ,lat-cur in (osmpbf:lat ,dense)
                for ,lon-cur in (osmpbf:lon ,dense)
                for ,id-prev = 0 then ,id
                for ,lat-prev = 0 then ,lat-acc
                for ,lon-prev = 0 then ,lon-acc
                for ,tag-start = 0 then ,tag-end
                for ,tag-end = (position 0 ,key-vals :start ,tag-start)
                for ,id = (+ ,id-cur ,id-prev)
                for ,lat-acc = (+ ,lat-cur ,lat-prev)
                for ,lon-acc = (+ ,lon-cur ,lon-prev)
                for ,lat = (degran-lat ,lat-acc ,block)
                for ,lon = (degran-lon ,lon-acc ,block)
                for ,tags = (unless (= ,tag-start ,tag-end)
                              (let ((,clo-start ,tag-start)
                                    (,clo-end ,tag-end))
                                (lambda ()
                                  (loop for (,key ,val) on (subseq ,key-vals ,clo-start ,clo-end) by #'cddr
                                     collect (cons (trivial-utf-8:utf-8-bytes-to-string
                                                    (elt ,string-table ,key))
                                                   (trivial-utf-8:utf-8-bytes-to-string
                                                    (elt ,string-table ,val)))))))
                do (progn ,@body))))))

(defmacro iterate-ways2 ((id tags refs block) &body body)
  (with-gensyms (group string-table object o-string-table)
    `(loop for ,group in (osmpbf:primitivegroup ,block)
        with ,string-table = (string-table ,block)
        do (loop for ,object in (osmpbf:ways ,group)
              for ,id = (osmpbf:id ,object)
              for ,o-string-table = ,string-table
              for ,tags = (lambda () (get-tags ,object ,o-string-table))
              for ,refs = (lambda () (delta-decode (osmpbf:refs ,object)))
              do (progn ,@body)))))

(defmacro iterate-rels2 ((id tags roles memids types block) &body body)
  (with-gensyms (group string-table relation)
    `(loop for ,group in (osmpbf:primitivegroup ,block)
        with ,string-table = (string-table ,block)
        do (loop for ,relation in (osmpbf:relations ,group)
              for ,id = (osmpbf:id ,relation)
              for ,tags = (lambda () (get-tags ,relation ,string-table))
              for ,roles = (lambda () (sid (osmpbf:roles-sid ,relation) ,string-table))
              for ,memids = (lambda () (delta-decode (osmpbf:memids ,relation)))
              for ,types = (lambda () (osmpbf:types ,relation))
              do (progn ,@body)))))

(defun undelta (seq)
  (when (> (length seq) 1)
    (loop for i from 1 below (length seq)
       do (incf (elt seq i)
                (elt seq (1- i)))))
  seq)

(defparameter +supported-features+ '("OsmSchema-V0.6" "DenseNodes"))
(defun ensure-features (header)
  (when-let (missing-features (set-difference (osmpbf:required-features header)
                                              +supported-features+
                                              :test 'equal))
    (error "Features ~A not supported" missing-features))
  (setf *header-block* header))

(defun unpack-block (blob)
  (if-let (uncompressed (osmpbf:raw blob))
    (let ((size (array-total-size uncompressed)))
      (ensure-buffer size)
      (dotimes (i (array-total-size uncompressed))
        (setf (row-major-aref *read-buffer* i)
              (row-major-aref uncompressed i)))
      size)
    (if-let (zlib-data (osmpbf:zlib-data blob))
      (let ((size (osmpbf:raw-size blob)))
        (ensure-buffer size)
        (chipz:decompress *read-buffer* (chipz:make-dstate :zlib) zlib-data)
        size)
      (error "Unsupported block compression"))))

(defun type-match (type-expr type)
  (etypecase type-expr
    (null t)
    (symbol (eql type-expr type))
    (list (member type type-expr))))

(defun expr-match (template value)
  (etypecase template
    (null t)
    (string (equal template value))
    (list (destructuring-bind (type expr) template
            (let ((pos (search expr value)))
              (ecase type
                (:pre (eql pos 0))
                (:suf (eql pos (max 0 (- (length value)
                                         (length expr)))))
                (:inf (not (null pos)))
                (:eql (and (eql pos 0)
                           (= (length value)
                              (length expr))))))))))

(defun query-match (query type keys vals)
  (or (null query)
      (loop for expr in query
         when (destructuring-bind (tp &optional keyp valp) expr
                (and
                 (type-match tp type)
                 (or (and (null keyp) (null valp))
                     (loop for key in keys
                        for val in vals
                        when (and (expr-match keyp key)
                                  (expr-match valp val))
                        return t))))
         return t)))

(defun query-match2 (query type tags-thunk)
  (or (null query)
      (loop for expr in query
         when (destructuring-bind (tp &optional keyp valp) expr
                (and
                 (type-match tp type)
                 (or (and (null keyp) (null valp))
                     (and tags-thunk
                          (loop for (key . val) in (if (functionp tags-thunk) (funcall tags-thunk) tags-thunk)
                             when (and (expr-match keyp key)
                                       (expr-match valp val))
                             return t)))))
         return t)))

(defun query-has-type (query type)
  (or (null query)
      (loop for expr in query
         when (type-match (car expr) type)
         return t)))

(defmacro iterate ((object filespec tags-query) &body body)
  (let ((has-rels (query-has-type (cadr tags-query) :rel))
        (has-ways (query-has-type (cadr tags-query) :way))
        (has-nodes (query-has-type (cadr tags-query) :node)))
    (with-gensyms (id lat lon keys vals query)
      `(let ((,query ,tags-query))
         (iterate-osm (,filespec)
           ,@(when has-rels `((iterate-rels (,object *primitive-block*)
                                (when (query-match ,query :rel (osmpbf:keys ,object) (osmpbf:vals ,object))
                                  ,@body))))
           ,@(when has-ways `((iterate-ways (,object *primitive-block*)
                                (when (query-match ,query :way (osmpbf:keys ,object) (osmpbf:vals ,object))
                                  ,@body))))
           ,@(when has-nodes
               `((iterate-nodes (,object *primitive-block*)
                   (when (query-match ,query :node (osmpbf:keys ,object) (osmpbf:vals ,object))
                     ,@body))
                 (iterate-dense (,id ,lat ,lon ,keys ,vals *primitive-block*)
                   (when (query-match ,query :node ,keys ,vals)
                     (let ((,object (make-instance 'osmpbf:node :id ,id :lat ,lat :lon ,lon
                                                   :keys ,keys :vals ,vals)))
                       ,@body))))))))))

(defstruct obj id tags)
(defstruct (node (:include obj)) lat lon)
(defstruct (way (:include obj)) refs)
(defstruct (rel (:include obj)) roles memids types)

(defmacro iterate2 ((object filespec query) &body body)
  (let ((has-rels (query-has-type (cadr query) :rel))
        (has-ways (query-has-type (cadr query) :way))
        (has-nodes (query-has-type (cadr query) :node)))
    (with-gensyms (g-query g-id g-tags g-roles g-memids g-types g-refs g-lat g-lon)
      `(let ((,g-query ,query))
         (iterate-osm (,filespec)
           ,@(when has-rels `((iterate-rels2 (,g-id ,g-tags ,g-roles ,g-memids ,g-types *primitive-block*)
                                (when (query-match2 ,g-query :rel ,g-tags)
                                  (let ((,object (make-rel :id ,g-id :tags ,g-tags
                                                           :roles ,g-roles :memids ,g-memids :types ,g-types)))
                                    ,@body)))))
           ,@(when has-ways `((iterate-ways2 (,g-id ,g-tags ,g-refs *primitive-block*)
                                (when (query-match2 ,g-query :way ,g-tags)
                                  (let ((,object (make-way :id ,g-id :tags ,g-tags
                                                           :refs ,g-refs)))
                                    ,@body)))))
           ,@(when has-nodes `((iterate-dense2 (,g-id ,g-tags ,g-lat ,g-lon *primitive-block*)
                                 (when (query-match2 ,g-query :node ,g-tags)
                                   (let ((,object (make-node :id ,g-id :tags ,g-tags
                                                             :lat ,g-lat :lon ,g-lon)))
                                     ,@body))))))))))

(defparameter +max-rel-pass+ 5 "max rel-2-rel passes")
(defmacro with-related ((object filespec query) &body body)
  (with-gensyms (g-object g-filespec g-objects g-idx-rels g-idx-ways g-idx-nodes g-rel g-type g-memid g-indexed
                          g-dirty g-try g-indexed-list)
    `(let ((,g-filespec ,filespec)
           ,g-objects
           (,g-idx-rels (make-hash-table))
           (,g-idx-ways (make-hash-table))
           (,g-idx-nodes (make-hash-table)))
       (print "iterate queried")
       (iterate2 (,g-object ,g-filespec ,query)
         (typecase ,g-object
           (node (let ((,object ,g-object)) ,@body))
           (way (push ,g-object ,g-objects)
                (setf (obj-tags ,g-object) (funcall (obj-tags ,g-object))
                      (way-refs ,g-object) (funcall (way-refs ,g-object)))
                (loop for ,g-rel in (way-refs ,g-object)
                   do (push ,g-object (gethash ,g-rel ,g-idx-nodes))))
           (rel (push ,g-object ,g-objects)
                (setf (obj-tags ,g-object) (funcall (obj-tags ,g-object))
                      (rel-roles ,g-object) (funcall (rel-roles ,g-object))
                      (rel-memids ,g-object) (funcall (rel-memids ,g-object))
                      (rel-types ,g-object) (funcall (rel-types ,g-object)))
                (loop for ,g-memid in (rel-memids ,g-object)
                   for ,g-type in (rel-types ,g-object)
                   do (push ,g-object (gethash ,g-memid (ecase ,g-type
                                                          (:node ,g-idx-nodes)
                                                          (:way ,g-idx-ways)
                                                          (:relation ,g-idx-rels))))))))
       (when (plusp (hash-table-count ,g-idx-rels))
         (loop
            for ,g-dirty = nil
            for ,g-try from 1 upto +max-rel-pass+
            do (format t "iterate rels pass ~a~%" ,g-try)
            do (iterate2 (,g-object ,g-filespec '((:rel)))
                 (when-let ((,g-indexed-list (gethash (obj-id ,g-object) ,g-idx-rels)))
                   (setf (rel-memids ,g-object) (funcall (rel-memids ,g-object))
                         (rel-types ,g-object) (funcall (rel-types ,g-object)))
                   (loop for ,g-indexed in ,g-indexed-list
                      do (loop
                            for ,g-memid on (rel-memids ,g-indexed)
                            for ,g-type in (rel-types ,g-indexed)
                            when (and (eql ,g-type :relation)
                                      (equal (car ,g-memid) (obj-id ,g-object)))
                            do (setf (car ,g-memid) ,g-object) and ;; replace memid ref with current rel
                            do (loop for ,g-memid in (rel-memids ,g-object) ;; add it's references to index
                                  for ,g-type in (rel-types ,g-object)
                                  do (push ,g-object
                                           (gethash ,g-memid (ecase ,g-type
                                                               (:node ,g-idx-nodes)
                                                               (:way ,g-idx-ways)
                                                               (:relation (prog1 ,g-idx-rels
                                                                            (setf ,g-dirty t)))))))))
                   (remhash (obj-id ,g-object) ,g-idx-rels))) ;; remove processed rel from index
            while ,g-dirty))
       (when (plusp (hash-table-count ,g-idx-ways))
         (print "iterate ways")
         (iterate2 (,g-object ,g-filespec '((:way)))
           (when-let ((,g-indexed-list (gethash (obj-id ,g-object) ,g-idx-ways)))
             (setf (way-refs ,g-object) (funcall (way-refs ,g-object)))
             (loop for ,g-indexed in ,g-indexed-list
                do (loop
                      for ,g-memid on (rel-memids ,g-indexed)
                      for ,g-type in (rel-types ,g-indexed)
                      when (and (eql ,g-type :way)
                                (equal (car ,g-memid) (obj-id ,g-object)))
                      do (setf (car ,g-memid) ,g-object) and ;; replace memid ref with current way
                      do (loop for ,g-rel in (way-refs ,g-object) ;; add it's refs to index
                            do (push ,g-object (gethash ,g-rel ,g-idx-nodes)))))
             (remhash (obj-id ,g-object) ,g-idx-ways)))) ;; remove processed way from index
       (when (plusp (hash-table-count ,g-idx-nodes))
         (print "iterate nodes")
         (iterate2 (,g-object ,g-filespec '((:node)))
           (when-let ((,g-indexed-list (gethash (obj-id ,g-object) ,g-idx-nodes)))
             (loop for ,g-indexed in ,g-indexed-list
                  do (etypecase ,g-indexed
                       (rel (loop
                               for ,g-memid on (rel-memids ,g-indexed)
                               for ,g-type in (rel-types ,g-indexed)
                               when (and (eql ,g-type :node)
                                         (equal (car ,g-memid) (obj-id ,g-object)))
                               do (setf (car ,g-memid) ,g-object))) ;; replace rel's memid  with current node
                       (way (nsubstitute ,g-object (obj-id ,g-object) (way-refs ,g-indexed))))) ;; replace way's ref with current node
             (remhash (obj-id ,g-object) ,g-idx-nodes))))
       (loop for ,object in (nreverse ,g-objects) ;; Iterate original matched objects
          do (progn ,@body)))))