{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DeriveDataTypeable     #-}
{-# LANGUAGE DeriveGeneric          #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# OPTIONS_HADDOCK not-home #-}

module Servant.API.Stream (
    Stream,
    StreamGet,
    StreamPost,
    StreamBody,
    StreamBody',
    -- * Source
    --
    -- | 'SourceIO' are equivalent to some *source* in streaming libraries.
    SourceIO,
    ToSourceIO (..),
    FromSourceIO (..),
    -- ** Auxiliary classes
    SourceToSourceIO (..),
    -- * Framing
    FramingRender (..),
    FramingUnrender (..),
    -- ** Strategies
    NoFraming,
    NewlineFraming,
    NetstringFraming,
    ) where


import           Control.Applicative
                 ((<|>))
import           Control.Monad.IO.Class
                 (MonadIO (..))
import qualified Data.Attoparsec.ByteString       as A
import qualified Data.Attoparsec.ByteString.Char8 as A8
import qualified Data.ByteString                  as BS
import qualified Data.ByteString.Lazy             as LBS
import qualified Data.ByteString.Lazy.Char8       as LBS8
import           Data.List.NonEmpty
                 (NonEmpty (..))
import           Data.Monoid
                 ((<>))
import           Data.Proxy
                 (Proxy)
import           Data.Typeable
                 (Typeable)
import           GHC.Generics
                 (Generic)
import           GHC.TypeLits
                 (Nat)
import           Network.HTTP.Types.Method
                 (StdMethod (..))
import           Servant.Types.SourceT

-- | A Stream endpoint for a given method emits a stream of encoded values at a
-- given @Content-Type@, delimited by a @framing@ strategy.
-- Type synonyms are provided for standard methods.
--
data Stream (method :: k1) (status :: Nat) (framing :: *) (contentType :: *) (a :: *)
  deriving (Typeable, Generic)

type StreamGet  = Stream 'GET 200
type StreamPost = Stream 'POST 200

-- | A stream request body.
type StreamBody = StreamBody' '[]

data StreamBody' (mods :: [*]) (framing :: *) (contentType :: *) (a :: *)
  deriving (Typeable, Generic)

-------------------------------------------------------------------------------
-- Sources
-------------------------------------------------------------------------------

-- | Stream endpoints may be implemented as producing a @'SourceIO' chunk@.
--
-- Clients reading from streaming endpoints can be implemented as consuming a
-- @'SourceIO' chunk@.
--
type SourceIO = SourceT IO

-- | 'ToSourceIO' is intended to be implemented for types such as Conduit, Pipe,
-- etc. By implementing this class, all such streaming abstractions can be used
-- directly as endpoints.
class ToSourceIO chunk a | a -> chunk where
    toSourceIO :: a -> SourceIO chunk

-- | Auxiliary class for @'ToSourceIO' x ('SourceT' m x)@ instance.
class SourceToSourceIO m where
    sourceToSourceIO :: SourceT m a -> SourceT IO a

instance SourceToSourceIO IO where
    sourceToSourceIO = id

-- | Relax to use auxiliary class, have m
instance SourceToSourceIO m => ToSourceIO chunk (SourceT m chunk) where
    toSourceIO = sourceToSourceIO

instance ToSourceIO a (NonEmpty a) where
    toSourceIO (x :| xs) = fromStepT (Yield x (foldr Yield Stop xs))

instance ToSourceIO a [a] where
    toSourceIO = source

-- | 'FromSourceIO' is intended to be implemented for types such as Conduit,
-- Pipe, etc. By implementing this class, all such streaming abstractions can
-- be used directly on the client side for talking to streaming endpoints.
class FromSourceIO chunk a | a -> chunk where
    fromSourceIO :: SourceIO chunk -> a

instance MonadIO m => FromSourceIO a (SourceT m a) where
    fromSourceIO = sourceFromSourceIO

sourceFromSourceIO :: forall m a. MonadIO m => SourceT IO a -> SourceT m a
sourceFromSourceIO src =
    SourceT $ \k ->
    k $ Effect $ liftIO $ unSourceT src (return . go)
  where
    go :: StepT IO a -> StepT m a
    go Stop        = Stop
    go (Error err) = Error err
    go (Skip s)    = Skip (go s)
    go (Effect ms) = Effect (liftIO (fmap go ms))
    go (Yield x s) = Yield x (go s)

-- This fires e.g. in Client.lhs
-- {-# OPTIONS_GHC -ddump-simpl -ddump-rule-firings -ddump-to-file #-}
{-# NOINLINE [2] sourceFromSourceIO #-}
{-# RULES "sourceFromSourceIO @IO" sourceFromSourceIO = id :: SourceT IO a -> SourceT IO a #-}

-------------------------------------------------------------------------------
-- Framing
-------------------------------------------------------------------------------

-- | The 'FramingRender' class provides the logic for emitting a framing strategy.
-- The strategy transforms a @'SourceT' m a@ into @'SourceT' m 'LBS.ByteString'@,
-- therefore it can prepend, append and intercalate /framing/ structure
-- around chunks.
--
-- /Note:/ as the @'Monad' m@ is generic, this is pure transformation.
--
class FramingRender strategy where
    framingRender :: Monad m => Proxy strategy -> (a -> LBS.ByteString) -> SourceT m a -> SourceT m LBS.ByteString

-- | The 'FramingUnrender' class provides the logic for parsing a framing
-- strategy.
class FramingUnrender strategy where
    framingUnrender :: Monad m => Proxy strategy -> (LBS.ByteString -> Either String a) -> SourceT m BS.ByteString -> SourceT m a

-------------------------------------------------------------------------------
-- NoFraming
-------------------------------------------------------------------------------

-- | A framing strategy that does not do any framing at all, it just passes the
-- input data This will be used most of the time with binary data, such as
-- files
data NoFraming

instance FramingRender NoFraming where
    framingRender _ = fmap

-- | As 'NoFraming' doesn't have frame separators, we take the chunks
-- as given and try to convert them one by one.
--
-- That works well when @a@ is a 'ByteString'.
instance FramingUnrender NoFraming where
    framingUnrender _ f = mapStepT go
      where
        go Stop        = Stop
        go (Error err) = Error err
        go (Skip s)    = Skip (go s)
        go (Effect ms) = Effect (fmap go ms)
        go (Yield x s) = case f (LBS.fromStrict x) of
            Right y  -> Yield y (go s)
            Left err -> Error err

-------------------------------------------------------------------------------
-- NewlineFraming
-------------------------------------------------------------------------------

-- | A simple framing strategy that has no header, and inserts a
-- newline character after each frame.  This assumes that it is used with a
-- Content-Type that encodes without newlines (e.g. JSON).
data NewlineFraming

instance FramingRender NewlineFraming where
    framingRender _ f = fmap (\x -> f x <> "\n")

instance FramingUnrender NewlineFraming where
    framingUnrender _ f = transformWithAtto $ do
        bs <- A.takeWhile (/= 10)
        () <$ A.word8 10 <|> A.endOfInput
        either fail pure (f (LBS.fromStrict bs))

-------------------------------------------------------------------------------
-- NetstringFraming
-------------------------------------------------------------------------------

-- | The netstring framing strategy as defined by djb:
-- <http://cr.yp.to/proto/netstrings.txt>
--
-- Any string of 8-bit bytes may be encoded as @[len]":"[string]","@.  Here
-- @[string]@ is the string and @[len]@ is a nonempty sequence of ASCII digits
-- giving the length of @[string]@ in decimal. The ASCII digits are @<30>@ for
-- 0, @<31>@ for 1, and so on up through @<39>@ for 9. Extra zeros at the front
-- of @[len]@ are prohibited: @[len]@ begins with @<30>@ exactly when
-- @[string]@ is empty.
--
-- For example, the string @"hello world!"@ is encoded as
-- @<31 32 3a 68 65 6c 6c 6f 20 77 6f 72 6c 64 21 2c>@,
-- i.e., @"12:hello world!,"@.
-- The empty string is encoded as @"0:,"@.
--
data NetstringFraming

instance FramingRender NetstringFraming where
    framingRender _ f = fmap $ \x ->
        let bs = f x
        in LBS8.pack (show (LBS8.length bs)) <> ":" <> bs <> ","

instance FramingUnrender NetstringFraming where
    framingUnrender _ f = transformWithAtto $ do
        len <- A8.decimal
        _ <- A8.char ':'
        bs <- A.take len
        _ <- A8.char ','
        either fail pure (f (LBS.fromStrict bs))