#lang racket

(define globals (make-hash))

(define-struct amb-function [argument body env])

(define (extend-environment name value env)
  (cons (cons name value) env))

(define (lookup-variable name env)
  (cond
    [(assq name env) => cdr]
    [(hash-ref globals name)]))

(define (amb-function* arguments body env)
  (when (null? arguments)
    (error "amb-function requires at least one argument"))
  (if (null? (cdr arguments))
      (amb-function (car arguments) body env)
      (amb-function (car arguments)
                    (amb-function* (cdr arguments) body env)
                    env)))

(define (stream-mappend function stream)
  (stream-fold stream-append empty-stream
               (stream-map function stream)))

(define (amb-apply1 function argument)
  (stream-mappend
   (lambda (x)
     (cond
       [(amb-function? x)
        (evaluate (amb-function-body x)
                  (extend-environment (amb-function-argument x)
                                      (stream argument)
                                      (amb-function-env x)))]
       [(procedure? x)
        (stream (x argument))]
       [else (error "not a function" x)]))
   function))

(define (amb-apply function arguments)
  (if (null? arguments)
      function
      (let ([app1 (stream-mappend (lambda (x)
                                    (amb-apply1 function x))
                                  (car arguments))])
        (amb-apply app1 (cdr arguments)))))

(define (amb values env)
  (if (null? values)
      empty-stream
      (stream-append (evaluate (car values) env)
                     (amb (cdr values) env))))

(define (evaluate value env)
  (match value
    [`(lambda ,arguments ,body)
     (evaluate
      (let loop ([arguments arguments])
        (cond
         [(null? arguments)
          (error "lambda takes at least one argument")]
         [(null? (cdr arguments))
          `(lambda1 ,(car arguments) ,body)]
         [else
          `(lambda1 ,(car arguments)
                    ,(loop (cdr arguments)))]))
      env)]
    [`(lambda1 ,argument ,body)
     (stream (amb-function argument body env))]
    [`(let ,binds ,body)
     (let ([new-value `((lambda ,(map first binds) ,body) . ,(map second binds))])
       (evaluate new-value env))]
    [`(begin . ,exprs)
     (for/last ([expr exprs])
       (evaluate expr env))]
    [`(quote ,value) (stream value)]
    [`(amb . ,values)
     (amb values env)]
    [`(define (,name . ,args) ,body)
     (evaluate
      `(define ,name (lambda ,args ,body))
      env)]
    [`(define ,name ,value)
     (hash-set! globals name
                (evaluate value env))
     (stream name)]
    [`(if ,test ,then ,else)
     ;; Try to avoid repeating evaluation.
     (let ([then (delay (evaluate then env))]
           [else (delay (evaluate else env))])
       (stream-mappend (lambda (x)
                         (if x
                             (force then)
                             (force else)))
                       (evaluate test env)))]
    [`(,function . ,args)
     (let ([function (evaluate function env)]
           [args (map (lambda (x) (evaluate x env)) args)])
       (amb-apply function args))]
    [_
     (if (symbol? value)
         (lookup-variable value env)
         (stream value))]))

(define (curry2 proc)
  (lambda (x)
    (lambda (y)
      (proc x y))))

(define default-environment
  `((cons . ,(stream (curry2 cons)))
    (car  . ,(stream car))
    (cdr  . ,(stream cdr))
    (add1 . ,(stream add1))
    (print . ,(stream print))
    (>    . ,(stream (curry2 >)))
    (=    . ,(stream (curry2 =)))
    (*    . ,(stream (curry2 *)))
    (+    . ,(stream (curry2 +)))
    (~    . ,(stream -))
    (-    . ,(stream (curry2 -)))
    (+-   . ,(stream (amb-function 'x
                                   '(amb x (~ x))
                                   '())))
    (not  . ,(stream not))
    (range . ,(stream (amb-function 'low
                                    '(lambda (high) (if (> low high)
                                                        (amb)
                                                        (amb low
                                                             (range (add1 low)
                                                                    high))))
                                    '())))))
(for ([pair default-environment])
  (match pair
   [(cons name value)
    (hash-set! globals name value)]))

(define (repl (environment '()))
  (let loop ()
    (let ([input (read)])
      (with-handlers ([exn:fail? (lambda (e) (displayln e))])
        (time
         (for ([value (evaluate input environment)])
           (pretty-print value)))))
    (loop)))

(repl)