#lang plait
(require (typed-in racket/base [gensym : (-> Symbol)]))
(define-type-alias Label Symbol)

(define-type LExpr
  (e-num [label : Label] [x : Number])
  (e-bool [label : Label] [x : Boolean])
  (e-str [label : Label] [x : String])
  (e-op [label : Label]
        [op : Operator]
        [left : LExpr]
        [right : LExpr])
  (e-un-op [label : Label]
           [op : UnaryOperator]
           [LExpr : LExpr])
  (e-if [label : Label]
        [cond : LExpr]
        [consq : LExpr]
        [altern : LExpr])
  (e-lam [label : Label]
         [param-label : Label]
         [param : Symbol]
         [body : LExpr])
  (e-app [label : Label]
         [func : LExpr]
         [arg : LExpr])
  (e-id [label : Label] [name : Symbol])
  (sugar-and [label : Label]
             [left : LExpr]
             [right : LExpr])
  (sugar-or [label : Label]
            [left : LExpr]
            [right : LExpr])
  (sugar-let [label : Label]
             [id : Symbol]
             [value : LExpr]
             [body : LExpr])
  (e-empty [label : Label]))

(define-type Operator
  (op-plus)
  (op-num-eq)
  (op-append)
  (op-str-eq)
  (op-link))

(define-type UnaryOperator
  (op-first)
  (op-rest)
  (op-is-empty))

(define-type-alias LEnv (Hashof Symbol Label))

(define-type TypeInferenceError
  (ti-err-fails-occurs-check [replace : Term] [with : Term])
  (ti-err-constructor-mismatch [t1 : Term] [t2 : Term])
  (ti-err-unbound-id [name : Symbol]))

(define-syntax-rule (~a arg ...)
  (foldl (lambda (val string) (string-append string val)) "" (list (to-string arg) ...)))
  
(define (raise-error [err : TypeInferenceError])
  (type-case TypeInferenceError err
    [(ti-err-fails-occurs-check replace with)
     (error 'ti-err-fails-occurs-check (~a replace " occurs in " with))]
    [(ti-err-constructor-mismatch t1 t2) (error 'ti-err-constructor-mismatch (~a t1 " does not unify with " t2))]
    [(ti-err-unbound-id name) (error 'ti-err-unbound-id (~a name " is unbound"))]))
  
(define (parse [input : S-Exp]): LExpr
  (cond
    [(s-exp-match? `true input) (e-bool (gensym) #t)]
    [(s-exp-match? `false input) (e-bool (gensym) #f)]
    [(s-exp-number? input) (e-num (gensym) (s-exp->number input))]
    [(s-exp-string? input) (e-str (gensym) (s-exp->string input))]
    [(s-exp-match? `{if ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 4)
           (e-if (gensym) (parse (second inlst)) (parse (third inlst)) (parse (fourth inlst)))
           (error '+ "incorrect number of args to +")))]
    [(s-exp-match? `{+ ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-plus) (parse (second inlst)) (parse (third inlst)))
           (error '+ "incorrect number of args to +")))]
    [(s-exp-match? `{num= ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-num-eq) (parse (second inlst)) (parse (third inlst)))
           (error 'num= "incorrect number of args to num=")))]
    [(s-exp-match? `{++ ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-append) (parse (second inlst)) (parse (third inlst)))
           (error '++ "incorrect number of args to ++")))]
    [(s-exp-match? `{str= ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-str-eq) (parse (second inlst)) (parse (third inlst)))
           (error 'str= "incorrect number of args to str=")))]
    [(s-exp-match? `{lam SYMBOL ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-lam (gensym) (gensym) (s-exp->symbol (second inlst)) (parse (third inlst)))
           (error 'lam "lambdas should only have one body")))]
    [(s-exp-match? `{lam ANY ...} input)
     (error 'lam "lambda parameters must be symbols")]
    [(s-exp-match? `{and ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (sugar-and (gensym) (parse (second inlst)) (parse (third inlst)))
           (error 'and "incorrect number of args to and")))]
    [(s-exp-match? `{or ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (sugar-or (gensym) (parse (second inlst)) (parse (third inlst)))
           (error 'or "incorrect number of args to or")))]
    [(s-exp-match? `{let {SYMBOL ANY} ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (sugar-let
            (gensym)
            (s-exp->symbol (first (s-exp->list (second inlst))))
            (parse (second (s-exp->list (second inlst))))
            (parse (third inlst)))
           (error 'let "incorrect number of args to let")))]
    [(s-exp-match? `empty input) (e-empty (gensym))]
    [(s-exp-match? `{link ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-link) (parse (second inlst)) (parse (third inlst)))
           (error 'link "incorrect number of args to link")))]
    [(s-exp-match? `{is-empty ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 2)
           (e-un-op (gensym) (op-is-empty) (parse (second inlst)))
           (error 'is-empty "incorrect number of args to is-empty")))]
    [(s-exp-match? `{first ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 2)
           (e-un-op (gensym) (op-first) (parse (second inlst)))
           (error 'first "incorrect number of args to first")))]
    [(s-exp-match? `{rest ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 2)
           (e-un-op (gensym) (op-rest) (parse (second inlst)))
           (error 'rest "incorrect number of args to rest")))]
    [(s-exp-match? `{ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 2)
           (e-app (gensym) (parse (first inlst)) (parse (second inlst)))
           (error 'app "incorrect number of args to app")))]
    [(s-exp-symbol? input) (e-id (gensym) (s-exp->symbol input))]))

(define-type Term
  (t-var [label : Label])
  (t-con [head : Symbol] [args : (Listof Term)]))
(define-type Constraint
  (eq [l : Term] [r : Term]))
(define-type-alias Substitution (Hashof Label Term))
(define-type-alias LEnv (Hashof Symbol Label))

(define (t-bool) (t-con 'Bool empty))
(define (t-num) (t-con 'Num empty))
(define (t-str) (t-con 'Str empty))
(define (t-list [elem : Term]) (t-con 'List (list elem)))
(define (t-fun [arg : Term] [ret : Term]) (t-con '-> (list arg ret)))
(define (get-label [e : LExpr]) : Label
  (type-case LExpr e
    [(e-bool lbl v) lbl]
    [(e-num lbl v) lbl]
    [(e-str lbl v) lbl]
    [(e-op lbl op l r) lbl]
    [(e-un-op lbl op e) lbl]
    [(e-if lbl cond consq altern) lbl]
    [(e-lam lbl param-lbl param body) lbl]
    [(e-app lbl func arg) lbl]
    [(e-id lbl name) lbl]
    [(sugar-and lbl l r) lbl]
    [(sugar-or lbl l r) lbl]
    [(sugar-let lbl id value body) lbl]
    [(e-empty lbl) lbl]))
(define (get-t-var [e : LExpr]) (t-var (get-label e)))
; Normalize consumes a Term t and produces a Term
; which is structurally identical to t but renames the t-var labels
; such that:
; 1) the first t-var label is renamed throughout with 'a
; 2) if β is renamed to γ
;       and (t-var 'δ) is the next unseen t-var after (t-var 'β)
;       and ϵ directly follows γ in the English alphabet,
;    then δ will be renamed to ϵ
(define (normalize [t : Term]) : Term
  (letrec ([char->symbol (λ (c) (string->symbol (list->string (cons c empty))))] ; y u no exist
           [αs (box (map char->symbol (string->list "abcdefghijklmnopqrstuvwxyz")))]
           [pop
            (λ ()
              (let ([α (first (unbox αs))])
                (begin (set-box! αs (rest (unbox αs))) α)))]
           [θ (make-hash empty)]
           [_n
            (λ (t)
              (type-case Term t
                [(t-var lbl)
                 (type-case (Optionof Symbol) (hash-ref θ lbl)
                   [(none)
                    (let ([α (pop)])
                      (begin (hash-set! θ lbl α) (t-var α)))]
                   [(some α) (t-var α)])]
                [(t-con head args) (t-con head (map _n args))]))])
    (_n t)))

(define (ormap [f : ('A -> Boolean)] [l : (Listof 'A)]) : Boolean
  (foldl (λ (c a) (or a (f c))) #f l))