#lang scheme/base
;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; DBD-MEMCACHED - DBI interface to memcached in both single & multi-instance mode.
;;
;; Bonzai Lab, LLC.  All rights reserved.
;;
;; Licensed under LGPL.
;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; memcached.ss - implementation of a single-instance memcached client
;; yc 9/23/2009 - first version
(require (planet bzlib/base) 
         (planet bzlib/net)
         scheme/list)

(define (cas? v)
  (number? v))

(define (display-line out fmt . args) 
  (display (apply format (string-append fmt "\r\n") args)
           out)
  (flush-output out))

(define (display/noreply out fmt noreply? . args)
  (apply display-line out (if noreply? (string-append fmt " noreply") fmt) args))

(define (cmd-store! out type/cas key flags exp-time bytes (noreply? #f))
  (case type/cas 
    ((set add replace append prepend)
     (display/noreply out "~a ~a ~a ~a ~a" noreply? type/cas key flags exp-time (bytes-length bytes)))
    (else
     (display/noreply out "cas ~a ~a ~a ~a ~a" noreply? key flags exp-time (bytes-length bytes) type/cas)))
  (display-line out "~a" bytes))

;; anything else is considered an error...
(define (response-store in (noreply? #f))
  (if (not noreply?)
      (let ((ln (read-line in 'return-linefeed)))
        (string->symbol (string-downcase ln)))
      'stored))

;; handling the store command!
(define (store! in out type/cas key value #:flags (flags 0) #:exp-time (exp-time 0) #:noreply? (noreply? #f))
    (cmd-store! out type/cas key flags exp-time value noreply?)
    (let ((response (response-store in noreply?)))
      (case response
        ((stored not_stored) response)
        ((exists not_found)
         (if (cas? type/cas) 
             (error 'store-cas "~a ~a" key response)
             (error 'store "invalid response: ~a" response)))
        (else (error 'store "invalid response: ~a" response)))))


(define (memcached-set! client key value #:flags (flags 0) #:exp-time (exp-time 0) #:noreply? (noreply? #f))
  (store! (client-in client) (client-out client) 'set key value #:flags flags #:exp-time exp-time #:noreply? noreply?))

(define (memcached-add! client key value #:flags (flags 0) #:exp-time (exp-time 0) #:noreply? (noreply? #f))
  (store! (client-in client) (client-out client) 'add key value #:flags flags #:exp-time exp-time #:noreply? noreply?))

(define (memcached-replace! client key value #:flags (flags 0) #:exp-time (exp-time 0) #:noreply? (noreply? #f))
  (store! (client-in client) (client-out client) 'replace key value #:flags flags #:exp-time exp-time #:noreply? noreply?))

(define (memcached-append! client key value #:flags (flags 0) #:exp-time (exp-time 0) #:noreply? (noreply? #f))
  (store! (client-in client) (client-out client) 'append key value #:flags flags #:exp-time exp-time #:noreply? noreply?))

(define (memcached-prepend! client key value #:flags (flags 0) #:exp-time (exp-time 0) #:noreply? (noreply? #f))
  (store! (client-in client) (client-out client) 'prepend key value #:flags flags #:exp-time exp-time #:noreply? noreply?))

(define (memcached-cas! client key value cas #:flags (flags 0) #:exp-time (exp-time 0) #:noreply? (noreply? #f))
  (store! (client-in client) (client-out client) cas key value #:flags flags #:exp-time exp-time #:noreply? noreply?))

;; ;;;;;;;;
;; retrieve API
;; generate the request.
(define (cmd-get out type keys) 
  (define (keys-helper)
    (let ((out (format "~a" keys)))
      (substring out 1 (sub1 (string-length out)))))
  (display-line out "~a ~a" (case type
                              ((get gets) type)
                              (else (error 'cmd-get "unknown get type: ~a" type)))
                (keys-helper)))

;; parse the response.
(define (response-get in)
  (define (end? ln)
    (string-ci=? ln "end"))
  (define (value? ln)
    (regexp-match #px"^(?i:VALUE) ([^\\s]+) (\\d+) (\\d+)\\s?(\\d+)?$" ln))
  (define (helper acc)
    (let ((ln (read-line in 'return-linefeed)))
      (if (or (eof-object? ln) (end? ln)) ;; we are done.
          (reverse acc)
          (if-it (value? ln) 
                 (let ((key (string->symbol (second it)))
                       (flags (string->number (third it)))
                       (len (string->number (fourth it)))
                       (cas (if-it (fifth it)
                                   (string->number it)
                                   #f)))
                   (let ((bytes (read-bytes len in))) 
                     (read-line in 'return-linefeed) ;; remove the last \r\n
                     (helper (cons (list key bytes flags cas) acc))))
                 (error 'response-get "Invalid response: ~a" ln)))))
  (helper '()))

(define (get in out type keys) 
  (cmd-get out type keys) 
  (response-get in)) 
;;(trace get)

(define (memcached-get client key . keys)
  (get (client-in client) (client-out client) 'get (cons key keys)))

(define (memcached-gets client key . keys)
  (get (client-in client) (client-out client) 'gets (cons key keys)))

;; deletion API.
;;

(define (cmd-delete! out key (time 0) (noreply? #f))
  (display/noreply out "delete ~a ~a" noreply? key time))

(define (response-delete! in (noreply? #f))
  (if (not noreply?)
      (let ((ln (read-line in 'return-linefeed)))
        (string->symbol (string-downcase ln)))
      'deleted))

(define (delete! in out key (time 0) (noreply? #f))
  (cmd-delete! out key time noreply?)
  (let ((resp (response-delete! in noreply?)))
    (case resp
      ((deleted not_found) resp)
      (else (error 'delete! "invalid response: ~a" resp)))))

(define (memcached-delete! client key (time 0) (noreply? #f))
  (delete! (client-in client) (client-out client) key time noreply?))

;; incr/decr API
(define (cmd-incr! out type key val (noreply? #f))
  (display/noreply out "~a ~a ~a" noreply? 
                   (case type ((incr decr) type) 
                     (else (error 'cmd-incr! "invalid cmd ~a" type))) 
                   key val))

(define (response-incr! in (noreply? #f))
  (if (not noreply?)
      (let ((ln (read-line in 'return-linefeed)))
        (if (string-ci=? "not_found" ln) 'not_found
            (string->number ln)))
      (void)))

(define (incr! in out type key val (noreply? #f))
  (cmd-incr! out type key val noreply?)
  (response-incr! in noreply?)) 

(define (memcached-incr! client key val (noreply? #f))
  (incr! (client-in client) (client-out client) 'incr key val noreply?))

(define (memcached-decr! client key val (noreply? #f))
  (incr! (client-in client) (client-out client) 'decr key val noreply?))

;; flush-all
(define (cmd-flush-all! out (time 0) (noreply? #f))
  (display/noreply out "flush_all ~a" noreply? time))

(define (response-flush-all! in (noreply? #f))
  (if (not noreply?)
      (let ((ln (read-line in 'return-linefeed)))
        (if (string-ci=? ln "ok") 'ok
            (error 'flush-all! "invalid response: ~a" ln)))
      'ok)) 

(define (flush-all! in out (time 0) (noreply? #f))
  (cmd-flush-all! out time noreply?)
  (response-flush-all! in noreply?)) 

(define (memcached-flush-all! client (time 0) (noreply? #f))
  (flush-all! (client-in client) (client-out client) time noreply?)) 

;; version
(define (cmd-version out)
  (display-line out "version"))

(define (response-version in)
  (let ((ln (read-line in 'return-linefeed)))
    (let ((v (regexp-match #px"^(?i:version) (.+)$" ln)))
      (if (not v) #f
          (cadr v)))))

(define (version in out)
  (cmd-version out)
  (response-version in))

(define (memcached-version client)
  (version (client-in client) (client-out client)))

(define (key? k)
  (cond ((or (bytes? k) (string? k)) ;; must not have space...
         (regexp-match #px"^[^\\s]+$" k))
        ((symbol? k)
         (key? (symbol->string k)))
        (else #f)))

(define (flags? f)
  (and (integer? f) (<= 0 f 65536)))

(define storage-api/c
  (->* (client? key? bytes?)
                          (#:flags flags? #:exp-time exact-nonnegative-integer?
                                   #:noreply? boolean?)
                          any))

(provide (rename-out (client-connect memcached-connect)
                     (client-disconnect memcached-disconnect)))

;; statistics API.
(provide/contract 
 (key? (-> any/c any))
 (flags? (-> any/c any))
 (memcached-set! storage-api/c)
 (memcached-add! storage-api/c)
 (memcached-replace! storage-api/c)
 (memcached-append! storage-api/c)
 (memcached-prepend! storage-api/c)
 (memcached-cas! (->* (client? key? bytes? exact-nonnegative-integer?)
                          (#:flags flags? #:exp-time exact-nonnegative-integer?
                                   #:noreply? boolean?)
                          any))
 (memcached-get (->* (client? key?)
                     ()
                     #:rest (listof key?)
                     any))
 (memcached-gets (->* (client? key?)
                      ()
                      #:rest (listof key?)
                      any))
 (memcached-delete! (->* (client? key?)
                         (exact-nonnegative-integer? boolean?)
                         any))
 (memcached-incr! (->* (client? key? bytes?)
                       (boolean?)
                       any))
 (memcached-decr! (->* (client? key? bytes?)
                       (boolean?)
                       any))
 (memcached-flush-all! (->* (client?) 
                            (exact-nonnegative-integer? boolean?)
                            any))
 (memcached-version (-> client? any))
 )


(define (test-memcached-client)
  (let ((c (client-connect "localhost" 11211)))
    (begin0 ;; (memcached-set! c 'test-me #"this is a test data")
      ;; (memcached-get c 'test-me 'hello 'world)
      ;; (memcached-delete! c 'test-me)
      ;; (memcached-version c)
      (memcached-flush-all! c)
      (client-disconnect c))))