#lang typed/racket

(require (for-syntax syntax/parse racket/list)
         typed-compose
         "graph-typed.rkt")

(provide Variable VariableMapping GeneralPair
         eval-with eval1-with
         extract-symbols
         any->string stringify-variable-mapping string->any map-sexp
         read-org-sexp unorg unstringify-pairs
         read-org-variable-mapping unorgv read-symbol-list drop-first-last
         list-sets->list-strings
         pretty-print-set pretty-print-set-sets
         ;; Syntax
         auto-hash-ref/explicit auto-hash-ref/:)

(module+ test
  (require typed/rackunit))

;;; ===================
;;; HashTable injection
;;; ===================

(define-type Variable Symbol)
(define-type (VariableMapping A) (Immutable-HashTable Variable A))

(: eval-with (-> (VariableMapping Any) Any AnyValues))
(define (eval-with ht expr)
  (parameterize ([current-namespace (make-base-namespace)])
    (for ([(x val) (in-hash ht)]) (namespace-set-variable-value! x val))
    (eval expr)))

(: eval1-with (-> (VariableMapping Any) Any Any))
(define (eval1-with ht expr)
  (call-with-values (λ () (eval-with ht expr))
                    (λ args (car args))))

(module+ test
  (test-case "eval-with"
    (check-equal? (eval1-with (hash 'a 1 'b 2) '(+ a b 1))
                  4)
    (define ht : (VariableMapping Integer) (hash 'a 1 'b 2))
    (define expr : Any '(+ a b 1))
    (check-equal? (eval1-with ht expr)
                  4)))

(define-syntax (auto-hash-ref/explicit stx)
  (syntax-parse stx
    [(_ (ht:id xs:id ...) body:expr)
     #`(let #,(for/list ([x (syntax->list #'(xs ...))])
                #`[#,x (hash-ref ht '#,x)])
         body)]))

(module+ test
  (test-case "auto-hash-ref/explicit"
    (define mytable #hash((a . 3) (b . 4)))
    (check-equal? (auto-hash-ref/explicit (mytable b a)
                                          (* a b))
                  12)
    (define ht #hash((a . #t) (b . #f)))
    (check-equal? (auto-hash-ref/explicit (ht a b)
                                          (and (not a) b))
                  #f)))

(define-syntax (auto-hash-ref/: stx)
  (syntax-parse stx
    [(_ ht:id body)
     (let* ([names/: (collect-colons (syntax->datum #'body))])
       #`(let #,(for/list ([x names/:])
                  ;; put x in the same context as body
                  #`[#,(datum->syntax #'body x)
                     (hash-ref ht '#,(strip-colon x))])
           body))]))

;;; The helper functions for auto-hash-ref/:.
(begin-for-syntax
  ;; Collect all the symbols starting with a colon in datum.
  (define (collect-colons datum)
    (remove-duplicates
     (flatten
      (for/list ([token datum])
        (cond
          [(symbol? token)
           (let ([name (symbol->string token)])
             (if (eq? #\: (string-ref name 0))
                 token
                 '()))]
          [(list? token)
           (collect-colons token)]
          [else '()])))))

  ;; Strip the leading colon off x.
  (define (strip-colon x)
    (let ([x-str (symbol->string x)])
      (if (eq? #\: (string-ref x-str 0))
          (string->symbol (substring x-str 1))
          x))))

(module+ test
  (test-case "auto-hash-ref/:"
    (define ht1 #hash((x . #t) (y . #t) (t . #f)))
    (define z #t)
    (check-equal? (auto-hash-ref/: ht1
                                   (and :x (not :y) z (or (and :t) :x)))
                  #f)
    (define ht2 #hash((a . 1) (b . 2)))
    (check-equal? (auto-hash-ref/: ht2 (+ :a (* 2 :b)))
                  5)))

;;; ==============================
;;; Analysis of quoted expressions
;;; ==============================

;;; Produces a list of symbols appearing in the quoted expression
;;; passed in the first argument.
(: extract-symbols (-> Any (Listof Symbol)))
(define (extract-symbols form)
  (: extract-rec (-> Any (Listof Any)))
  (define (extract-rec form)
    (match form
      [(? symbol?) (list form)]
      [(? list?)
       (flatten (for/list : (Listof Any)
                    ([x form])
                  (extract-symbols x)))]
      [else '()]))
  (cast (extract-rec form) (Listof Symbol)))

(module+ test
  (test-case "extract-symbols"
    (check-equal? (extract-symbols '(1 (2 3) x (y z 3)))
                  '(x y z))))


;;; =========================
;;; Org-mode interoperability
;;; =========================

(: any->string (-> Any String))
(define (any->string x)
  (with-output-to-string (λ () (display x))))

(module+ test
  (test-case "any->string"
   (check-equal? (any->string 'a) "a")
   (check-equal? (any->string '(a 1 (x y))) "(a 1 (x y))")
   (check-equal? (any->string "hello") "hello")))

(: stringify-variable-mapping (-> (VariableMapping Any) (VariableMapping String)))
(define (stringify-variable-mapping ht)
  (for/hash : (VariableMapping String)
      ([(key val) (in-hash ht)]) (values key (any->string val))))

(module+ test
  (test-case "stringify-variable-mapping"
    (define mp (stringify-variable-mapping #hash((a . (and a b)) (b . (not b)))))
    (check-equal? (hash-ref mp 'a) "(and a b)")
    (check-equal? (hash-ref mp 'b) "(not b)")))

(: string->any (-> String Any))
(define (string->any str)
  (with-input-from-string str (λ () (read))))

(module+ test
  (test-case "string->any"
    (check-equal? (string->any "(or b (not a))") '(or b (not a)))
    (check-equal? (string->any "14") 14)))

;;; Given a sexp, converts all "#f" to #f and "#t" to #t.
;;;
;;; When I read Org-mode tables, I pump them through a call to the
;;; prin1 because the elisp sexp seems incompatible with Racket.  On
;;; the other hand, Racket Booleans seem to upset elisp a little, so
;;; prin1 wraps them in additional double quotes.  This function
;;; removes those quotes.
(: handle-org-booleans (-> Any Any))
(define/match (handle-org-booleans datum)
  [("#t") #t]
  [("#f") #f]
  [((? list?)) (map handle-org-booleans datum)]
  [ (_) datum])

(: map-sexp (-> (-> Any Any) Any Any))
(define (map-sexp func sexp)
  (match sexp
    [(? list?) (map ((curry map-sexp) func) sexp)]
    [datum (func datum)]))

(module+ test
  (test-case "map-sexp"
    (check-equal? (map-sexp (λ (x) (add1 (cast x Number))) '(1 2 (4 10) 3))
                  '(2 3 (5 11) 4))))

(: read-org-sexp (-> String Any))
(define read-org-sexp
  (compose ((curry map-sexp) (match-lambda
                               [(and (? string?) str) (string->any str)]
                               [x x]))
           string->any))
(define unorg read-org-sexp)

(module+ test
  (test-case "read-org-sexp"
    (check-equal? (read-org-sexp "((\"a\" \"(and a b)\") (\"b\" \"(or b (not a))\"))")
                  '((a (and a b)) (b (or b (not a)))))
    (check-equal? (unorg "(#t \"#t\" \"#t \" '(1 2 \"#f\"))")
                  '(#t #t #t '(1 2 #f)))))

(define-type (GeneralPair A B) (U (Pair A B) (List A B)))

(: unstringify-pairs (-> (Listof (GeneralPair String Any))
                         (Listof (GeneralPair Symbol Any))))
(define (unstringify-pairs pairs)
  (for/list ([p pairs])
    (match p
      [(list key val)
       (cons (string->symbol key) (if (string? val)
                                      (string->any val)
                                      val))]
      [(cons key val)
       (cons (string->symbol key) (if (string? val)
                                      (string->any val)
                                      val))])))

(module+ test
  (test-case "unstringify-pairs"
    (check-equal? (unstringify-pairs '(("a" . "1") ("b" . "(and a (not b))")))
                  '((a . 1) (b . (and a (not b)))))
    (check-equal? (unstringify-pairs '(("a" . 1) ("b" . "(and a (not b))")))
                  '((a . 1) (b . (and a (not b)))))))

(: read-org-variable-mapping (-> String (VariableMapping Any)))
(define read-org-variable-mapping
  (multi-compose
   (λ ([pairs : (Listof (Pair Symbol Any))])
     (make-immutable-hash pairs))
   (λ (sexp)
     (unstringify-pairs (cast sexp (Listof (GeneralPair String Any)))))
   string->any))

;;; A synonym for read-org-variable-mapping.
(define unorgv read-org-variable-mapping)

(module+ test
  (test-case "read-org-variable-mapping"
    (define m1 (read-org-variable-mapping "((\"a\" \"(and a b)\") (\"b\" \"(or b (not a))\"))"))
    (define m2 (read-org-variable-mapping "((\"a\" . \"(and a b)\") (\"b\" . \"(or b (not a))\"))"))
    (define m3 (unorgv "((\"a\" . \"(and a b)\") (\"b\" . \"(or b (not a))\"))"))
    (check-equal? (hash-ref m1 'a) '(and a b))
    (check-equal? (hash-ref m2 'a) '(and a b))
    (check-equal? (hash-ref m3 'a) '(and a b))
    (check-equal? (hash-ref m1 'b) '(or b (not a)))
    (check-equal? (hash-ref m2 'b) '(or b (not a)))
    (check-equal? (hash-ref m3 'b) '(or b (not a)))))

(: dotit (-> Graph Void))
(define dotit (compose display graphviz))

(: read-symbol-list (-> String (Listof Symbol)))
(define (read-symbol-list str)
  (cast (string->any (string-append "(" str ")")) (Listof Symbol)))

(module+ test
  (test-case "read-symbol-list"
    (check-equal? (read-symbol-list "a b c") '(a b c))))

(: drop-first-last (-> String String))
(define (drop-first-last str)
  (substring str 1 (- (string-length str) 1)))

(module+ test
  (test-case "drop-first-last"
    (check-equal? (drop-first-last "(a b)") "a b")))

(: list-sets->list-strings (-> (Listof (Setof Any)) (Listof String)))
(define (list-sets->list-strings lst)
  (map (multi-compose drop-first-last
                      any->string
                      (λ ([x : (Setof Any)])
                        (set->list x))) lst))

(module+ test
  (test-case "list-sets->list-strings"
    (check-equal? (list-sets->list-strings (list (set 'x 'y) (set 'z) (set) (set 't)))
                  '("y x" "z" "" "t"))))

(: pretty-print-set (-> (Setof Any) String))
(define (pretty-print-set s)
  (string-join (sort (set-map s any->string) string<?)))

(module+ test
  (test-case "pretty-print-set"
    (check-equal? (pretty-print-set (set 'a 'b 1)) "1 a b")))

(: pretty-print-set-sets (-> (Setof (Setof Any)) String))
(define (pretty-print-set-sets ms)
  (string-join (for/list ([m ms]) : (Listof String)
                         (format "{~a}" (pretty-print-set m))) ""))

(module+ test
  (test-case "pretty-print-set-sets"
    (check-equal? (pretty-print-set-sets (set (set 'a 'b) (set 'c))) "{a b}{c}")))