{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}

module Language.Yanais.Fusion (
    Step(..),
    stepState,
    runSteps,
    StepFun(..),
    Stream,
    runStream,
    embedTfs,
    step
) where

import Control.Applicative (Alternative(..))
import Prelude hiding (filter)
import Data.Bifunctor
import qualified Data.Text.Internal.Fusion as TF

data Step s a = Done | Skip !s | Yield !a !s

instance Functor (Step s) where
    fmap _ Done = Done
    fmap _ (Skip s) = Skip s
    fmap f (Yield a s) = Yield (f a) s
    {-# INLINE fmap #-}

instance Bifunctor Step where
    bimap _ _ Done = Done
    bimap f _ (Skip s) = Skip (f s)
    bimap f g (Yield a s) = Yield (g a) (f s)
    {-# INLINE bimap #-}

stepState :: Step s a -> Maybe s
stepState Done = Nothing
stepState (Skip st) = Just st
stepState (Yield _ st) = Just st

data StepFun s a = StepFun (s -> Step s a)

instance Functor (StepFun s) where
    fmap f inif = inif >>= (\x -> pure (f x))
    {-# INLINE fmap #-}

instance Applicative (StepFun s) where
    pure x = StepFun $ \st -> Yield x st
    {-# INLINE pure #-}

    f <*> v = f >>= (\fx -> v >>= (\vx -> StepFun $ Yield (fx vx)))

instance Monad (StepFun s) where
    return = pure
    (StepFun f1) >>= f2 = StepFun go
      where
        go st = case f1 st of
          Done -> Done
          Skip st2 -> Skip st2
          Yield x st2 -> let (StepFun f3) = (f2 x) in f3 st2

-- this makes `guard` work
instance Alternative (StepFun s) where
    empty = StepFun $ Skip
    {-# INLINE empty #-}

    StepFun f1 <|> StepFun f2 = StepFun $ \st ->
      case f1 st of
        Yield a st2 -> Yield a st2
        -- backtracking
        Done -> f2 st
        Skip _ -> f2 st

runSteps :: StepFun s a -> s -> [a]
runSteps (StepFun sf) = go
  where
    go st = case sf st of
      Done -> []
      Skip st2 -> go st2
      Yield a st2 -> a:(go st2)

data Stream a = forall s. Stream
    (StepFun s a)   -- ^ stepper function
    !s              -- ^ current state

instance Functor Stream where
    fmap f (Stream step_ cst) = Stream (fmap f step_) cst

runStream :: Stream a -> [a]
runStream (Stream sf ctx) = runSteps sf ctx
{-# INLINE runStream #-}

smap :: TF.Step s a -> Step s a
smap  TF.Done         = Done
smap (TF.Skip st)     = Skip st
smap (TF.Yield av st) = Yield av st

-- | embed a stream from Data.Text into our stream type
embedTfs :: TF.Stream x -> Stream x
embedTfs (TF.Stream step_ st _) = Stream (StepFun (smap . step_)) st

-- | step an inner stream;
-- forwards skips and turns Done into Nothing-yields
step :: StepFun (Stream a) (Maybe a)
step = StepFun $ \(Stream (StepFun f) st) -> case f st of
    Done -> Yield Nothing (Stream (StepFun f) st)
    Skip st2 -> Skip (Stream (StepFun f) st2)
    Yield x st2 -> Yield (Just x) (Stream (StepFun f) st2)