#lang racket

(provide pc2ts compile/run)

(define reg-funcs '())
(define reg-pc 'no_pc)
(define dismount-var 'no_dis)
(define construct-var 'no_const)

(define global-decls '())
(define current-global-decls '())
(define main-def "")
(define union-defs "")
(define js-keywords
  ; Need to update
  '(False class from or None continue global pass True
          def if raise and del import return as elif
          in try assert else is while async except
          lambda with await finally nonlocal yield break
          for not))

(define add-func
  (lambda (func)
    (set! reg-funcs (cons func reg-funcs))))

(define is-func?
  (lambda (func)
    (assv func reg-funcs)))

(define is-global?
  (λ (var)
    (cond
      [(equal? (global var) reg-pc) #t]
      [(memv var reg-regs)
       => (λ (p) #t)]
      [else #f])))

(define reg-unions '())

(define check-args
  (lambda (union args)
    (cond
      [(null? args) #t]
      [(memq (car args) (cdr args))
       (error 'define-union "duplicated variant `~a' in union `~a'\n" (car args) (car union))]
      [else (check-args union (cdr args))])))

(define add-union
  (lambda (union)
    (if (not (lookup-union (car union)))
        (begin (check-args union (cadr union)) 
               (set! reg-unions (cons union reg-unions)))
        (error 'define-union "duplicated definition of union-type `~a'\n" (car union)))))

(define reg-regs '())

(define init-storage
  (lambda ()
    (set! reg-funcs '())
    (set! reg-unions '())
    (set! reg-regs '())))

(define new-safe-char
  (lambda (char)
    (cond
      [(eq? #\? char) "_"]
      [(eq? #\! char) "_"]
      [(eq? #\. char) "_"]
      [(eq? #\+ char) "_"]
      [(eq? #\- char) "_"]
      [(eq? #\* char) "_"]
      [(eq? #\/ char) "_"]
      [(eq? #\< char) "_"]
      [(eq? #\> char) "_"]
      [(eq? #\: char) "_"]
      [(eq? #\$ char) "_"]
      [(eq? #\% char) "_"]
      [(eq? #\^ char) "_cap"]
      [(eq? #\& char) "_"]
      [(eq? #\~ char) "_"]
      [(eq? #\_ char) "_"]
      [(and (char>=? char #\0) (char<=? char #\9))
       (string-append "r" (list->string `(,char)))]
      [else (list->string `(,char))])))

(define safe 
  (lambda (sym)
    (let ([str-sym (raw-safe sym)])
      (cond
        [(memv sym js-keywords)
         => (λ (p) (string-append "_" str-sym))]
        [else str-sym]))))

(define raw-safe
  (lambda (sym)
    (if (symbol? sym)
        (let loop ([l (string->list (symbol->string sym))])
          (cond
            [(null? l) ""]
            [else (string-append
                   (new-safe-char (car l))
                   (loop (cdr l)))]))
        sym)))

; TODO: Remove this
(define global
  (lambda (sym)
    (string-append "" (safe sym))))

(define join
  (lambda (lst separator)
    (let loop ([lst lst]
               [result ""]
               [is-first? #t])
      (cond
        [(null? lst) result]
        [is-first? (loop (cdr lst)
                         (format "~a" (car lst)) #f)]
        [else (loop (cdr lst)
                    (string-append
                     result
                     (format "~a~a" separator (car lst)))
                    #f)]))))

(define file->list
  (lambda (fname)
    (let ([file (open-input-file fname)])
      (let ([data
             (let recurse ([decl (read file)])
               (if (eof-object? decl)
                   '()
                   (cons decl (recurse (read file)))))])
        (close-input-port file)
        data))))

(define pc2ts
  (lambda (file-name source-name)
    ;; WARNING: pc2ts will erase existing files when generating new ones!
    (when (file-exists? source-name) (delete-file source-name))
    (init-storage)
    (let ([decl* (file->list file-name)])
      (let ([src (open-output-file source-name)])
        (dynamic-wind
         (lambda () #f)
         (lambda ()
           ;; write a generated header file to header-name
           (display (pc2ts-header decl*) src)
           (check-correct-info)
           ;; write a generated source file source-name
           (display (pc2ts-source) src)
           (set! global-decls '())
           (set! union-defs ""))
         (lambda ()
           (close-output-port src)))))))

(define check-correct-info
  (lambda ()
    (begin
      (if (null? reg-regs)
          (display "Warning: you have defined no registers.\n")
          (void)))))

(define pc2ts-append
  (lambda args
    (apply string-append
           (map (lambda (elt)
                  (cond
                    [(symbol? elt) (format "~a" elt)]
                    [(number? elt) (format "~s" elt)]
                    [(string? elt) elt]
                    [else (error 'pc2ts-append "Invalid argument ~s" elt)]))
                args))))

(define pc2ts-gen-unions
  (lambda (union)
    (let ([name (safe (car union))]
          [tag* (cadr union)]
          [field** (caddr union)])
      (apply string-append
             (map (lambda (tag field*)
                    (let* ([safe-tag (safe tag)]
                           [fnname (cond
                                     [(memv tag js-keywords)
                                      => (λ (p) (pc2ts-append name safe-tag))]
                                     [else (pc2ts-append name "_" safe-tag)])])
                      (pc2ts-append
                       (pc2ts-fn-proto fnname field*) " {\n"
                       "    return UnionType(UnionEnums." name "." safe-tag ", {"
                       (string-join (map (λ (v)
                                           (string-append
                                            "\n        \""
                                            (safe v) "\": " (safe v)))
                                         field*)
                                    ", ")
                       "\n    });\n}\n")))
                  tag* field**)))))

;; added by wy for constructor argument name binding
;; lookup-arg looks up the argument name of name.tag at position pos
(define lookup-union
  (lambda (name)
    (let loop ([reg reg-unions])
      (cond
        [(null? reg) #f]
        [(eq? name (caar reg)) (car reg)]
        [else (loop (cdr reg))]))))

(define get-arg-list
  (lambda (name tag)
    (let ([u (lookup-union name)])
      (if (not u) (error 'lookup-union 
                         "union type `~a' not defined\n" name)
          (let loop ([tags (cadr u)] [args (caddr u)])
            (cond
              [(null? tags) 
               (error 'lookup-arg
                      "union type `~a' doesn't have a tag `~a'~n" name tag)]
              [(eq? tag (car tags)) (car args)]
              [else (loop (cdr tags) (cdr args))]))))))

(define lookup-arg
  (lambda (name tag pos)
    (list-ref (get-arg-list name tag) pos)))

(define check-union-case
  (lambda (expr name type case)
    (cond
      [(and (null? type) (not (null? case)))
       (let ([s (open-output-string)])
         (pretty-print expr s)
         (error 'union-case  "~a\nsuperfluous cases for union type `~a': ~a"
                (get-output-string s) name case))]
      [(and (null? case) (not (null? type)))
       (let ([s (open-output-string)])
         (pretty-print expr s)
         (error 'union-case  "~a\nunmatched cases for union type `~a': ~a"
                (get-output-string s) name type))]
      [(and (null? type) (null? case)) #t]
      [(not (memq (car case) type))
       (let ([s (open-output-string)])
         (pretty-print expr s)
         (error 'union-case "~a\nvariant `~a' is not in union type `~a'"
                (get-output-string s) (car case) name))]
      [(memq (car case) (cdr case))
       (let ([s (open-output-string)])
         (pretty-print expr s)
         (error 'union-case  "~a\nduplicated cases `~a' in union-case of type `~a'"
                (get-output-string s) (car case) name))]
      [else (check-union-case expr name (remq (car case) type) (cdr case))])))

(define case-env
  (lambda (env var*)
    (let loop ([env env] [var* var*])
      (if (null? var*)
          env
          (extend-env (car var*) (car var*) (loop env (cdr var*)))))))

(define handle-union-case-case
  (lambda (name env u_obj)
    (lambda (template body)
      (match template 
        [`(,tag . ,var*) #:when (list? var*)
                         (let ([sname (safe name)]
                               [stag (safe tag)])
                           (let ([given (length var*)] [expected (length (get-arg-list name tag))])
                             (if (not (= given expected))
                                 (error 'union-case
                                        "~a\nwrong number of arguments to constructor `~a' of union-type `~a': expected: ~a, given: ~a"
                                        template tag name expected given)
                                 (pc2ts-append
                                  "        case UnionEnums." sname "." stag ":\n"
                                  (let loop ([var* var*] [n 0])
                                    (cond
                                      [(null? var*) ""]
                                      [else (string-append 
                                             (pc2ts-append
                                              "            var " (safe (car var*)) " = " u_obj
                                              "." (safe (lookup-arg name tag n)) ";\n")
                                             (loop (cdr var*) (add1 n)))]))
                                  ((parse-function-body #t (case-env env var*) 3) body)
                                  "            break;\n\n"))))]
        ;; Cannot possibly be effective, commented JBH 12/13
        ;; [else (string-append "default {\n"
        ;;         ((parse-function-body #t (case-env env var*)) body)
        ;;         "}\n")]
        ))))

(define get-last
  (lambda (ls)
    (cond
      ((null? ls) #f)
      ((null? (cdr ls)) (car ls))
      (else (get-last (cdr ls))))))

;; this is for error checking
(define get-body
  (lambda (c)
    (match c
      [`(,test ,body) body])))

(define remove-last
  (lambda (ls)
    (match ls
      [`((else ,body)) '()]
      [`((,test ,body) . ,c*) `((,test ,body) . ,(remove-last c*))])))

(define apply-env
  (lambda (env x)
    (match env
      [`(empty-env) (error 'empty-env "unbound variable: ~s" x)]
      [`(extend-env ,x^ ,a ,env)
       (if (eq? x^ x) a (apply-env env x))])))

(define extend-env
  (lambda (x a env)
    `(extend-env ,x ,a ,env)))

(define empty-env
  (lambda ()
    `(empty-env)))

(define tabs
  (lambda (n)
    (cond
      [(zero? n) ""]
      [else (string-append "    " (tabs (sub1 n)))])))

(define parse-function-body
  (lambda (tail env level)
    (if tail
        (lambda (expr)
          (match expr
            [`(error ,name ,msg) 
             (pc2ts-append
              (tabs level) "throw Error(\"" msg "\");\n")]
            [`(if ,test ,conseq ,alt) 
             (let ((test ((parse-function-body #f env (add1 level)) test))
                   (conseq ((parse-function-body #t env (add1 level)) conseq))
                   (alt ((parse-function-body #t env (add1 level)) alt)))
               (pc2ts-append
                (tabs level) "if (" test ") {\n"
                conseq
                (tabs level) "} else {\n"
                alt
                (tabs level) "}\n"))]
            [`(cond (else ,body)) 
             (let ((body ((parse-function-body #t env level) body)))
               body)]
            [`(cond . ,c*) 
             (let ((last (get-last c*))
                   (c* (remove-last c*)))
               (cond
                 [(eq? (car last) 'else)
                  (let* ((test0 ((parse-function-body #f env level) (caar c*)))
                         (body0 ((parse-function-body #t env (add1 level)) (get-body (car c*))))
                         (test* (map (parse-function-body #f env level) (map car (cdr c*))))
                         (body* (map (parse-function-body #t env (add1 level)) (map get-body (cdr c*))))
                         (body ((parse-function-body #t env (add1 level)) (cadr last))))
                    (pc2ts-append
                     (tabs level) "if (" test0 ") {\n" body0 "\n"
                     (apply string-append
                            (map (lambda (x y)
                                   (pc2ts-append (tabs level) "} else if (" x ") {\n"
                                                 y))
                                 test* body*))
                     (tabs level) "} else {\n" body "    }\n"))]
                 [else
                  (let* ((test0 ((parse-function-body #f env level) (caar c*)))
                         (body0 ((parse-function-body #t env level) (cadar c*)))
                         (test* (map (parse-function-body #f env level) (map car (cdr c*))))
                         (body* (map (parse-function-body #t env level) (map cadr (cdr c*)))))
                    (pc2ts-append
                     "if (" test0 ") {\n"
                     "    " body0 "\n"
                     (apply string-append
                            (map (lambda (x y)
                                   (pc2ts-append "} else if (" x ") {\n"
                                                 y))
                                 test* body*))
                     "}\n"))]))]
            [`(begin . ,expr*)
             (apply string-append (map (parse-function-body #t env level) expr*))]
            [`(set! ,var ,var1) #:when (eq? var var1) ""]
            [`(set! ,var ,val)
             (let ((val ((parse-function-body #f env level) val)))
               (if (equal? (safe var) reg-pc)
                   (pc2ts-append (tabs level) (safe var) " = " val ";\n")
                   (pc2ts-append (tabs level) (safe var) " = " val ";\n")))]
            [`(union-case ,val ,name . ,c*)
             (let ((template* (map car c*))
                   (body* (map get-body c*)))
               (if (not (check-union-case expr name
                                          (cadr (or (lookup-union name)
                                                    (error 'lookup-union 
                                                           "union type `~a' not defined ~n" name)))
                                          (map car template*)))
                   (error 'union-case "union-case doesn't match definition: `~a'\n"
                          name)
                   (letrec ([sname (safe name)]
                            [target_u_obj (safe val)]
                            [cases (apply string-append
                                          (map (handle-union-case-case name env target_u_obj)
                                               template*
                                               body*))])
                     (pc2ts-append
                      ; "    global " val "\n"
                      "    switch ((" target_u_obj " as UnionType).type) {\n"
                      cases
                      "    }\n"))))]
            [`(let ,bind* ,body) 
             (let ((lhs* (map car bind*))
                   (rhs* (map (parse-function-body #f env level) (map cadr bind*))))
               (pc2ts-append
                "\n"
                (apply string-append
                       (map (lambda (x y)
                              (pc2ts-append
                               "var " (safe x) " = " y ";\n"))
                            lhs* rhs*))
                body
                "\n"))]
            [`(printf ,str . ,parms*) 
             (string-append (tabs level)
                            "console.log(format("
                                  (join (cons (format "~s" str)
                                              (map safe
                                                   #;(λ (s)
                                                     (cond
                                                       [(is-global? s) (global s)]
                                                       [else (safe s)]))
                                                   parms*))
                                        ", ") "))\n")]
            [`(mount-trampoline ,construct ,dismount ,pc) 
             (set! construct-var (safe construct))
             (set! dismount-var (safe dismount))
             (pc2ts-append (tabs level)
                           "mount_tram();\n")]
            [`(dismount-trampoline ,dismount)
             (pc2ts-append (tabs level)
                           (safe dismount) "();\n")]
            [`(,func) #:when (is-func? func)
                      (pc2ts-append reg-pc " = " (safe func) ";\n")]
            [`,elsee
             (let ((elsee ((parse-function-body #f env level) elsee)))
               (pc2ts-append "return " elsee ";\n"))]
            ))
        (lambda (expr)
          (match expr
            ;; [(error ,name ,msg) 
            ;;  (pc2ts-append
            ;;   "fprintf(stderr, \"" msg "\");\n exit(1);\n")]
            [`#t  (pc2ts-append "true")]
            [`#f  (pc2ts-append "false")]
            
            [`,x #:when (symbol? x)
                 (safe (apply-env env x))
                 #;(letrec ([var (apply-env env x)]
                                             [safe-x (safe var)])
                                      (if (is-global? var)
                                          (begin
                                            (set! current-global-decls
                                                  (remove-duplicates
                                                   (cons (global var)
                                                         current-global-decls)))
                                            (global var))
                                          safe-x))]
            [`,x #:when (integer? x) (pc2ts-append x)]
            [`(zero? ,x) 
             (let ((x ((parse-function-body #f env level) x)))
               (pc2ts-append "(" x " === 0)"))]
            [`(and ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append "(" a " && " b ")"))]
            [`(or ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append "(" a " || " b ")"))]
            [`(not ,x) 
             (let ((x ((parse-function-body #f env level) x)))
               (pc2ts-append "(!" x ")"))]
            [`(< ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append "(" a " < " b ")"))]
            [`(> ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append "(" a " > " b ")"))]
            [`(<= ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append "(" a " <= " b ")"))]
            [`(>= ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append "(" a " >= " b ")"))]
            [`(+ ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append a " + " b))]
            [`(* ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append a " * " b))]
            [`(- ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append a " - " b))]
            [`(/ ,a ,b) 
             (let ((a ((parse-function-body #f env level) a))
                   (b ((parse-function-body #f env level) b)))
               (pc2ts-append a " / " b))]
            [`(sub1 ,a) 
             (let ((a ((parse-function-body #f env level) a)))
               (pc2ts-append "(" a " - 1)"))]
            [`(add1 ,a) 
             (let ((a ((parse-function-body #f env level) a)))
               (pc2ts-append "(" a " + 1)"))]
            [`(random ,x) 
             (let ((x ((parse-function-body #f env level) x)))
               (pc2ts-append "Math.trunc(Math.random() * " x ")"))]
            [`(if ,test ,conseq ,alt) 
             (let ((test ((parse-function-body #f env level) test))
                   (conseq ((parse-function-body #f env level) conseq))
                   (alt ((parse-function-body #f env level) alt)))
               (pc2ts-append "((" test ") ? (" conseq ") : (" alt "))"))]
            [`(,func . ,args*) #:when (symbol? func)
                               (let ((args* (map (parse-function-body #f env level) args*)))
                                 (pc2ts-append (tabs level)
                                               (safe func) "(" (join args* ", ") ")"))])))))

(define pc2ts-gen-funcs
  (lambda (env)
    (lambda (func)
      (let ([name (safe (car func))]
            [body (cadr func)])
        (if (equal? name "main")
            (begin
              (set! main-def (pc2ts-append
                              (pc2ts-append "function format(str: String, ... args: any[]) {\n"
                                            "    var i = 0;\n"
                                            "    return str.replace(/(~[savd])/g, function (match: Object) {\n"
                                            "        var argi = args[i];\n"
                                            "        i = i + 1;\n"
                                            "        return JSON.stringify(argi);\n"
                                            "    });\n"
                                            "}\n"
                                            "var _dismount_thunk: Function = function () {\n"
                                            "    " reg-pc " = undefined;\n"
                                            "};\n"
                                            )
                                       ((parse-function-body #t env 0) body)))
                              "")

            (begin
              (pc2ts-append
               (pc2ts-append "function " name "() {\n")
               ((parse-function-body #t env 1) body)
               "}\n\n")))))))

(define global-env
  (lambda ()
    (let loop ([env (empty-env)] 
               [reg (append (map car reg-funcs) reg-regs)])
      (if (null? reg)
          env
          (extend-env (car reg) (car reg) (loop env (cdr reg)))))))

(define pc2ts-source
  (lambda ()
    (let* ([s1 (apply string-append (map pc2ts-gen-unions reg-unions))]
           [s2 (apply string-append (map (pc2ts-gen-funcs (global-env)) reg-funcs))])
      (let ([s3 (pc2ts-append
                 "function mount_tram() {\n"
                 "    " dismount-var "= " construct-var "(_dismount_thunk)\n\n"
                 "    while (" reg-pc " !== undefined) {\n"
                 "        " reg-pc "();\n"
                 "    }\n"
                 "}\n\n")]
            [s4 (pc2ts-append
                 "function jumpout() {\n"
                 "    " reg-pc " = undefined\n"
                 "}")])
        (string-append
         "// Union functions\n"
         s1
         "// Generate functions\n"
         s2
         s3
         s4
         "\n"
         main-def)))))

(define pc2ts-header
  (lambda (decl*)
    (string-append
     (apply string-append
            (map pc2ts-header-parse decl*))
     "// Define the union classes\n"
     "type UnionType = { type: string } & Record<string, any>\n"
     "const UnionType = (tag: string, vals: Record<string, any>): UnionType => {\n"
     "    return { ...vals, type: tag }\n"
     "}\n"
     "namespace UnionEnums {\n"
     union-defs
     "}\n")))

(define pc2ts-header-parse
  (lambda (decl)
    (match decl
      [`(load ,file . ,file*)  ""]
      [`(exit)  ""]
      [`(display ,anything . ,anything*)  ""]
      [`(pretty-print ,anything . ,anything*)  ""]
      [`(define-registers . ,reg*) 
       (set! reg-regs reg*)
       (string-append
        "// Define the registers\n"
        (if (null? reg*)
            ""
            (string-append
             (join (map (λ (v)
                          (let ([global-new (safe v)])
                            #;(set! global-decls (cons global-new global-decls))
                            (string-append "var " global-new ": any;"))) reg*) "\n")
             "\nexport {}\n")))]
      [`(define-program-counter ,pc)
       (set! reg-pc (safe pc))
       (string-append
        "// Define the program counter\n"
        "var " reg-pc " : Function | undefined = undefined;\n\n")]
      [`(define-union ,name . ,c*) 
       (let ((tag* (map car c*))
             (field** (map cdr c*)))
         (add-union `(,name ,tag* ,field**))
         (let ([name (safe name)])
           (let ([enum-values
                  (apply string-append
                         (map
                          (lambda (tag field* index)
                            (let ([tag (safe tag)])
                              (format "        ~a = \"~a\",\n" tag (string-upcase
                                                                     (string-append
                                                                      name "_" tag)))))
                          tag* field** (range (length tag*))))])
             (set! union-defs
                   (string-append union-defs
                                  (pc2ts-append
                                   ;"class " name "_t(object):\n"
                                   "    export enum " name " {\n"
                                   enum-values
                                   "    }\n")))
             "")))]
      [`(define-label ,name ,body) ""
                                   (begin (add-func `(,name ,body))
                                          "" #;(string-append (if (equal? (safe name) "main")
                                                                  "int "
                                                                  "void ") (safe name) "();\n"))])))

(define pc2ts-fn-proto
  (lambda (fn-name param*)
    (let ([declare-params
           (lambda (param*)
             (join (map (lambda (param)
                          (format "~a : any" (safe param))) param*) ", "))])
      (pc2ts-append
       "function " (safe fn-name) "(" (declare-params param*)  ")"))))

(define compile/run
  (lambda (base-name)
    (let ([pc-file (string-append base-name ".pc")]
          [ts-file (string-append base-name ".ts")])
      (pc2ts pc-file ts-file)
      (system (string-append "ts-node ./" ts-file)))))