#lang racket/base
(require (for-syntax racket/base)
         racket/contract/base
         unstable/prop-contract
         "private/dict.rkt")

;; ----------------------------------------

(define-values (prop:dict/c dict/c-struct? dict/c-struct-ref)
  (make-struct-type-property 'dict/c #f))

(define-values (prop:dict/contract dict/contract? dict/contract-ref)
  (make-struct-type-property 'dict/contract #f
                             (list (cons prop:dict car)
                                   (cons prop:dict/c cadr))))

;; ----------------------------------------

(define (get-dict/c-type-key-c v)
  (vector-ref v 0))
(define (get-dict/c-type-value-c v)
  (vector-ref v 1))
(define (get-dict/c-type-iter-c v)
  (vector-ref v 2))
(define (get-dict/c-inst-key-c v)
  (vector-ref v 3))
(define (get-dict/c-inst-value-c v)
  (vector-ref v 4))
(define (get-dict/c-inst-iter-c v)
  (vector-ref v 5))

;; ----------------------------------------

(define (dict-key-contract d)
  (if (dict/c-struct? d)
      (let* ([cv (dict/c-struct-ref d)]
             [type-key-contract (get-dict/c-type-key-c cv)]
             [get-inst-key-contract (get-dict/c-inst-key-c cv)])
        (if get-inst-key-contract
            (and/c type-key-contract (get-inst-key-contract d))
            type-key-contract))
      (if (vector? d)
          exact-nonnegative-integer?
          any/c)))

(define (dict-value-contract d)
  (if (dict/c-struct? d)
      (let* ([cv (dict/c-struct-ref d)]
             [type-value-contract (get-dict/c-type-value-c cv)]
             [get-inst-value-contract (get-dict/c-inst-value-c cv)])
        (if get-inst-value-contract
            (and/c type-value-contract (get-inst-value-contract d))
            type-value-contract))
      any/c))

(define (dict-iter-contract d)
  (if (dict/c-struct? d)
      (let* ([cv (dict/c-struct-ref d)]
             [type-iter-contract (get-dict/c-type-iter-c cv)]
             [get-inst-iter-contract (get-dict/c-inst-iter-c cv)])
        (if get-inst-iter-contract
            (and/c type-iter-contract (get-inst-iter-contract d))
            type-iter-contract))
      any/c))

;; ----------------------------------------

(define dict-ref-contract
  (->i ([d dict?] [k (d) (dict-key-contract d)])
       ([default any/c])
       any)) ;; because default can be multi-valued procedure
(define dict-set!-contract
  (->i ([d (and/c dict? dict-mutable?)]
        [k (d) (dict-key-contract d)]
        [value (d) (dict-value-contract d)])
       [_ void?]))
(define dict-set-contract
  (->i ([d (and/c dict? dict-can-functional-set?)]
        [k (d) (dict-key-contract d)]
        [value (d) (dict-value-contract d)])
       [_ dict?]))
(define dict-remove!-contract
  (->i ([d (and/c dict? dict-mutable? dict-can-remove-keys?)]
        [k (d) (dict-key-contract d)])
       [_ void?]))
(define dict-remove-contract
  (->i ([d (and/c dict? dict-can-functional-set? dict-can-remove-keys?)]
        [k (d) (dict-key-contract d)])
       [_ dict?]))
(define dict-count-contract
  (-> dict? exact-nonnegative-integer?))
(define dict-iterate-first-contract
  (->i ([d dict?]) [_ (d) (or/c #f (dict-iter-contract d))]))
(define dict-iterate-next-contract
  (->i ([d dict?] [iter (d) (dict-iter-contract d)])
       [_ (d) (or/c #f (dict-iter-contract d))]))
(define dict-iterate-key-contract
  (->i ([d dict?] [iter (d) (dict-iter-contract d)])
       [_ (d) (dict-key-contract d)]))
(define dict-iterate-value-contract
  (->i ([d dict?] [iter (d) (dict-iter-contract d)])
       [_ (d) (dict-value-contract d)]))

(define prop:dict-contract
  (vector-immutable/c dict-ref-contract
                      (or/c #f dict-set!-contract)
                      (or/c #f dict-set-contract)
                      (or/c #f dict-remove!-contract)
                      (or/c #f dict-remove-contract)
                      dict-count-contract
                      dict-iterate-first-contract
                      dict-iterate-next-contract
                      dict-iterate-key-contract
                      dict-iterate-value-contract))

(define prop:dict/c-contract
  (vector-immutable/c contract?
                      contract?
                      contract?
                      (or/c #f (-> dict? contract?))
                      (or/c #f (-> dict? contract?))
                      (or/c #f (-> dict? contract?))))

(define even-length-list/c
  (flat-named-contract 'even-length-list/c
                       (lambda (l) (even? (length l)))))

;; ----------------------------------------

(provide/contract
 [prop:dict/contract
  (struct-type-property/c
   (list/c prop:dict-contract
           prop:dict/c-contract))]

 [dict?
  (-> any/c boolean?)]
 [dict-mutable?
  (-> dict? boolean?)]
 [dict-can-remove-keys?
  (-> dict? boolean?)]
 [dict-can-functional-set?
  (-> dict? boolean?)]

 [dict-has-key?
  (-> dict? any/c boolean?)]
 [dict-ref
  dict-ref-contract]
 [dict-ref!
  (->i ([d (and/c dict? dict-mutable?)]
        [k (d) (dict-key-contract d)]
        [default (d) (or/c (dict-value-contract d) (-> (dict-value-contract d)))]) ;; use if/c ?
       [_ (d) (dict-value-contract d)])]
 [dict-set!
  dict-set!-contract]
 [dict-set
  dict-set-contract]
 [dict-set*!
  (->i ([d (and/c dict? dict-mutable?)])
       #:rest [_ (d) (let ([key/c (dict-key-contract d)]
                           [val/c (dict-value-contract d)])
                       (letrec ([args/c
                                 (recursive-contract
                                  (or/c null
                                        (cons/c key/c (cons/c val/c args/c))))])
                         (and/c even-length-list/c
                                args/c)))]
       [_ void?])]
 [dict-set*
  (->i ([d (and/c dict? dict-can-functional-set?)])
       #:rest [_ (d) (let ([key/c (dict-key-contract d)]
                           [val/c (dict-value-contract d)])
                       (letrec ([args/c
                                 (recursive-contract
                                  (or/c null
                                        (cons/c key/c (cons/c val/c args/c))))])
                         (and/c even-length-list/c
                                args/c)))]
       [_ dict?])]
 [dict-update!
  (->i ([d (and/c dict? dict-mutable?)]
        [k (d) (dict-key-contract d)]
        [update (d) (-> (dict-value-contract d) (dict-value-contract d))])
       ([default (d) (or/c (dict-value-contract d) (-> (dict-value-contract d)))]) ;; use if/c
       [_ void?])]
 [dict-update
  (->i ([d (and/c dict? dict-can-functional-set?)]
        [k (d) (dict-key-contract d)]
        [update (d) (-> (dict-value-contract d) (dict-value-contract d))])
       ([default (d) (or/c (dict-value-contract d) (-> (dict-value-contract d)))]) ;; use if/c ?
       [_ dict?])]
 [dict-remove!
  dict-remove!-contract]
 [dict-remove
  dict-remove-contract]
 [dict-count
  dict-count-contract]
 [dict-iterate-first
  dict-iterate-first-contract]
 [dict-iterate-next
  dict-iterate-next-contract]
 [dict-iterate-key
  dict-iterate-key-contract]
 [dict-iterate-value
  dict-iterate-value-contract]

 [dict-map
  (->i ([d dict?] [proc (d) (-> (dict-key-contract d) (dict-value-contract d) any)])
       [_ list?])]
 [dict-for-each
  (->i ([d dict?] [proc (d) (-> (dict-key-contract d) (dict-value-contract d) any)])
       [_ void?])]

 [dict-keys
  (->i ([d dict?])
       [_ (d) (listof (dict-key-contract d))])]
 [dict-values
  (->i ([d dict?])
       [_ (d) (listof (dict-value-contract d))])]
 [dict->list
  (->i ([d dict?])
       [_ (d) (listof (cons/c (dict-key-contract d) (dict-value-contract d)))])])

(provide prop:dict

         make-custom-hash
         make-immutable-custom-hash
         make-weak-custom-hash

         in-dict
         in-dict-keys
         in-dict-values
         in-dict-pairs

         dict-key-contract
         dict-value-contract
         dict-iter-contract)
