module Control.Monad.Logic.Trans (SFKT(..), FK, SK, unSFKT, runSFKT, runSFKTOnce) where

import Prelude (class Applicative, class Apply, class Bind, class Functor, class Monad, bind, pure, ($), (<<<), (>>=), Unit, unit)
import Control.Monad.Logic.Class

import Control.Lazy (class Lazy)
import Control.MonadPlus (class MonadPlus, class Alternative, class Alt, class Plus, (<|>), empty)
import Control.Monad.Reader.Class (class MonadReader, local, class MonadAsk, ask)
import Control.Monad.State.Class (class MonadState, state)
import Control.Monad.Trans.Class (class MonadTrans, lift)
import Control.Monad.Rec.Class (class MonadRec, Step(..))
import Data.Functor (map)
import Data.Tuple.Nested ((/\))
import Data.List (List(Nil), (:))
import Data.Maybe (Maybe(..))

type FK :: Type -> Type
type FK ans = Unit -> ans

type SK :: Type -> Type -> Type
type SK ans a = { val :: a, fk :: FK ans } -> ans

newtype SFKT :: (Type -> Type) -> Type -> Type
newtype SFKT m a = SFKT (forall ans. SK (m ans) a -> FK (m ans) -> m ans)

unSFKT :: forall a m. SFKT m a -> forall ans. SK (m ans) a -> FK (m ans) -> m ans
unSFKT (SFKT f) = f

runSFKTOnce :: forall a m. Alternative m => SFKT m a -> m a
runSFKTOnce (SFKT f) = f (pure <<< _.val) (\_ -> empty)

runSFKT :: forall a m. Monad m => SFKT m a -> m (List a)
runSFKT (SFKT f) = f (\{val: a, fk} -> map (a:_) (fk unit)) (\_ -> pure Nil)

instance Functor (SFKT m) where
    map f m = SFKT (\sk -> unSFKT m (\{val: a, fk} -> sk { val: f a, fk: fk }))

instance Apply m => Apply (SFKT m) where
    apply mf ma = SFKT (\sk -> unSFKT mf (\{val: f, fk: fk} -> unSFKT ma (\{val: a, fk: fk'} -> sk { val: f a, fk: fk' }) fk))

instance Applicative m => Applicative (SFKT m) where
    pure a = SFKT (\sk fk -> sk {val: a, fk: fk})

instance Bind m => Bind (SFKT m) where
    bind m f = SFKT (\sk -> unSFKT m (\{val: a, fk} -> unSFKT (f a) sk fk))

instance Monad m => Monad (SFKT m)

instance Alt (SFKT m) where
    alt m1 m2 = SFKT (\sk fk -> unSFKT m1 sk (\_ -> unSFKT m2 sk fk))

instance Plus (SFKT m) where
    empty = SFKT (\_ fk -> fk unit)

instance Applicative m => Alternative (SFKT m)

instance Monad m => MonadPlus (SFKT m)

instance MonadTrans SFKT where
    lift m = SFKT (\sk fk -> m >>= (\a -> sk { val: a, fk: fk }))

instance Monad m => MonadLogic (SFKT m) where
    msplit ma = lift (unSFKT ma (\{val: a, fk} -> pure (Just (a /\ (lazyLift fk >>= reflect)))) (\_ -> pure Nothing))
        where lazyLift f = SFKT (\sk fk -> f unit >>= (\a -> sk {val: a, fk: fk }))
    interleave m1 m2 = do
        r <- msplit m1
        case r of
            Nothing -> m2
            Just (a /\ m1') -> pure a <|> interleave m2 m1'

instance MonadAsk r m => MonadAsk r (SFKT m) where
    ask = lift ask

instance MonadReader r m => MonadReader r (SFKT m) where
    local f m = SFKT (\sk -> unSFKT m (local f <<< sk))

instance MonadState s m => MonadState s (SFKT m) where
    state f = lift $ state f

instance Lazy (SFKT m a) where
    defer f = SFKT (\sk fk -> unSFKT (f unit) sk fk)

instance Monad m => MonadRec (SFKT m) where
    tailRecM :: forall acc b. (acc -> SFKT m (Step acc b)) -> acc -> SFKT m b
    tailRecM f a0 = SFKT \sk fk -> go {sk: sk, fk: fk} a0
        where
            go :: forall ans. { sk :: SK (m ans) b, fk :: FK (m ans) } -> acc -> m ans
            go r@{sk,fk} an =
                let
                    (SFKT sfktf) = f an
                    success {val: st, fk: fk'} =
                        case st of
                            Loop an' -> go r an'
                            Done b -> sk {val: b, fk: fk' }
                 in sfktf success fk