#lang racket/base
(require "intf.rkt" 
         "annotate.rkt"
         "colors.rkt"
         syntax/boundmap
         syntax/kerncase)
(provide annotate-contracts)

(define (annotate-contracts stx low-binders binding-inits)
  (define boundary-start-map (make-hash))
  (define internal-start-map (make-hash))
  (define domain-map (make-hash))
  (define range-map (make-hash))
  
  ;; my-coloring-plans : hash[stx -o-> (listof color)]
  (define my-coloring-plans (make-hash))
  ;; client-coloring-plans : hash[stx -o-> (listof color)]
  (define client-coloring-plans (make-hash))
  
  (let loop ([stx stx])
    (add-to-map stx 'racket/contract:contract-on-boundary boundary-start-map)
    (add-to-map stx 'racket/contract:internal-contract internal-start-map)
    (add-to-map stx 'racket/contract:negative-position domain-map)
    (add-to-map stx 'racket/contract:positive-position range-map)
    (syntax-case stx ()
      [(a . b) (loop #'a) (loop #'b)]
      [_ (void)]))
  
  ;; fill in the coloring-plans table for boundary contracts
  (for ([(start-k start-val) (in-hash boundary-start-map)])
    (for ([start-stx (in-list start-val)])
      (do-contract-traversal start-stx #t
                             my-coloring-plans client-coloring-plans
                             low-binders binding-inits
                             domain-map range-map
                             #t)))
  
  ;; fill in the coloring-plans table for internal contracts
  (for ([(start-k start-val) (in-hash internal-start-map)])
    (for ([start-stx (in-list start-val)])
      (do-contract-traversal start-stx #f
                             my-coloring-plans client-coloring-plans
                             low-binders binding-inits
                             domain-map range-map
                             #f)))
  
  ;; enact the coloring plans
  (for ((coloring-plans (in-list (list my-coloring-plans client-coloring-plans)))
        (mode (in-list '(my-obligations-mode client-obligations-mode))))
    (for ([(stx colors) (in-hash coloring-plans)])
      (cond
        [(and (member my-obligation-style-name colors)
              (member their-obligation-style-name colors))
         (color stx both-obligation-style-name mode)]
        [(member my-obligation-style-name colors)
         (color stx my-obligation-style-name mode)]
        [(member their-obligation-style-name colors)
         (color stx their-obligation-style-name mode)]
        [(member unk-obligation-style-name colors)
         (color stx unk-obligation-style-name mode)]))))

(define (do-contract-traversal start-stx boundary-contract?
                               my-coloring-plans client-coloring-plans
                               low-binders binding-inits domain-map range-map polarity)
  (let ploop ([stx start-stx]
              [polarity polarity])
    
    (define (call-give-up)
      (give-up start-stx boundary-contract? my-coloring-plans client-coloring-plans))
    
    (let ([main-prop (syntax-property stx 'racket/contract:contract)])
      (cond
        [main-prop
         ;; we've found a contract, now go color it and 
         ;; continue with the sub-contract expressions (as indicated via the properties)
         (let sloop ([prop main-prop])
           (cond
             [(pair? prop) (sloop (car prop)) (sloop (cdr prop))]
             [(vector? prop)
              (let ([id (vector-ref prop 0)]
                    [to-color-pos (vector-ref prop 1)]
                    [to-color-neg (vector-ref prop 2)])
                (for ((stx (in-list to-color-pos)))
                  (base-color stx polarity boundary-contract? my-coloring-plans client-coloring-plans))
                (for ((stx (in-list to-color-neg)))
                  (base-color stx (not polarity) boundary-contract? my-coloring-plans client-coloring-plans))
                (for ((stx (in-list (hash-ref domain-map id '()))))
                  (do-contract-traversal stx boundary-contract?
                                         my-coloring-plans client-coloring-plans
                                         low-binders binding-inits domain-map range-map (not polarity)))
                (for ((stx (in-list (hash-ref range-map id '()))))
                  (do-contract-traversal stx boundary-contract?
                                         my-coloring-plans client-coloring-plans
                                         low-binders binding-inits domain-map range-map polarity)))]))]
        
        [else
         ;; we didn't find a contract, but we might find one in a subexpression
         ;; so we need to go look for it (possibly giving up)
         (kernel-syntax-case stx #f
           [(#%expression expr)
            (ploop #'expr polarity)]
           [(module id name-id (#%plain-module-begin mod-level-form ...))
            (call-give-up)]
           [(begin tl-form ... last-one)
            (ploop #'last-one polarity)]
           [(#%provide pvd ...)
            (call-give-up)]
           [(define-values (id ...) expr)
            (call-give-up)]
           [(define-syntaxes (id ...) expr)
            (call-give-up)]
           [(define-values-for-syntax (id ...) expr)
            (call-give-up)]
           [(#%require rspec ...)
            (call-give-up)]
           [id
            (identifier? #'id)
            (if (known-predicate? #'id)
                (base-color #'id polarity boundary-contract? my-coloring-plans client-coloring-plans)
                (let ([binders (module-identifier-mapping-get low-binders #'id (λ () #f))])
                  (if binders
                      (begin
                        (base-color #'id polarity boundary-contract? my-coloring-plans client-coloring-plans)
                        (for ((binder (in-list (module-identifier-mapping-get low-binders #'id))))
                          (base-color binder polarity boundary-contract? my-coloring-plans client-coloring-plans)
                          (for ((rhs (in-list (module-identifier-mapping-get binding-inits binder))))
                            (ploop rhs polarity))))
                      (call-give-up))))]
           [const
            (let ([val (syntax-e #'const)])
              (or (boolean? val)
                  (number? val)
                  (string? val)
                  (char? val)
                  (regexp? val)))
            (base-color stx polarity boundary-contract? my-coloring-plans client-coloring-plans)]
           [(#%plain-lambda (id) expr ...)
            (identifier? #'id)
            (base-color stx polarity boundary-contract? my-coloring-plans client-coloring-plans)]
           [(#%plain-lambda id expr ...)
            (identifier? #'id)
            (base-color stx polarity boundary-contract? my-coloring-plans client-coloring-plans)]
           [(#%plain-lambda formals expr ...)
            (call-give-up)]
           [(case-lambda [formals expr] ...)
            ;; this should really only happen when the arity of the case-lambda includes 1
            ;; (otherwise we should call give-up)
            (base-color stx polarity boundary-contract? my-coloring-plans client-coloring-plans)]
           [(if a b c)
            ;; these calls are questionable. 
            ;; if we ultimately end up giving up in both
            ;; branches, maybe we should actually be coloring the entire thing
            ;; in the blank color, but we'll only color the then and else branches
            ;; in that color with this code.
            ;; on the other hand, recurring like this will mean that the two
            ;; branches are considered separately and thus calling give-up
            ;; on one side will not pollute the other side.
            (do-contract-traversal #'b boundary-contract?
                                   my-coloring-plans client-coloring-plans
                                   low-binders binding-inits domain-map range-map polarity)
            (do-contract-traversal #'c boundary-contract?
                                   my-coloring-plans client-coloring-plans
                                   low-binders binding-inits domain-map range-map polarity)]
           ;; [(begin expr ...) (void)]
           [(begin0 fst rst ...)
            (ploop #'fst polarity)]
           [(let-values bindings body ... last-one)
            (ploop #'last-one polarity)]
           [(letrec-values bindings body ... last-one)
            (ploop #'last-one polarity)]
           [(set! a b)
            (call-give-up)]
           [(quote stuff)
            (base-color stx polarity boundary-contract? my-coloring-plans client-coloring-plans)]
           [(quote-syntax stuff)
            (call-give-up)]
           [(with-continuation-mark a b c)
            (ploop #'c polarity)]
           [(#%plain-app f args ...)
            (call-give-up)]
           [(#%top . id)
            (call-give-up)]
           [(#%variable-reference ignored ...)
            (call-give-up)]
           [_ 
            (begin
              #;(error 'contract-traversal.rkt "unknown thing: ~s\n" stx)
              (call-give-up))
              ])]))))

;; add-to-map : syntax any hash[any -> (listof stx)]
;; looks at stx's property prop and updates map,
;; using the value of the property as the key
(define (add-to-map stx prop map)
  (let loop ([val (syntax-property stx prop)])
    (cond
      [(symbol? val)
       (hash-set! map val (cons stx (hash-ref map val '())))]
      [(pair? val)
       (loop (car val))
       (loop (cdr val))])))

;; returns #t if the result is known to be a predicate that should correspond to a
;; complete obligation for the contract. If it is some unknown variable, this variable
;; may refer to some other contract with nested obligations, so we have to return #f here.
;; approximate this by just asking 'did this identifier come from the core?' (which is known
;; to not bind any contracts (I hope))
(define (known-predicate? id)
  (let ([ib (identifier-binding id)])
    (and (list? ib)
         (let ([src (list-ref ib 0)])
           (let-values ([(base rel) (module-path-index-split src)])
             (member base '('#%kernel racket racket/base scheme scheme/base)))))))

(define (give-up stx boundary-contract? my-coloring-plans client-coloring-plans)
  (let loop ([stx stx])
    (when (syntax-original? stx)
      (blank-color stx boundary-contract? my-coloring-plans client-coloring-plans))
    
    (let oloop ([origin (syntax-property stx 'origin)])
      (cond
        [(pair? origin) (oloop (car origin)) (oloop (cdr origin))]
        [(syntax? origin) 
         (when (syntax-original? origin)
           (blank-color origin boundary-contract? my-coloring-plans client-coloring-plans))]))
    
    (syntax-case stx ()
      [(a . b) (loop #'a) (loop #'b)]
      [_ (void)])))

(define (base-color stx polarity boundary-contract? my-coloring-plans client-coloring-plans)
  (let ([style (if polarity my-obligation-style-name their-obligation-style-name)]
        [neg-style (if polarity their-obligation-style-name my-obligation-style-name)])
    (cond
      [boundary-contract?
       (make-a-coloring-plan stx style my-coloring-plans)
       (make-a-coloring-plan stx neg-style client-coloring-plans)]
      [else
       (make-a-coloring-plan stx style my-coloring-plans)])))

(define (blank-color stx boundary-contract? my-coloring-plans client-coloring-plans)
  (cond
    [boundary-contract?
     (make-a-coloring-plan stx unk-obligation-style-name my-coloring-plans)
     (make-a-coloring-plan stx unk-obligation-style-name client-coloring-plans)]
    [else
     (make-a-coloring-plan stx unk-obligation-style-name my-coloring-plans)]))

(define (make-a-coloring-plan stx plan coloring-plans)
  (hash-set! coloring-plans
             stx
             (cons 
              plan
              (hash-ref coloring-plans stx '()))))
