(define lua-tok-type (fn (t) (if (= t nil) "eof" (get t :type))))

(define lua-tok-value (fn (t) (if (= t nil) nil (get t :value))))

(define
  lua-binop-prec
  (fn
    (op)
    (cond
      ((= op "or") 1)
      ((= op "and") 2)
      ((= op "<") 3)
      ((= op ">") 3)
      ((= op "<=") 3)
      ((= op ">=") 3)
      ((= op "==") 3)
      ((= op "~=") 3)
      ((= op "..") 5)
      ((= op "+") 6)
      ((= op "-") 6)
      ((= op "*") 7)
      ((= op "/") 7)
      ((= op "%") 7)
      ((= op "^") 10)
      (else 0))))

(define lua-binop-right? (fn (op) (or (= op "..") (= op "^"))))

(define
  lua-parse
  (fn
    (src)
    (let
      ((tokens (lua-tokenize src)) (idx 0) (tok-len 0))
      (begin
        (set! tok-len (len tokens))
        (define peek-tok (fn () (nth tokens idx)))
        (define
          peek-tok-at
          (fn (n) (if (< (+ idx n) tok-len) (nth tokens (+ idx n)) nil)))
        (define advance-tok! (fn () (set! idx (+ idx 1))))
        (define
          check-tok?
          (fn
            (type value)
            (let
              ((t (peek-tok)))
              (and
                (= (lua-tok-type t) type)
                (or (= value nil) (= (lua-tok-value t) value))))))
        (define
          consume!
          (fn
            (type value)
            (if
              (check-tok? type value)
              (let ((t (peek-tok))) (begin (advance-tok!) t))
              (error
                (str
                  "lua-parse: expected "
                  type
                  " "
                  value
                  " got "
                  (lua-tok-type (peek-tok))
                  " "
                  (lua-tok-value (peek-tok)))))))
        (define at-keyword? (fn (kw) (check-tok? "keyword" kw)))
        (define at-op? (fn (op) (check-tok? "op" op)))
        (define
          at-binop?
          (fn
            ()
            (let
              ((t (peek-tok)))
              (cond
                ((and (= (lua-tok-type t) "keyword") (or (= (lua-tok-value t) "and") (= (lua-tok-value t) "or")))
                  true)
                ((and (= (lua-tok-type t) "op") (> (lua-binop-prec (lua-tok-value t)) 0))
                  true)
                (else false)))))
        (define parse-expr nil)
        (define parse-block nil)
        (define parse-prefixexp nil)
        (define parse-table-ctor nil)
        (define parse-unary nil)
        (define parse-args nil)
        (define parse-funcbody nil)
        (set!
          parse-args
          (fn
            ()
            (cond
              ((at-op? "(")
                (begin
                  (advance-tok!)
                  (if
                    (at-op? ")")
                    (begin (advance-tok!) (list))
                    (let
                      ((args (list (parse-expr))))
                      (begin
                        (define
                          more
                          (fn
                            ()
                            (when
                              (at-op? ",")
                              (begin
                                (advance-tok!)
                                (append! args (parse-expr))
                                (more)))))
                        (more)
                        (consume! "op" ")")
                        args)))))
              ((at-op? "{") (list (parse-table-ctor)))
              ((check-tok? "string" nil)
                (let
                  ((t (peek-tok)))
                  (begin
                    (advance-tok!)
                    (list (list (quote lua-str) (lua-tok-value t))))))
              (else (error "lua-parse: expected args")))))
        (set!
          parse-funcbody
          (fn
            ()
            (begin
              (consume! "op" "(")
              (let
                ((params (list)) (is-vararg false))
                (begin
                  (when
                    (not (at-op? ")"))
                    (begin
                      (cond
                        ((at-op? "...")
                          (begin (advance-tok!) (set! is-vararg true)))
                        (else
                          (let
                            ((t (peek-tok)))
                            (begin
                              (when
                                (not (= (lua-tok-type t) "ident"))
                                (error "lua-parse: expected name in params"))
                              (append! params (lua-tok-value t))
                              (advance-tok!)))))
                      (define
                        more-params
                        (fn
                          ()
                          (when
                            (and (at-op? ",") (not is-vararg))
                            (begin
                              (advance-tok!)
                              (cond
                                ((at-op? "...")
                                  (begin
                                    (advance-tok!)
                                    (set! is-vararg true)))
                                (else
                                  (let
                                    ((t (peek-tok)))
                                    (begin
                                      (when
                                        (not (= (lua-tok-type t) "ident"))
                                        (error "lua-parse: expected name"))
                                      (append! params (lua-tok-value t))
                                      (advance-tok!)))))
                              (more-params)))))
                      (more-params)))
                  (consume! "op" ")")
                  (let
                    ((body (parse-block)))
                    (begin
                      (consume! "keyword" "end")
                      (list (quote lua-function) params is-vararg body))))))))
        (define
          parse-primary
          (fn
            ()
            (let
              ((t (peek-tok)))
              (cond
                ((= (lua-tok-type t) "number")
                  (begin
                    (advance-tok!)
                    (list (quote lua-num) (lua-tok-value t))))
                ((= (lua-tok-type t) "string")
                  (begin
                    (advance-tok!)
                    (list (quote lua-str) (lua-tok-value t))))
                ((and (= (lua-tok-type t) "keyword") (= (lua-tok-value t) "nil"))
                  (begin (advance-tok!) (list (quote lua-nil))))
                ((and (= (lua-tok-type t) "keyword") (= (lua-tok-value t) "true"))
                  (begin (advance-tok!) (list (quote lua-true))))
                ((and (= (lua-tok-type t) "keyword") (= (lua-tok-value t) "false"))
                  (begin (advance-tok!) (list (quote lua-false))))
                ((and (= (lua-tok-type t) "op") (= (lua-tok-value t) "..."))
                  (begin (advance-tok!) (list (quote lua-vararg))))
                ((and (= (lua-tok-type t) "keyword") (= (lua-tok-value t) "function"))
                  (begin (advance-tok!) (parse-funcbody)))
                ((and (= (lua-tok-type t) "op") (= (lua-tok-value t) "{"))
                  (parse-table-ctor))
                ((or (= (lua-tok-type t) "ident") (and (= (lua-tok-type t) "op") (= (lua-tok-value t) "(")))
                  (parse-prefixexp))
                (else
                  (error
                    (str
                      "lua-parse: unexpected "
                      (lua-tok-type t)
                      " "
                      (lua-tok-value t))))))))
        (define parse-pow-chain
          (fn () (let ((lhs (parse-primary))) (parse-binop-rhs 10 lhs))))
        (set!
          parse-unary
          (fn
            ()
            (cond
              ((at-op? "-")
                (begin
                  (advance-tok!)
                  (list (quote lua-unop) "-" (parse-unary))))
              ((at-op? "#")
                (begin
                  (advance-tok!)
                  (list (quote lua-unop) "#" (parse-unary))))
              ((at-keyword? "not")
                (begin
                  (advance-tok!)
                  (list (quote lua-unop) "not" (parse-unary))))
              (else (parse-pow-chain)))))
        (define
          parse-binop-rhs
          (fn
            (min-prec lhs)
            (begin
              (define
                loop
                (fn
                  (lhs-acc)
                  (if
                    (at-binop?)
                    (let
                      ((op (lua-tok-value (peek-tok))))
                      (let
                        ((prec (lua-binop-prec op)))
                        (if
                          (< prec min-prec)
                          lhs-acc
                          (begin
                            (advance-tok!)
                            (let
                              ((rhs (parse-unary)))
                              (let
                                ((next-min (if (lua-binop-right? op) prec (+ prec 1))))
                                (let
                                  ((rhs2 (parse-binop-rhs next-min rhs)))
                                  (loop
                                    (list (quote lua-binop) op lhs-acc rhs2)))))))))
                    lhs-acc)))
              (loop lhs)))))
      (set!
        parse-expr
        (fn () (let ((lhs (parse-unary))) (parse-binop-rhs 1 lhs))))
      (set!
        parse-prefixexp
        (fn
          ()
          (let
            ((base nil))
            (begin
              (cond
                ((check-tok? "ident" nil)
                  (let
                    ((t (peek-tok)))
                    (begin
                      (advance-tok!)
                      (set! base (list (quote lua-name) (lua-tok-value t))))))
                ((at-op? "(")
                  (begin
                    (advance-tok!)
                    (set! base (list (quote lua-paren) (parse-expr)))
                    (consume! "op" ")")))
                (else (error "lua-parse: expected prefixexp")))
              (define
                more
                (fn
                  ()
                  (cond
                    ((at-op? ".")
                      (begin
                        (advance-tok!)
                        (let
                          ((t (peek-tok)))
                          (begin
                            (when
                              (not (= (lua-tok-type t) "ident"))
                              (error "lua-parse: expected name after ."))
                            (set!
                              base
                              (list (quote lua-field) base (lua-tok-value t)))
                            (advance-tok!)
                            (more)))))
                    ((at-op? "[")
                      (begin
                        (advance-tok!)
                        (let
                          ((e (parse-expr)))
                          (begin
                            (consume! "op" "]")
                            (set! base (list (quote lua-index) base e))
                            (more)))))
                    ((at-op? ":")
                      (begin
                        (advance-tok!)
                        (let
                          ((t (peek-tok)))
                          (begin
                            (when
                              (not (= (lua-tok-type t) "ident"))
                              (error "lua-parse: expected name after :"))
                            (let
                              ((name (lua-tok-value t)))
                              (begin
                                (advance-tok!)
                                (let
                                  ((args (parse-args)))
                                  (begin
                                    (set!
                                      base
                                      (list
                                        (quote lua-method-call)
                                        base
                                        name
                                        args))
                                    (more)))))))))
                    ((or (at-op? "(") (at-op? "{") (check-tok? "string" nil))
                      (let
                        ((args (parse-args)))
                        (begin
                          (set! base (list (quote lua-call) base args))
                          (more))))
                    (else nil))))
              (more)
              base))))
      (set!
        parse-table-ctor
        (fn
          ()
          (begin
            (consume! "op" "{")
            (let
              ((fields (list)))
              (begin
                (define
                  parse-field
                  (fn
                    ()
                    (cond
                      ((at-op? "[")
                        (begin
                          (advance-tok!)
                          (let
                            ((k (parse-expr)))
                            (begin
                              (consume! "op" "]")
                              (consume! "op" "=")
                              (let
                                ((v (parse-expr)))
                                (append! fields (list (quote lua-kv) k v)))))))
                      ((and (check-tok? "ident" nil) (let ((nxt (peek-tok-at 1))) (and (not (= nxt nil)) (= (lua-tok-type nxt) "op") (= (lua-tok-value nxt) "="))))
                        (let
                          ((t (peek-tok)))
                          (begin
                            (advance-tok!)
                            (advance-tok!)
                            (let
                              ((v (parse-expr)))
                              (append!
                                fields
                                (list
                                  (quote lua-kv)
                                  (list (quote lua-str) (lua-tok-value t))
                                  v))))))
                      (else
                        (let
                          ((v (parse-expr)))
                          (append! fields (list (quote lua-pos) v)))))))
                (when
                  (not (at-op? "}"))
                  (begin
                    (parse-field)
                    (define
                      more-fields
                      (fn
                        ()
                        (when
                          (and
                            (or (at-op? ",") (at-op? ";"))
                            (not (at-op? "}")))
                          (begin
                            (advance-tok!)
                            (when
                              (not (at-op? "}"))
                              (begin (parse-field) (more-fields)))))))
                    (more-fields)))
                (consume! "op" "}")
                (cons (quote lua-table) fields))))))
      (define
        parse-explist
        (fn
          ()
          (let
            ((exps (list (parse-expr))))
            (begin
              (define
                more
                (fn
                  ()
                  (when
                    (at-op? ",")
                    (begin
                      (advance-tok!)
                      (append! exps (parse-expr))
                      (more)))))
              (more)
              exps))))
      (define
        parse-namelist
        (fn
          ()
          (let
            ((names (list)))
            (begin
              (let
                ((t (peek-tok)))
                (begin
                  (when
                    (not (= (lua-tok-type t) "ident"))
                    (error "lua-parse: expected name"))
                  (append! names (lua-tok-value t))
                  (advance-tok!)))
              (define
                more
                (fn
                  ()
                  (when
                    (at-op? ",")
                    (begin
                      (advance-tok!)
                      (let
                        ((t (peek-tok)))
                        (begin
                          (when
                            (not (= (lua-tok-type t) "ident"))
                            (error "lua-parse: expected name"))
                          (append! names (lua-tok-value t))
                          (advance-tok!)))
                      (more)))))
              (more)
              names))))
      (define
        parse-if
        (fn
          ()
          (begin
            (consume! "keyword" "if")
            (let
              ((cnd (parse-expr)))
              (begin
                (consume! "keyword" "then")
                (let
                  ((then-body (parse-block))
                    (elseifs (list))
                    (else-body nil))
                  (begin
                    (define
                      elseif-loop
                      (fn
                        ()
                        (when
                          (at-keyword? "elseif")
                          (begin
                            (advance-tok!)
                            (let
                              ((c (parse-expr)))
                              (begin
                                (consume! "keyword" "then")
                                (let
                                  ((b (parse-block)))
                                  (begin
                                    (append! elseifs (list c b))
                                    (elseif-loop)))))))))
                    (elseif-loop)
                    (when
                      (at-keyword? "else")
                      (begin (advance-tok!) (set! else-body (parse-block))))
                    (consume! "keyword" "end")
                    (list (quote lua-if) cnd then-body elseifs else-body))))))))
      (define
        parse-while
        (fn
          ()
          (begin
            (consume! "keyword" "while")
            (let
              ((cnd (parse-expr)))
              (begin
                (consume! "keyword" "do")
                (let
                  ((body (parse-block)))
                  (begin
                    (consume! "keyword" "end")
                    (list (quote lua-while) cnd body))))))))
      (define
        parse-repeat
        (fn
          ()
          (begin
            (consume! "keyword" "repeat")
            (let
              ((body (parse-block)))
              (begin
                (consume! "keyword" "until")
                (let
                  ((cnd (parse-expr)))
                  (list (quote lua-repeat) body cnd)))))))
      (define
        parse-do
        (fn
          ()
          (begin
            (consume! "keyword" "do")
            (let
              ((body (parse-block)))
              (begin (consume! "keyword" "end") (list (quote lua-do) body))))))
      (define parse-for-num-rest
        (fn (name)
          (begin
            (consume! "op" "=")
            (let ((start (parse-expr)))
              (begin
                (consume! "op" ",")
                (let ((stop (parse-expr)) (step nil))
                  (begin
                    (when (at-op? ",")
                      (begin
                        (advance-tok!)
                        (set! step (parse-expr))))
                    (consume! "keyword" "do")
                    (let ((body (parse-block)))
                      (begin
                        (consume! "keyword" "end")
                        (list (quote lua-for-num) name start stop step body))))))))))
      (define parse-for-in-names
        (fn (names)
          (cond
            ((at-op? ",")
              (begin
                (advance-tok!)
                (let ((nt (peek-tok)))
                  (begin
                    (when (not (= (lua-tok-type nt) "ident"))
                      (error "lua-parse: expected name after , in for"))
                    (let ((nm (lua-tok-value nt)))
                      (begin
                        (advance-tok!)
                        (parse-for-in-names (append names (list nm)))))))))
            (else names))))
      (define parse-for-in-exps
        (fn (exps)
          (cond
            ((at-op? ",")
              (begin
                (advance-tok!)
                (parse-for-in-exps (append exps (list (parse-expr))))))
            (else exps))))
      (define parse-for-in-rest
        (fn (names)
          (begin
            (consume! "keyword" "in")
            (let ((exps (parse-for-in-exps (list (parse-expr)))))
              (begin
                (consume! "keyword" "do")
                (let ((body (parse-block)))
                  (begin
                    (consume! "keyword" "end")
                    (list (quote lua-for-in) names exps body))))))))
      (define parse-for
        (fn ()
          (begin
            (consume! "keyword" "for")
            (let ((t (peek-tok)))
              (begin
                (when (not (= (lua-tok-type t) "ident"))
                  (error "lua-parse: expected name in for"))
                (let ((name (lua-tok-value t)))
                  (begin
                    (advance-tok!)
                    (cond
                      ((at-op? "=") (parse-for-num-rest name))
                      (else
                        (parse-for-in-rest (parse-for-in-names (list name))))))))))))
      (define
        parse-funcname
        (fn
          ()
          (let
            ((t (peek-tok)))
            (begin
              (when
                (not (= (lua-tok-type t) "ident"))
                (error "lua-parse: expected name after function"))
              (let
                ((base (list (quote lua-name) (lua-tok-value t))))
                (begin
                  (advance-tok!)
                  (define
                    dots
                    (fn
                      ()
                      (when
                        (at-op? ".")
                        (begin
                          (advance-tok!)
                          (let
                            ((tt (peek-tok)))
                            (begin
                              (when
                                (not (= (lua-tok-type tt) "ident"))
                                (error "lua-parse: expected name after ."))
                              (set!
                                base
                                (list
                                  (quote lua-field)
                                  base
                                  (lua-tok-value tt)))
                              (advance-tok!)
                              (dots)))))))
                  (dots)
                  (let
                    ((is-method false) (method-name nil))
                    (begin
                      (when
                        (at-op? ":")
                        (begin
                          (advance-tok!)
                          (let
                            ((tt (peek-tok)))
                            (begin
                              (when
                                (not (= (lua-tok-type tt) "ident"))
                                (error "lua-parse: expected name after :"))
                              (set! is-method true)
                              (set! method-name (lua-tok-value tt))
                              (advance-tok!)))))
                      (list base is-method method-name)))))))))
      (define
        parse-function-decl
        (fn
          ()
          (begin
            (consume! "keyword" "function")
            (let
              ((parts (parse-funcname)))
              (let
                ((base (nth parts 0))
                  (is-method (nth parts 1))
                  (m-name (nth parts 2)))
                (let
                  ((func (parse-funcbody)))
                  (if
                    is-method
                    (let
                      ((target (list (quote lua-field) base m-name))
                        (new-params (cons "self" (nth func 1))))
                      (list
                        (quote lua-function-decl)
                        target
                        (list
                          (quote lua-function)
                          new-params
                          (nth func 2)
                          (nth func 3))))
                    (list (quote lua-function-decl) base func))))))))
      (define
        parse-local
        (fn
          ()
          (begin
            (consume! "keyword" "local")
            (cond
              ((at-keyword? "function")
                (begin
                  (advance-tok!)
                  (let
                    ((t (peek-tok)))
                    (begin
                      (when
                        (not (= (lua-tok-type t) "ident"))
                        (error "lua-parse: expected name after local fn"))
                      (let
                        ((name (lua-tok-value t)))
                        (begin
                          (advance-tok!)
                          (let
                            ((func (parse-funcbody)))
                            (list (quote lua-local-function) name func))))))))
              (else
                (let
                  ((names (parse-namelist)) (exps (list)))
                  (begin
                    (when
                      (at-op? "=")
                      (begin (advance-tok!) (set! exps (parse-explist))))
                    (list (quote lua-local) names exps))))))))
      (define
        parse-return
        (fn
          ()
          (begin
            (consume! "keyword" "return")
            (let
              ((exps (list)))
              (begin
                (when
                  (not
                    (or
                      (at-keyword? "end")
                      (at-keyword? "else")
                      (at-keyword? "elseif")
                      (at-keyword? "until")
                      (check-tok? "eof" nil)
                      (at-op? ";")))
                  (set! exps (parse-explist)))
                (when (at-op? ";") (advance-tok!))
                (list (quote lua-return) exps))))))
      (define
        parse-assign-or-call
        (fn
          ()
          (let
            ((pexpr (parse-prefixexp)))
            (cond
              ((or (at-op? "=") (at-op? ","))
                (let
                  ((lhs (list pexpr)))
                  (begin
                    (define
                      more
                      (fn
                        ()
                        (when
                          (at-op? ",")
                          (begin
                            (advance-tok!)
                            (append! lhs (parse-prefixexp))
                            (more)))))
                    (more)
                    (consume! "op" "=")
                    (let
                      ((rhs (parse-explist)))
                      (list (quote lua-assign) lhs rhs)))))
              ((or (= (first pexpr) (quote lua-call)) (= (first pexpr) (quote lua-method-call)))
                (list (quote lua-call-stmt) pexpr))
              (else (error "lua-parse: expected '=' or call"))))))
      (define
        parse-stmt
        (fn
          ()
          (cond
            ((at-op? ";") (begin (advance-tok!) nil))
            ((at-keyword? "if") (parse-if))
            ((at-keyword? "while") (parse-while))
            ((at-keyword? "repeat") (parse-repeat))
            ((at-keyword? "do") (parse-do))
            ((at-keyword? "for") (parse-for))
            ((at-keyword? "function") (parse-function-decl))
            ((at-keyword? "local") (parse-local))
            ((at-keyword? "return") (parse-return))
            ((at-keyword? "break")
              (begin (advance-tok!) (list (quote lua-break))))
            (else (parse-assign-or-call)))))
      (define
        block-end?
        (fn
          ()
          (or
            (at-keyword? "end")
            (at-keyword? "else")
            (at-keyword? "elseif")
            (at-keyword? "until")
            (check-tok? "eof" nil))))
      (set!
        parse-block
        (fn
          ()
          (let
            ((stmts (list)))
            (begin
              (define
                loop
                (fn
                  ()
                  (when
                    (not (block-end?))
                    (let
                      ((s (parse-stmt)))
                      (begin
                        (when (not (= s nil)) (append! stmts s))
                        (cond
                          ((and (not (= s nil)) (= (first s) (quote lua-return)))
                            nil)
                          (else (loop))))))))
              (loop)
              (cons (quote lua-block) stmts)))))
      (parse-block))))


(define
  lua-parse-expr
  (fn
    (src)
    (let
      ((tree (lua-parse (str "return " src))))
      (let ((ret (nth tree 1))) (first (nth ret 1))))))
