{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | This module exports 'ToSourceIO' and 'FromSourceIO' for 'Proxy' and 'SafeT' instances.
module Servant.Pipes
  ( PipesToSourceIO (..)
  )
where

import Control.Monad.IO.Class (MonadIO (..))
import Control.Monad.Trans.Control (liftBaseWith)
import Pipes (ListT (..))
import Pipes.Internal (Proxy (..), X, closed)
import Pipes.Safe (SafeT, runSafeT)
import Servant.API.Stream
import qualified Servant.Types.SourceT as S

-- | Helper class to implement @'ToSourceIO' 'Proxy'@ instance
-- for various monads.
class PipesToSourceIO m where
  pipesToSourceIO :: Proxy X () () b m () -> SourceIO b

instance PipesToSourceIO IO where
  pipesToSourceIO ma = S.SourceT ($ go ma)
    where
      go :: Proxy X () () b IO () -> S.StepT IO b
      go (Pure ()) = S.Stop
      go (M p) = S.Effect (fmap go p)
      go (Request v _) = closed v
      go (Respond b n) = S.Yield b (go (n ()))

instance m ~ IO => PipesToSourceIO (SafeT m) where
  pipesToSourceIO ma =
    S.SourceT $ \k ->
      runSafeT $ liftBaseWith $ \runSafe ->
        k (go runSafe ma)
    where
      go
        :: (forall x. SafeT m x -> m x)
        -> Proxy X () () b (SafeT m) ()
        -> S.StepT IO b
      go _ (Pure ()) = S.Stop
      go runSafe (M p) = S.Effect $ runSafe $ fmap (go runSafe) p
      go _ (Request v _) = closed v
      go runSafe (Respond b n) = S.Yield b (go runSafe (n ()))

instance
  (PipesToSourceIO m, a ~ (), a' ~ X, b' ~ (), r ~ ())
  => ToSourceIO b (Proxy a' a b' b m r)
  where
  toSourceIO = pipesToSourceIO

instance PipesToSourceIO m => ToSourceIO a (ListT m a) where
  toSourceIO = pipesToSourceIO . enumerate

instance
  (MonadIO m, a ~ (), a' ~ X, b' ~ (), r ~ ())
  => FromSourceIO b (Proxy a' a b' b m r)
  where
  fromSourceIO src = pure $ M $ liftIO $ S.unSourceT src (pure . go)
    where
      go :: S.StepT IO b -> Proxy X () () b m ()
      go S.Stop = Pure ()
      go (S.Error err) = M (liftIO (fail err))
      go (S.Skip s) = go s -- drives
      go (S.Effect ms) = M (liftIO (fmap go ms))
      go (S.Yield x s) = Respond x (const (go s))
  {-# SPECIALIZE INLINE fromSourceIO :: SourceIO x -> IO (Proxy X () () x IO ()) #-}

instance MonadIO m => FromSourceIO a (ListT m a) where
  fromSourceIO src = Select <$> liftIO (fromSourceIO src)
