(.module:
  [lux #*
   [control
    monad]
   [data
    [text ("text/" Equivalence<Text>)
     format]
    [maybe ("maybe/" Monad<Maybe>)]
    [product]
    ["e" error]
    [collection
     [list ("list/" Functor<List> Fold<List> Monoid<List>)]
     [dictionary [plist]]]]]
  [// (#+ Operation Compiler)
   ["compiler" //
    [extension]
    [//
     [reference (#+ Register Variable)]]]])

(type: Local (Bindings Text [Type Register]))
(type: Foreign (Bindings Text [Type Variable]))

(def: (local? name scope)
  (-> Text Scope Bit)
  (|> scope
      (get@ [#.locals #.mappings])
      (plist.contains? name)))

(def: (local name scope)
  (-> Text Scope (Maybe [Type Variable]))
  (|> scope
      (get@ [#.locals #.mappings])
      (plist.get name)
      (maybe/map (function (_ [type value])
                   [type (#reference.Local value)]))))

(def: (captured? name scope)
  (-> Text Scope Bit)
  (|> scope
      (get@ [#.captured #.mappings])
      (plist.contains? name)))

(def: (captured name scope)
  (-> Text Scope (Maybe [Type Variable]))
  (loop [idx +0
         mappings (get@ [#.captured #.mappings] scope)]
    (case mappings
      #.Nil
      #.None

      (#.Cons [_name [_source-type _source-ref]] mappings')
      (if (text/= name _name)
        (#.Some [_source-type (#reference.Foreign idx)])
        (recur (inc idx) mappings')))))

(def: (reference? name scope)
  (-> Text Scope Bit)
  (or (local? name scope)
      (captured? name scope)))

(def: (reference name scope)
  (-> Text Scope (Maybe [Type Variable]))
  (case (..local name scope)
    (#.Some type)
    (#.Some type)

    _
    (..captured name scope)))

(def: #export (find name)
  (-> Text (Operation (Maybe [Type Variable])))
  (extension.lift
   (function (_ state)
     (let [[inner outer] (|> state
                             (get@ #.scopes)
                             (list.split-with (|>> (reference? name) not)))]
       (case outer
         #.Nil
         (#.Right [state #.None])

         (#.Cons top-outer _)
         (let [[ref-type init-ref] (maybe.default (undefined)
                                                  (..reference name top-outer))
               [ref inner'] (list/fold (: (-> Scope [Variable (List Scope)] [Variable (List Scope)])
                                          (function (_ scope ref+inner)
                                            [(#reference.Foreign (get@ [#.captured #.counter] scope))
                                             (#.Cons (update@ #.captured
                                                              (: (-> Foreign Foreign)
                                                                 (|>> (update@ #.counter inc)
                                                                      (update@ #.mappings (plist.put name [ref-type (product.left ref+inner)]))))
                                                              scope)
                                                     (product.right ref+inner))]))
                                       [init-ref #.Nil]
                                       (list.reverse inner))
               scopes (list/compose inner' outer)]
           (#.Right [(set@ #.scopes scopes state)
                     (#.Some [ref-type ref])]))
         )))))

(def: #export (with-local [name type] action)
  (All [a] (-> [Text Type] (Operation a) (Operation a)))
  (function (_ [bundle state])
    (case (get@ #.scopes state)
      (#.Cons head tail)
      (let [old-mappings (get@ [#.locals #.mappings] head)
            new-var-id (get@ [#.locals #.counter] head)
            new-head (update@ #.locals
                              (: (-> Local Local)
                                 (|>> (update@ #.counter inc)
                                      (update@ #.mappings (plist.put name [type new-var-id]))))
                              head)]
        (case (compiler.run' [bundle (set@ #.scopes (#.Cons new-head tail) state)]
                             action)
          (#e.Success [[bundle' state'] output])
          (case (get@ #.scopes state')
            (#.Cons head' tail')
            (let [scopes' (#.Cons (set@ #.locals (get@ #.locals head) head')
                                  tail')]
              (#e.Success [[bundle' (set@ #.scopes scopes' state')]
                           output]))

            _
            (error! "Invalid scope alteration/"))

          (#e.Error error)
          (#e.Error error)))

      _
      (#e.Error "Cannot create local binding without a scope."))
    ))

(do-template [<name> <val-type>]
  [(def: <name>
     (Bindings Text [Type <val-type>])
     {#.counter +0
      #.mappings (list)})]

  [init-locals   Nat]
  [init-captured Variable]
  )

(def: (scope parent-name child-name)
  (-> (List Text) Text Scope)
  {#.name     (list& child-name parent-name)
   #.inner    +0
   #.locals   init-locals
   #.captured init-captured})

(def: #export (with-scope name action)
  (All [a] (-> Text (Operation a) (Operation a)))
  (function (_ [bundle state])
    (let [parent-name (case (get@ #.scopes state)
                        #.Nil
                        (list)
                        
                        (#.Cons top _)
                        (get@ #.name top))]
      (case (action [bundle (update@ #.scopes
                                     (|>> (#.Cons (scope parent-name name)))
                                     state)])
        (#e.Error error)
        (#e.Error error)

        (#e.Success [[bundle' state'] output])
        (#e.Success [[bundle' (update@ #.scopes
                                       (|>> list.tail (maybe.default (list)))
                                       state')]
                     output])
        ))
    ))

(def: #export next-local
  (Operation Register)
  (extension.lift
   (function (_ state)
     (case (get@ #.scopes state)
       #.Nil
       (#e.Error "Cannot get next reference when there is no scope.")
       
       (#.Cons top _)
       (#e.Success [state (get@ [#.locals #.counter] top)])))))

(def: (ref-to-variable ref)
  (-> Ref Variable)
  (case ref
    (#.Local register)
    (#reference.Local register)
    
    (#.Captured register)
    (#reference.Foreign register)))

(def: #export (environment scope)
  (-> Scope (List Variable))
  (|> scope
      (get@ [#.captured #.mappings])
      (list/map (function (_ [_ [_ ref]]) (ref-to-variable ref)))))