#lang racket

(provide (struct-out peer) (struct-out subnet)
         string->ip
         ip->string
         ip-netmask
         netmask-ip
         string->subnet
         subnet->string
         string->peer
         subnet->bl
         subnet-flip-last
         subnet-drop-last
         subnet-disaggregate
         ip->bl)

(require racket/struct)

;; Support functions for dealing with IP addresses, subnets, and peers

(module+ test
  (require rackunit))

;; An IP is a Nat

;; A PeerType is one of
;; - 'prov
;; - 'cust
;; - 'peer

;; A Peer is a (peer IP PeerType)
(struct peer [ip type]
  #:transparent
  #:methods gen:custom-write
  [(define write-proc
     (make-constructor-style-printer
      (lambda (x) 'peer)
      (lambda (x) (list (ip->string (peer-ip x)) (peer-type x)))))])

;; A Subnet is a (subnet IP Nat)
(struct subnet [ip mask]
  #:transparent
  #:methods gen:custom-write
  [(define write-proc
     (make-constructor-style-printer
      (lambda (x) 'subnet)
      (lambda (x) (list (subnet->string x)))))])

;; Str -> IP
;; Parses the IP address contained in str
(define (string->ip str)
  (define parts (reverse (string-split str ".")))
  (for/sum ([part (in-list parts)]
            [i (in-naturals)])
    (arithmetic-shift (string->number part)
                      (* i 8))))

;; IP -> Str
;; Serializes the IP address ip to a string
(define (ip->string ip)
  (define parts
    (reverse (for/list ([i (in-range 4)])
               (number->string (bitwise-and (arithmetic-shift ip (* i -8))
                                            255)))))
  (string-join parts "."))

;; IP -> Nat
;; Calculates how many mask bits are set in the given netmask represented as an IP address
;; Example: 255.0.0.0 -> 8
(define (ip-netmask ip)
  (- 32
     (for/sum ([i (in-range 32)])
       #:break (positive? (bitwise-and ip (arithmetic-shift 1 i)))
       1)))

;; Nat -> IP
;; Creates an IP address representing the given netmask size
;; Example: 8 -> 255.0.0.0
(define (netmask-ip mask)
  (arithmetic-shift (sub1 (arithmetic-shift 1 mask))
                    (- 32 mask)))

;; Str -> Subnet
;; Parses a subnet from string
(define (string->subnet str)
  (match-define (list ipstr maskstr) (string-split str "/"))
  (subnet (string->ip ipstr)
          (string->number maskstr)))

;; Subnet -> Str
;; Serializes a subnet to string
(define (subnet->string sub)
  (format "~a/~a"
          (ip->string (subnet-ip sub))
          (subnet-mask sub)))

;; Subnet -> [Listof Bool]
;; Converts a subnet into a list of booleans representing its bits
(define (subnet->bl sub)
  (ip->bl (subnet-ip sub) (subnet-mask sub)))

;; [Listof Bool] -> [Listof Bool]
;; Flips the last bit of the given bit list
(define (bl-flip-last bl)
;; cursed line
  (match bl
    [`(,f ... ,l) `(,@f ,(not l))]))

;; Subnet -> Subnet
;; Provides a subnet with the last bit flipped
(define (subnet-flip-last sub)
  (subnet (bitwise-xor (arithmetic-shift 1 (- 32 (subnet-mask sub))) (subnet-ip sub))
          (subnet-mask sub)))

(define (subnet-drop-last sub)
  (match-define (subnet ip mask) sub)
  (subnet (arithmetic-shift (arithmetic-shift ip (- mask 33)) (- 33 mask))
          (sub1 mask)))

;; Subnet Subnet -> [Listof Subnet]
;; Given an aggregate subnet and a component part, return an optimal list of
;; subnets filling the address space represented by (agg - part)
(define (subnet-disaggregate agg part)
  (define diff (- (subnet-mask part) (subnet-mask agg)))
  (define part-flip (subnet-flip-last part))
  (cond
    [(= 1 diff) (list part-flip)]
    [else (cons part-flip (subnet-disaggregate agg (subnet-drop-last part)))]))

;; IP -> [Listof Bool]
;; Converts an IP address into a list of booleans representing its bits
(define (ip->bl ip [bits 32])
  (for/list ([i (in-range bits)])
        (= 1 (bitwise-and 1 (arithmetic-shift ip (- i 31))))))

;; Str -> Peer
;; Parses a peer from string
(define (string->peer str)
  (match-define (list ip type) (string-split str "-"))
  (peer (string->ip ip) (string->symbol type)))

;; ================================================================================

(module+ test
  (check-equal? (ip->string (string->ip "123.84.0.67")) "123.84.0.67")
  (check-equal? (ip->string (string->ip "0.0.0.0")) "0.0.0.0")

  (check-equal? (ip->string (netmask-ip 23)) "255.255.254.0")
  (check-equal? (ip->string (netmask-ip 8)) "255.0.0.0")

  (check-equal? (ip-netmask (string->ip "255.255.254.0")) 23)
  (check-equal? (ip-netmask (string->ip "255.0.0.0")) 8)
  (check-equal? (subnet->bl (string->subnet "3.3.0.0/24"))
                '(#f #f #f #f #f #f #t #t
                  #f #f #f #f #f #f #t #t
                  #f #f #f #f #f #f #f #f)))