#lang plai

;;; TYPED MINI-SCHEME

(print-only-errors #t)

;;;;;;;;;;;;;;;;;;; AST ;;;;;;;;;;;;;;;;;;;
(define-type Expr
  [num (n number?)]
  [str (s string?)]
  [bool (b boolean?)]
  [id (name symbol?)]
  [ifc (test Expr?) (then Expr?) (else Expr?)]
  [fun (params (listof typed-param?)) (body Expr?)]
  [app (fun-expr Expr?) (arg-exprs (listof Expr?))]
  [lcal (defs (listof Def?)) (body Expr?)])

(define (typed-param? v)
  (and (pair? v)
       (symbol? (car v))
       (Type? (cdr v))))

(define-type Def
  [def (name symbol?) (val-expr Expr?)])

(define-type Type
  [TNum]
  [TStr]
  [TBool]
  [TFun (args (listof Type?)) (ret Type?)])
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


;;;;;;;;;;;;;;;;;;; PARSER ;;;;;;;;;;;;;;;;;;;
(define (parse sexp)
  (cond
    [(number? sexp) (num sexp)]
    [(symbol? sexp) (id sexp)]
    [(string? sexp) (str sexp)]
    [(boolean? sexp) (bool sexp)]
    [(list? sexp)
     (case (first sexp)
       [(fun) (fun (map (λ (p) (cons (first p)
                                     (parse-type (third p))))
                        (second sexp))
                   (parse (third sexp)))]
       [(if) (ifc (parse (second sexp))
                   (parse (third sexp))
                   (parse (fourth sexp)))]
       [(local) (lcal (map parse-def (second sexp))
                      (parse (third sexp)))]
       [else (app (parse (first sexp))
                  (map parse (rest sexp)))])]))

(define (parse-def sexp)
  (case (first sexp) 
    [(define) (def (second sexp) (parse (third sexp)))]))

(define (parse-type s)
  (if (list? s) 
      (TFun (map parse-type (take s (- (length s) 2)))
            (parse-type (last s)))
      (case s
        [(Num) (TNum)]
        [(Bool) (TBool)]
        [(Str) (TStr)])))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;;;;;;;;;;;;;;;;;; VALUES ;;;;;;;;;;;;;;;;;;;
(define-type Val
  [numV (n number?)]
  [strV (s string?)]
  [boolV (b boolean?)]
  [procV (p procedure?)]
  [closureV (params (listof symbol?)) (body Expr?) (env Env?)])

(define (unwrap v)
  (type-case Val v
    [numV (n) n]
    [strV (s) s]
    [boolV (b) b]
    [procV (p) p]
    [closureV (p b e) (error 'unwrap "cannot unwrap closureV")]))

(define (wrap v)
  (cond
    [(number? v) (numV v)]
    [(string? v) (strV v)]
    [(boolean? v) (boolV v)]
    [(procedure? v) (procV v)]
    [else (error 'wrap "cannot wrap: ~a" v)]))


;;;;;;;;;;;;;;;;;;; ENVIRONMENT ;;;;;;;;;;;;;;;;;;;
(define-type Env
  [mtEnv]
  [aEnv (bindings (listof pair?)) (rest Env?)])

;; env-lookup :: Sym Env -> Val
(define (env-lookup id env)
  (type-case Env env
    [mtEnv () (error 'env-lookup "no binding for identifier: ~a " id)]
    [aEnv (bds r)
          (let ([res (assoc id bds)])
            (if res
                (cdr res)
                (env-lookup id r)))]))

;; update-env! :: Env Sym Val -> Void
;; imperative update of env, adding/overring the binding for id.
(define (update-env! env id val)
  (set-aEnv-bindings! env (cons (cons id val) (aEnv-bindings env))))

(define (initEnv)
  (aEnv
   (list 
    (cons '+ (procV +))
    (cons '- (procV -))
    (cons '* (procV *))
    (cons '/ (procV /))
    (cons 'equal? (procV equal?))
    (cons 'zero?  (procV zero?))
    (cons '< (procV < ))
    (cons '> (procV >))
    (cons 'not (procV not))
    (cons 'and (procV (λ (x y) (and x y))))
    (cons 'or  (procV (λ (x y) (or x))))
    (cons 'string<? (procV string<?))
    (cons 'string-length (procV string-length))
    (cons 'string-append (procV string-append))
    (cons 'number->string (procV number->string))
    ) 
   (mtEnv)))


(define (initTEnv)
  (aEnv
   (list 
    (cons '+ (TFun (list (TNum) (TNum)) (TNum)))
    (cons '- (TFun (list (TNum) (TNum)) (TNum)))
    (cons '* (TFun (list (TNum) (TNum)) (TNum)))
    (cons '/ (TFun (list (TNum) (TNum)) (TNum)))
    (cons 'equal? (TFun (list (TNum) (TNum)) (TBool)))
    (cons 'zero?  (TFun (list (TNum)) (TBool)))
    (cons '<   (TFun (list (TNum) (TNum)) (TBool)))
    (cons '>   (TFun (list (TNum) (TNum)) (TBool)))
    (cons 'not (TFun (list (TBool)) (TBool)))
    (cons 'and (TFun (list (TBool) (TBool)) (TBool)))
    (cons 'or  (TFun (list (TBool) (TBool)) (TBool)))
    (cons 'string<? (TFun (list (TStr) (TStr)) (TBool)))
    (cons 'string-length (TFun (list (TStr)) (TNum)))
    (cons 'string-append (TFun (list (TStr) (TStr)) (TStr)))
    (cons 'number->string (TFun (list (TNum)) (TStr)))
    )
   (mtEnv)))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;;;;;;;;;;;;;;;;;; TYPING ;;;;;;;;;;;;;;;;;;;
(define (typeof expr tenv)
  (type-case Expr expr
    [num (n) (TNum)]
    [str (s) (TStr)]
    [bool (b) (TBool)]
    [fun (tparams body) 
         (TFun (map cdr tparams)
               (typeof body (aEnv tparams tenv)))]
    [ifc (c t f)
         (let ([tc (typeof c tenv)]
               [tt (typeof t tenv)]
               [tf (typeof f tenv)])
           (check "if-condition" tc (TBool))
           (check "if-branch" tt tf)
           tt)]           
    [id (v) (env-lookup v tenv)]
    [app (fun-expr arg-exprs)
         (type-case Type (typeof fun-expr tenv)
           [TFun (tparams tret) 
                 (let ([targs (map (λ (e) (typeof e tenv)) arg-exprs)])
                   (for-each check
                             (make-list (length tparams) "argument type")
                             tparams targs)
                   tret)]
           [else (error 'typeof "expected function, got ~v" (typeof fun-expr tenv))])]
    [lcal (defs body)
      (let ([new-tenv (aEnv '() tenv)])
        (begin 
          (for-each (λ (d) (typeof-def d new-tenv)) defs)
          (typeof body new-tenv)))]))
    
(define (typeof-def d tenv)
  (type-case Def d
    [def (id val-expr)
      (update-env! tenv id (typeof val-expr tenv))]))

(define (check ctx t1 t2)
  (when (not (equal? t1 t2))
    (error 'typeof "type error in ~a: got ~v expected ~v" ctx t1 t2)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    

;;;;;;;;;;;;;;;;;;; INTERPRETER ;;;;;;;;;;;;;;;;;;;
(define (interp expr env)
  (type-case Expr expr
    [num (n) (numV n)]
    [str (s) (strV s)]
    [bool (b) (boolV b)]
    [ifc (c t f)
         (if (boolV-b (interp c env))
             (interp t env)
             (interp f env))]
    [id (v) (env-lookup v env)]
    [fun (tparams body) 
         (closureV (map car tparams) body env)]
    [app (fun-expr arg-exprs)
         (let* ([fun-val (interp fun-expr env)]
                [arg-vals (map (λ (e) (interp e env)) arg-exprs)])
           (apply-fun fun-val arg-vals))]
    [lcal (defs body)
          (let ([new-env (aEnv '() env)])
            (begin 
              (for-each (λ (d) (interp-def d new-env)) defs)
              (interp body new-env)))]))

(define (apply-fun f vs)
  (type-case Val f
    [closureV (ps body env)
              (interp body (aEnv (map cons ps vs) env))]
    [procV (p) (wrap (apply p (map unwrap vs)))]
    [else (error 'interp "not a function")]))

(define (interp-def d env)
  (type-case Def d
    [def (id val-expr)
      (update-env! env id (interp val-expr env))]))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; running a program (after typechecking it)
(define (run p)
  (let ([prog (parse p)])
    (typeof prog (initTEnv))
    (interp prog (initEnv))))



;;;;;;;
;; some tests
(test (run '5) (numV 5))
(test (run '(+ 5 5)) (numV 10))
(test (run '(+ 1 ((fun ([x : Num] [y : Num]) (+ x y)) 1 2))) (numV 4))
(test (run '(((fun ([x : (Num -> Num)]) x) 
                    (fun ([x : Num]) (+ x 5))) 3)) (numV 8))
(test (run '(local ((define x 10)
                    (define y x))
              (+ x y)))
      (numV 20))

(test (run '(local 
                ((define one? (fun ([x : Num]) 
                                (zero? (- x 1)))))
              (if (one? 1)
                  "hola"
                  "chao")))
      (strV "hola"))

(test/exn (run '(+ 5 "hola")) "type")
(test/exn (run '(string-length 4)) "type")
(test/exn (run '(if 5 "hola" "hola")) "type")
(test/exn (run '(if (zero? 0) "hola" 1)) "type")
(test/exn (run '(+ 1 ((fun ([x : Num] [y : Num]) (+ x y)) 1 "hola"))) "type")
(test/exn (run '(+ 1 ((fun ([x : Num] [y : Num]) "hola") 1 2))) "type")