(local fennel (require :fennel))

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

(fn kv-table? [tbl] ; this should be exported to macros but is not ;_;
    (and (table? tbl) (not (sequence? tbl))))

(fn commasep [list f ?sep]
    (table.concat (icollect [_ v (ipairs list)] (f v)) (or ?sep ", ")))
(fn kvcommasep [tbl f ?sep]
    (commasep (. (getmetatable tbl) :keys) #(f $1 (. tbl $1)) ?sep))

(fn indent [str scope]
    (.. (string.rep " " scope.indent) str))

(fn indented [f scope ?f-will-indent]
    (set scope.indent (+ scope.indent 2))
    (local result (if ?f-will-indent (f) (indent (f) scope)))
    (set scope.indent (- scope.indent 2))
    result)

(fn block [pre stmts post scope f]
    (let [block (indented #(commasep stmts #(indent (f $1) scope) "\n") scope true)]
        (.. pre "\n" block "\n" (indent post scope))))

(local extract {})
(fn extract.quotes-in-table [into tbl inputs locals]
    (collect [k v (pairs tbl) &into into] k (extract.quotes v inputs locals)))

(fn extract.ensure-input [input inputs]
    (var found false)
    (each [_ sym (ipairs inputs)]
        (when (= sym input) (set found true)))
    (when (not found) (table.insert inputs input))
    input)

(fn extract.quotes [form inputs locals]
    (case form
        (where [q ref] (list? form) (sym? q) (sym? ref) (= (tostring q) :quote))
        (let [localref (. locals (tostring ref))]
            (if localref (extract.ensure-input (sym localref) inputs)
                (error (.. "Unknown local: " (tostring ref)))))

        (where l (list? l))     (extract.quotes-in-table (list) l inputs locals)
        (where s (sequence? s)) (extract.quotes-in-table (sequence) s inputs locals)
        (where t (kv-table? t)) (extract.quotes-in-table (setmetatable {} (getmetatable t)) t inputs locals)

        _ form))

(fn new-scope []
    {:locals {}
     :env {}
     :input []
     :indent 0
     :push (fn [self] (set self.locals (setmetatable {} {:__index self.locals})))
     :pop (fn [self] (set self.locals (. (getmetatable self.locals) :__index)))
     :with (fn [self f]
             (self:push)
             (local result (f))
             (self:pop)
             result)
     :expr (fn [self expr]
             (if (= expr (sym :nil))                                   :nil
                 (and (sym? expr) (not (multi-sym? expr)))             (self:env-ref expr)
                 (or (= (type expr) :number) (= (type expr) :boolean)) (tostring expr)
                 (let [name (safesym :inline-expr)
                       arglist (sequence)
                       expr (extract.quotes expr arglist self.locals)
                       bare-ref? (= (length arglist) 0)
                       expr (if bare-ref? expr `(fn ,arglist ,expr))
                       ref  (if bare-ref? name (.. name "(" (commasep arglist #(tostring $1)) ")"))]                
                    (table.insert self.input {: name : expr})
                    ref)))
     :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 comp {})
(local forms {})

(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] (list? form))
        (.. (comp.expr head scope) "(" (commasep rest #(comp.expr $1 scope)) ")")

        (where [array] (sequence? form) (= (length form) 1))
        (.. "@" (comp.expr array scope))

        (where [array & indexes] (sequence? form))
        (.. "(" (comp.expr array scope) ")" (table.concat (icollect [_ i (ipairs indexes)] (.. "[" (comp.expr i scope) "]"))))

        (where struct (kv-table? struct))
        (.. "{ " (kvcommasep struct #(.. (tostring $1) " = " (comp.expr $2 scope))) " }")

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

        (where _ (= form (sym :nil)))
        :nil

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

        (where lit (or (= (type lit) :string) (= (type lit) :number) (= (type lit) :nil) (= (type lit) :boolean)))
        (scope:expr lit)

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

(local type-constructors {})

(fn comp.type [typ scope]
    (case typ
        (where [innertyp index] (sequence? typ))
        (.. "(" (comp.type innertyp scope) ")[" (comp.type index scope) "]")

        (where [innertyp] (sequence? typ) (= (length typ) 1))
        (.. "&" (comp.type innertyp scope))

        (where struct (kv-table? struct))
        (.. "struct { " (kvcommasep struct #(comp.type-field $1 $2 scope)) " }")

        (where [head & args] (list? typ) (sym? head) (. type-constructors (tostring head)))
        ((. type-constructors (tostring head)) args scope)

        (where [head & args] (list? typ))
        (.. (comp.type head scope) "(" (commasep args #(comp.type $1 scope)) ")")

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

        _ (scope:expr typ)))

(fn comp.type-field [name typ scope]
    (case name
        (where name (sym? name) (not (multi-sym? name)))
        (.. (tostring name) " : " (comp.type typ scope))
        
        :union
        (.. "union { " (kvcommasep typ #(comp.type-field $1 $2 scope)) " }") ; unions shouldn't nest :/
        
        _ (error (.. "Invalid field name: " (view name)))))

(fn find-ival [tbl pred]
    (var ival nil)
    (each [i val (ipairs tbl)]
        (when (pred val) (set ival i)))
    ival)

(fn split-arglist [arglist]
    (let [iarg-return-sep (find-ival arglist #(and (sym? $1) (= (tostring $1) ":")))
          in  (icollect [i arg (ipairs arglist) &until (= i iarg-return-sep)] arg)
          out (when (not= iarg-return-sep nil) (icollect [i arg (ipairs arglist)] (when (> i iarg-return-sep) arg)))]
        (values in out)))

(fn arglist-type [arglist scope] (.. "{ " (commasep (or arglist []) #(comp.type $1 scope)) " }"))
(tset type-constructors :-> (fn [argtyps scope] 
    (let [(in out) (split-arglist argtyps)] 
        (.. (arglist-type in scope) " -> " (arglist-type (or out []) scope)))))
(tset type-constructors :&  (fn [[typ] scope]    (.. :& (comp.type typ scope))))
(tset type-constructors :$  (fn [typs scope]     (.. "tuple(" (commasep typs #(comp.type $1 scope)) ")")))
(fn type-constructors.unquote [[expr] scope] (scope:expr expr))

(fn forms.raw-escape [[text] scope] text)
(fn type-constructors.raw-escape [[text] scope] text)

(fn forms.var [defn scope]
    (fn rvaluelist [initvals]
        (if (> (length initvals) 0) 
            (.. " = " (commasep initvals #(comp.expr $1 scope))) 
            ""))
    (case defn
        (where [names & initvals] (list? names))
        (.. "var " (commasep names #(scope:addlocal $1)) (rvaluelist initvals))
        (where [nametyps & initvals] (kv-table? nametyps))
        (.. "var " (kvcommasep nametyps #(.. (scope:addlocal $1) " : " (comp.type $2 scope))) (rvaluelist initvals))
        [name typ initval] (.. "var " (scope:addlocal name) " : " (comp.type typ scope) " = " (comp.expr initval scope)) 
        [name initval] (.. "var " (scope:addlocal name) " = " (comp.expr initval scope))))

(fn forms.do [stmts scope]
    (scope:with #(block :do stmts :end scope #(comp.expr $1 scope))))

(fn comp.def [[arglist & stmts] scope]
    (if (> (length stmts) 0)
        (scope:with #(let [(in out) (split-arglist arglist)
                        argpairs (fcollect [i 1 (length in) 2] {:name (. in i) :type (. in (+ i 1))})
                        rettyp  (if out (.. ": { " (commasep out #(comp.type $1 scope)) " }") "")
                        argdefs (commasep argpairs #(.. (scope:addlocal $1.name) " : " (comp.type $1.type scope)))]
                        (block (.. "terra (" argdefs ")" rettyp) stmts :end scope #(comp.expr $1 scope))))
        (.. "(function () local terra anonfn :: " ((. type-constructors "->") arglist scope) " return anonfn end)()")))

(set forms.def comp.def)

(fn forms.return [vals scope] (.. "return " (commasep vals #(comp.expr $1 scope))))

(fn forms.cast [[typ expr] scope]
    (.. "([" (comp.type typ scope)"](" (comp.expr expr scope) "))"))
(fn forms.unquote [[expr] scope] (.. "([" (scope:expr expr) "])"))

(tset forms :& (fn [[expr] scope] (.. :& (comp.expr expr scope))))
(tset forms :$ (fn [items scope] (.. "{ " (commasep items #(comp.expr $1 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  :^ :<< :>>])]
    (def-infix op op))
(def-infix := :==)
(def-infix :not= "~=")
(fn forms.not [[expr] scope] (.. "not (" (comp.expr expr scope) ")"))
(fn forms.set [[left right] scope]
    (if (list? left) (.. (commasep left #(comp.expr $1 scope)) " = (" (comp.expr right scope) ")")
        (.. (comp.expr left scope) " = (" (comp.expr right scope) ")")))
(fn forms.tset [args scope]
    (let [iright (length args)
          right (. args iright)]
        (tset args iright nil)
        (.. ((. forms ".") args scope) " = (" (comp.expr right scope) ")")))

(tset forms "." (fn [[struct & fields] scope]
    (.. (comp.expr struct scope) (commasep fields #(.. ".[" (scope:expr $1) "]") ""))))
(tset forms ":" (fn [[obj field & args] scope]
    (.. (comp.expr obj scope) ":[" (scope:expr field) "](" (commasep args #(comp.expr $1 scope)) ")")))

(fn forms.if [params scope]
    (let [has-else? (= (% (length params) 2) 1)
          ielse (when has-else? (length params))
          clauses (fcollect [i 1 (length params) 2]
                    (if (= i ielse) 
                        {:pre (indent :else scope)
                         :clause (. params i)}
                        {:pre (.. (if (= i 1) :if (indent :elseif scope)) " " (comp.expr (. params i) scope) " then")
                         :clause (. params (+ i 1))}))
          blocks (icollect [_ clause (ipairs clauses)]
                    (scope:with #(.. clause.pre "\n" (indented #(comp.expr clause.clause scope) scope))))]
        (.. (table.concat blocks "\n") "\n" (indent :end scope))))

(fn comp.quote [stmts scope]
    (if (= (length stmts) 0) (error "nothing to quote")
        (= (length stmts) 1) (.. "`(" (comp.expr (. stmts 1) scope) ")")
        (let [last-expr (. stmts (length stmts))]
            (tset stmts (length stmts) nil)
            (.. (block :quote stmts :in scope #(comp.expr $1 scope)) "\n"
                (indented #(comp.expr last-expr scope) scope) "\n"
                (indent :end scope)))))

(fn comp.global [[initial-val] scope]
    (.. "global(" (scope:expr initial-val) ")"))

(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 ttype [expr] (build expr comp.type))
(fn tquote [...] (build [...] comp.quote))
(fn def [...] (build [...] comp.def))
(fn static [initial-value] (build [initial-value] comp.global))
(fn eval [expr] `(,(def `[] `(return ,expr))))

(fn compile [compiler-name expr] (build expr (. comp (tostring compiler-name))))
(fn uncompile [compiler-name expr]
    (let [scope (new-scope)
          compiler (. comp (tostring compiler-name))]
        (compiler expr scope)))

(fn printform [form] `(do (print ,(view (macroexpand form))) ,form))

{: ttype :q tquote : def : static : eval : compile : uncompile : printform}