(;module: {#;doc "Algebraic effects."}
  lux
  (lux (control ["F" functor]
                applicative
                monad)
       [io #- run]
       (data (coll [list "List/" Monad<List> Monoid<List>])
             [number "Nat/" Codec<Text,Nat>]
             text/format
             error
             [ident "Ident/" Eq<Ident>]
             [text])
       [compiler]
       [macro]
       (macro [ast]
              ["s" syntax #+ syntax: Syntax]
              (syntax [common]))
       [type]))

## [Type]
(type: #export (Eff F a)
  {#;doc "A Free Monad implementation for algebraic effects."}
  (#Pure a)
  (#Effect (F (Eff F a))))

(sig: #export (Handler E M)
  {#;doc "A way to interpret effects into arbitrary monads."}
  (: (Monad M)
     monad)
  (: (All [a] (-> (E a) (M a)))
     handle))

## [Values]
(struct: #export (Functor<Eff> dsl)
  (All [F] (-> (F;Functor F) (F;Functor (Eff F))))
  (def: (map f ea)
    (case ea
      (#Pure a)
      (#Pure (f a))
      
      (#Effect value)
      (#Effect (:: dsl map (map f) value)))))

(struct: #export (Applicative<Eff> dsl)
  (All [F] (-> (F;Functor F) (Applicative (Eff F))))
  (def: functor (Functor<Eff> dsl))

  (def: (wrap a)
    (#Pure a))
  
  (def: (apply ef ea)
    (case [ef ea]
      [(#Pure f) (#Pure a)]
      (#Pure (f a))

      [(#Pure f) (#Effect fa)]
      (#Effect (:: dsl map
                   (:: (Functor<Eff> dsl) map f)
                   fa))

      [(#Effect ff) _]
      (#Effect (:: dsl map
                   (lambda [f] (apply f ea))
                   ff))
      )))

(struct: #export (Monad<Eff> dsl)
  (All [F] (-> (F;Functor F) (Monad (Eff F))))
  (def: applicative (Applicative<Eff> dsl))

  (def: (join efefa)
    (case efefa
      (#Pure efa)
      (case efa
        (#Pure a)
        (#Pure a)

        (#Effect fa)
        (#Effect fa))
      
      (#Effect fefa)
      (#Effect (:: dsl map
                   (:: (Monad<Eff> dsl) join)
                   fefa))
      )))

(type: #hidden (|@ L R)
  (All [a] (| (L a) (R a))))

(def: #hidden (combine-functors left right)
  (All [L R]
    (-> (F;Functor L) (F;Functor R)
        (F;Functor (|@ L R))))
  (struct
   (def: (map f l|r)
     (case l|r
       (+0 l) (+0 (:: left map f l))
       (+1 r) (+1 (:: right map f r)))
     )))

(def: #hidden (combine-handlers Monad<M> left right)
  (All [L R M]
    (-> (Monad M)
        (Handler L M) (Handler R M)
        (Handler (|@ L R) M)))
  (struct
   (def: monad Monad<M>)
   
   (def: (handle l|r)
     (case l|r
       (#;Left l) (:: left handle l)
       (#;Right r) (:: right handle r)
       ))))

## [Syntax]
(syntax: #export (|E [effects (s;many s;any)])
  {#;doc (doc "A way to combine smaller effect into a larger effect."
              (type: EffABC (|E EffA EffB EffC)))}
  (wrap (list (` ($_ ;;|@ (~@ effects))))))

(syntax: #export (|F [functors (s;many s;any)])
  {#;doc (doc "A way to combine smaller effect functors into a larger functor."
              (def: Functor<EffABC>
                (Functor EffABC)
                (|F Functor<EffA> Functor<EffB> Functor<EffC>)))}
  (wrap (list (` ($_ ;;combine-functors (~@ functors))))))

(syntax: #export (|H monad [handlers (s;many s;any)])
  {#;doc (doc "A way to combine smaller effect handlers into a larger handler."
              (def: Handler<EffABC,IO>
                (Handler EffABC io;IO)
                (|H io;Monad<IO>
                    Handler<EffA,IO> Handler<EffB,IO> Handler<EffC,IO>)))}
  (do @
    [g!combiner (compiler;gensym "")]
    (wrap (list (` (let [(~ g!combiner) (;;combine-handlers (~ monad))]
                     ($_ (~ g!combiner) (~@ handlers))))))))

(type: Op
  {#name Text
   #inputs (List AST)
   #output AST})

(def: op^
  (Syntax Op)
  (s;form (s;either ($_ s;seq
                        s;local-symbol
                        (s;tuple (s;some s;any))
                        s;any)
                    ($_ s;seq
                        s;local-symbol
                        (:: s;Monad<Syntax> wrap (list))
                        s;any))))

(syntax: #export (effect: [exp-lvl common;export-level]
                   [name s;local-symbol]
                   [ops (s;many op^)])
  {#;doc (doc "Define effects by specifying which operations and constants a handler must provide."
              (effect: #export EffA
                (opA [Nat Text] Bool)
                (fieldA Nat))

              "In this case, 'opA' will be a function (-> Nat Text Bool)."
              "'fieldA' will be a value provided by a handler.")}
  (do @
    [g!output (compiler;gensym "g!output")
     #let [op-types (List/map (lambda [op]
                                (let [g!tag (ast;tag ["" (get@ #name op)])
                                      g!inputs (` [(~@ (get@ #inputs op))])
                                      g!output (` (-> (~ (get@ #output op)) (~ g!output)))]
                                  (` ((~ g!tag) (~ g!inputs) (~ g!output)))))
                              ops)
           type-name (ast;symbol ["" name])
           type-def (` (type: (~@ (common;gen-export-level exp-lvl))
                         ((~ type-name) (~ g!output))
                         (~@ op-types)))
           op-tags (List/map (|>. (get@ #name) [""] ast;tag (list) ast;tuple)
                             ops)
           functor-def (` (struct: (~@ (common;gen-export-level exp-lvl)) (~' _) (F;Functor (~ type-name))
                            (def: ((~' map) (~' f) (~' fa))
                              (case (~' fa)
                                (^template [(~' <tag>)]
                                  ((~' <tag>) (~' params) (~' cont))
                                  ((~' <tag>) (~' params) (. (~' f) (~' cont))))
                                ((~@ op-tags))))
                            ))
           function-defs (List/map (lambda [op]
                                     (let [g!name (ast;symbol ["" (get@ #name op)])
                                           g!tag (ast;tag ["" (get@ #name op)])
                                           g!params (: (List AST)
                                                       (case (list;size (get@ #inputs op))
                                                         +0 (list)
                                                         s (|> (list;n.range +0 (n.dec s))
                                                               (List/map (|>. Nat/encode
                                                                              (format "_")
                                                                              [""]
                                                                              ast;symbol)))))]
                                       (` (def: (~@ (common;gen-export-level exp-lvl)) ((~ g!name) (~@ g!params))
                                            (-> (~@ (get@ #inputs op))
                                                ((~ type-name) (~ (get@ #output op))))
                                            ((~ g!tag) [(~@ g!params)] ;id)))))
                                   ops)]]
    (wrap (list& type-def
                 functor-def
                 function-defs))))

(type: Translation
  {#effect Ident
   #target-type AST
   #target-monad AST})

(def: translation^
  (Syntax Translation)
  (s;form (do s;Monad<Syntax>
            [_ (s;this! (' =>))]
            (s;seq s;symbol
                   (s;tuple (s;seq s;any
                                   s;any))))))

(syntax: #export (handler: [exp-lvl common;export-level]
                   [name s;local-symbol]
                   [[effect target-type target-monad] translation^]
                   [defs (s;many (common;def *compiler*))])
  {#;doc (doc "Define effect handlers by implementing the operations and values of an effect."
              (handler: _
                (=> EffA [IO Monad<IO>])
                (def: (opA length sample)
                  (:: Monad<IO> wrap (n.< length
                                          (size sample))))

                (def: fieldA (:: Monad<IO> wrap +10)))

              "Since a name for the handler was not specified, 'handler:' will generate the name as Handler<EffA,IO>.")}
  (do @
    [(^@ effect [e-module _]) (compiler;un-alias effect)
     g!input (compiler;gensym "g!input")
     g!cont (compiler;gensym "g!cont")
     g!value (compiler;gensym "value")
     g!wrap (compiler;gensym "wrap")
     #let [g!cases (|> defs
                       (List/map (lambda [def]
                                   (let [g!tag (ast;tag [e-module (get@ #common;def-name def)])
                                         g!args (List/map (|>. [""] ast;symbol)
                                                          (get@ #common;def-args def))
                                         eff-calc (case (get@ #common;def-type def)
                                                    #;None
                                                    (get@ #common;def-value def)

                                                    (#;Some type)
                                                    (` (: (~ type) (~ (get@ #common;def-value def)))))
                                         invocation (case g!args
                                                      #;Nil
                                                      eff-calc

                                                      _
                                                      (` ((~ eff-calc) (~@ g!args))))]
                                     (list (` ((~ g!tag) [(~@ g!args)] (~ g!cont)))
                                           (` (do (~ target-monad)
                                                [(~' #let) [(~ g!wrap) (~' wrap)]
                                                 (~ g!value) (~ invocation)]
                                                ((~ g!wrap) ((~ g!cont) (~ g!value)))))
                                           ))))
                       List/join)]]
    (wrap (list (` (struct: (~@ (common;gen-export-level exp-lvl)) (~ (ast;symbol ["" name]))
                     (;;Handler (~ (ast;symbol effect)) (~ target-type))
                     (def: (~' monad) (~ target-monad))
                     
                     (def: ((~' handle) (~ g!input))
                       (case (~ g!input)
                         (~@ g!cases))
                       )))))))

(def: #export (with-handler handler body)
  {#;doc "Handles an effectful computation with the given handler to produce a monadic value."}
  (All [E M a] (-> (Handler E M) (Eff E a) (M a)))
  (case body
    (#Pure value)
    (:: handler wrap value)

    (#Effect effect)
    (do (get@ #monad handler)
      [result (:: handler handle effect)]
      (with-handler handler result))
    ))

(def: (un-apply type-app)
  (-> Type Type)
  (case type-app
    (#;AppT effect value)
    effect

    _
    (error! (format "Wrong type format: " (%type type-app)))))

(def: (clean-effect effect)
  (-> Type Type)
  (case effect
    (#;UnivQ env body)
    (#;UnivQ (list) body)

    _
    (error! (format "Wrong effect format: " (%type effect)))))

(def: g!functor AST (ast;symbol ["" "\t@E\t"]))

(syntax: #export (doE functor [bindings (s;tuple (s;some s;any))] body)
  {#;doc (doc "An alternative to the 'do' macro for monads."
              (with-handler Handler<EffABC,IO>
                (doE Functor<EffABC>
                  [a (lift fieldA)
                   b (lift fieldB)
                   c (lift fieldC)]
                  (wrap ($_ n.+ a b c)))))}
  (do @
    [g!output (compiler;gensym "")]
    (wrap (list (` (let [(~ g!functor) (~ functor)]
                     (do (Monad<Eff> (~ g!functor))
                       [(~@ bindings)
                        (~ g!output) (~ body)]
                       (#;;Pure (~ g!output)))))))))

(def: (flatten-effect-stack stack)
  (-> Type (List Type))
  (case stack
    (#;SumT left right)
    (List/append (flatten-effect-stack left)
                 (flatten-effect-stack right))

    (^ (#;AppT branches (#;VarT _)))
    (flatten-effect-stack branches)

    (^ (#;AppT (#;AppT (#;NamedT (ident-for ;;|@) _)
                       left)
               right))
    (#;Cons left (flatten-effect-stack right))

    (^ (#;AppT (#;AppT (#;NamedT (ident-for ;;Eff) _)
                       effect)
               param))
    (list effect)

    _
    (list stack)
    ))

(def: (same-effect? expected actual)
  (case [expected actual]
    [(#;NamedT e-name _) (#;NamedT a-name _)]
    (Ident/= e-name a-name)

    _
    false))

(def: (nest-effect idx total base)
  (-> Nat Nat AST AST)
  (cond (n.= +0 idx)
        (` (+0 (~ base)))

        (n.> +2 total)
        (` (+1 (~ (nest-effect (n.dec idx) (n.dec total) base))))

        ## else
        (` (+1 (~ base)))
        ))

(syntax: #export (lift [value (s;alt s;symbol
                                     s;any)])
  {#;doc (doc "A way to (automatically) lift effectful fields and operations from simple effects into the larger space of composite effects."
              (with-handler Handler<EffABC,IO>
                (doE Functor<EffABC>
                  [a (lift fieldA)
                   b (lift fieldB)
                   c (lift fieldC)]
                  (wrap ($_ n.+ a b c)))))}
  (case value
    (#;Left var)
    (do @
      [input (compiler;find-type var)
       output compiler;expected-type]
      (case [input output]
        (^=> [(#;AppT eff0 _) (#;AppT stackT0 recT0)]
             [(type;apply-type stackT0 recT0) (#;Some unfoldT0)]
             [stackT0 (^ (#;AppT (#;NamedT (ident-for ;;Eff) _)
                                 stackT1))]
             [(type;apply-type stackT1 recT0) (#;Some unfoldT1)]
             [(flatten-effect-stack unfoldT1) stack]
             [(|> stack list;enumerate
                  (list;find (lambda [[idx effect]]
                               (same-effect? effect eff0))))
              (#;Some [idx _])])
        (wrap (list (` (#;;Effect (:: (~ g!functor) (~' map) (~' wrap)
                                      (~ (nest-effect idx (list;size stack) (ast;symbol var))))))))

        _
        (compiler;fail (format "Invalid type to lift: " (%type output)))))

    (#;Right node)
    (do @
      [g!value (compiler;gensym "")]
      (wrap (list (` (let [(~ g!value) (~ node)]
                       (;;lift (~ g!value)))))))))