(local fennel (require :fennel))

(fn safesym [name]
    (let [str (tostring (gensym name))]
        (str:gsub "[^%w_]" #(string.format "_%02x" ($:byte)))))

(fn new-scope []
    {:locals {}
     :env {}
     :input []
     :push (fn [self] (set self.locals (setmetatable {} {:__index self.locals})))
     :pop (fn [self] (set self.locals (. (getmetatable self.locals) :__index)))
     :expr (fn [self expr]
        (let [name (safesym :inline-expr)]
            (table.insert self.input {: name : expr})
            name))
     :env-ref (fn [self symbol]
        (let [name (tostring symbol)
              loc (. self.env name)]
            (if loc loc
                (let [envloc (safesym name)]
                    (tset self.env name envloc)
                    (table.insert self.input {:name envloc :expr (sym name)})
                    envloc))))
     :reference (fn [self symbol]
        (let [name (tostring symbol)
              loc (. self.locals name)]
            (if loc loc (self:env-ref symbol))))
     :addlocal (fn [self symbol]
        (let [name (tostring symbol)
              loc (safesym name)]
            (tset self.locals name loc)
            loc))})

(local forms {})

(fn commasep [list f]
    (table.concat (icollect [_ v (ipairs list)] (f v)) ", "))

(fn comp-expr [form scope]
    (case form
        (where [head & rest] (sym? head) (list? form) (. forms (tostring head))) 
        ((. forms (tostring head)) rest scope)

        (where [head & rest] (sym? head) (list? form)) 
        (.. (comp-expr head scope) "(" (commasep rest #(comp-expr $1 scope)) ")")

        (where [escape] (sequence? form))
        (scope:expr escape)

        (where str (= (type str) :string))
        (scope:expr str)

        (where num (= (type num) :number))
        (tostring num)

        (where multisym (multi-sym? multisym))
        (string.gsub (tostring multisym) "[^.:]+" #(scope:reference $1) 1)

        (where symbol (sym? symbol))
        (scope:reference symbol)

        nil           
        "nil"
        
        _ (error (.. "Failed to parse expression: " (view form)))))

(fn comp-type [typ scope] 
    (case typ
        (where typlist (sequence? typlist))
        (.. "{" (commasep typlist #(comp-type $1 scope)) "}")

        (where [vec typ n] (list? typ) (= (type n) :number) (= vec (sym :vec)))
        (.. "vector(" (comp-type typ scope) ", " (tostring n) ")")

        (where tuple (list? tuple))
        (.. "tuple(" (commasep tuple #(comp-type $1 scope)) ")")

        (where [to in out] (list? typ) (= to (sym :->)))
        (.. (comp-type in scope) " -> " (comp-type out scope))

        (where symbol (sym? symbol))
        (tostring symbol)

        (where symbol (= (type symbol) :string))
        symbol
        
        _ (error (.. "Failed to parse type: " (view typ)))))

(fn forms.rawterra [[text] scope] text)
(fn forms.var [[name typ] scope]
    (.. "var " (scope:addlocal name) " " (comp-type typ scope)))

(fn forms.do [stmts scope]
    (scope:push)
    (local lines (icollect [_ stmt (ipairs stmts)] (comp-expr stmt scope)))
    (scope:pop)
    (table.concat lines "\n"))

(fn forms.def [[arglist & stmts] scope]
    (scope:push)
    (let [argpairs (fcollect [i 1 (length arglist) 2] {:name (. arglist i) :type (. arglist (+ i 1))})
          _ (print (view argpairs))
          argdefs (commasep argpairs #(.. (scope:addlocal $1.name) " : " (comp-type $1.type scope)))
          fulldef (.. "terra (" argdefs ")\n" (forms.do stmts scope) "\nend")]
        (scope:pop)
        fulldef))
(fn forms.return [[expr] scope] (.. "return " (comp-expr expr scope)))
(fn def-infix [fnlop luaop]
    (tset forms fnlop (fn [[left right] scope]
        (.. "(" (comp-expr left scope) " " luaop " " (comp-expr right scope) ")"))))
(each [_ op (ipairs [:+ :- :* :/ :% :< :<= :> :>= :and :or :not :^ :<< :>>])]
    (def-infix op op))
(fn forms.cast [[typ expr] scope]
    (.. "([" (comp-type typ scope)"]" (comp-expr expr scope) ")"))

(def-infix := :==)
(def-infix :not= "~=")
(def-infix :set :=)

(fn build [expr compiler]
    (let [scope (new-scope)
          terra-expr (compiler expr scope)
          env-setup (if (= (length scope.input) 0) ""
                        (.. "local " (commasep scope.input #$1.name) " = ...\n"))
          compiled (.. env-setup "return " terra-expr)
          call `((terralib.loadstring ,compiled))]
        (icollect [_ {: expr} (ipairs scope.input) &into call] expr)
        call))

(fn terra [expr] (build expr comp-expr))
(fn ttype [expr] (build expr comp-type))

(fn def [...] (terra `(,(sym :def) ,...)))
(fn q [...] (terra `(,(sym :quote) ,...)))
(fn unterra [...] (view (macroexpand (terra ...))))

{: terra : ttype : def : q : unterra}