#lang racket

(require "iputil.rkt")

(provide make-rt rt-update! rt-lookup rt-dump)

(struct rt-node [edge0 edge1 data] #:transparent #:mutable)
(struct rt-edge [label target] #:transparent)

(define empty-node-data (gensym 'MEOW))

(define (make-rt)
  (rt-node #f #f empty-node-data))

;; Helper functions for rt-node struct access based on whether the edge is 1 or 0
(define (rt-getter bit) (if bit rt-node-edge1 rt-node-edge0))
(define (rt-setter bit) (if bit set-rt-node-edge1! set-rt-node-edge0!))

(define (bl-common-len bl1 bl2)
  (for/sum ([b1 (in-list bl1)]
            [b2 (in-list bl2)])
       #:break (not (equal? b1 b2))
       1))

(define (rt-partial-iterate node key [visited '()])
  (cond
    [(empty? key) (list 'exact node)]
    [else
     (let* ([bit (first key)]
            [getter (rt-getter bit)]
            [next-edge (getter node)]
            [next-label (and next-edge (rt-edge-label next-edge))]
            [next-target (and next-edge (rt-edge-target next-edge))]
            [next-common-len (and next-label (bl-common-len next-label key))])
       (cond
         [(and next-edge (= next-common-len (length next-label)))
          (rt-partial-iterate next-target (drop key next-common-len) (cons node visited))]
         [next-edge (list 'partial node next-edge next-common-len key (cons node visited))]
         [else (list 'no-match node key (cons node visited))]))]))

(define (rt-update! node key updater failure-result)
  (define (insert-node! node key)
    (let* ([bit (first key)]
           [setter! (rt-setter bit)])
      (setter! node (rt-edge key (rt-node #f #f (updater (failure-result)))))))
  (define (split-node! node key orig-edge prefix-len)
    (let* ([bit (first key)]
           [setter! (rt-setter bit)]
           [orig-label (rt-edge-label orig-edge)]
           [orig-target (rt-edge-target orig-edge)]
           [new-orig-edge (rt-edge (drop orig-label prefix-len) orig-target)]
           [new-insert-edge (rt-edge (drop key prefix-len)
                                     (rt-node #f #f (updater (failure-result))))]
           [diff-bit (list-ref key prefix-len)]
           [common-node (rt-node (if diff-bit new-orig-edge new-insert-edge)
                                 (if diff-bit new-insert-edge new-orig-edge)
                                 empty-node-data)]
           [common-edge (rt-edge (take key prefix-len) common-node)])
      (setter! node common-edge)))
  (match (rt-partial-iterate node key)
    [(list 'exact node)
     (set-rt-node-data!
       node
       (let ([d (rt-node-data node)])
         (updater
           (if (eq? d empty-node-data) (failure-result) d))))]
    [(list 'partial node orig-edge prefix-len partial-key visited)
     (split-node! node partial-key orig-edge prefix-len)]
    [(list 'no-match node partial-key visited)
     (insert-node! node partial-key)]))

(define (rt-lookup node key failure-result)
  (define (find-first-with-data nodes)
    (or (for/first ([node (in-list nodes)]
                    #:when (not (eq? empty-node-data (rt-node-data node))))
          (rt-node-data node))
        (failure-result)))
  (find-first-with-data
   (match (rt-partial-iterate node key)
     [(list 'exact node)
      (list node)]
     [(list _ ... visited)
      visited])))

; (define test (make-rt))
; (define (test-insert! x)
;   (rt-update!
;     test
;     (map (curry equal? #\1) (string->list x))
;     (lambda (_) x)
;     (lambda () x)))
; (test-insert! "0001")
; ; (test-insert! "1000")
; ; (test-insert! "1010")
; (test-insert! "0011")
; (test-insert! "0000")

(define (rt-dump node [prefix ""])
  (displayln (format "~a node ~a" prefix (rt-node-data node)))
  (define edge0 (rt-node-edge0 node))
  (define edge1 (rt-node-edge1 node))
  (when edge0
    (displayln (format "~a  edge0 ~a" prefix (rt-edge-label edge0)))
    (rt-dump (rt-edge-target edge0) (string-append prefix "  ")))
  (when edge1
    (displayln (format "~a  edge1 ~a" prefix (rt-edge-label edge1)))
    (rt-dump (rt-edge-target edge1) (string-append prefix "  "))))
;(rt-dump test)