{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# 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
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.Kind (Type)
import Data.List.NonEmpty (NonEmpty (..))
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 :: Type) (contentType :: Type) (a :: Type)
  deriving (Generic, Typeable)

type StreamGet = Stream 'GET 200

type StreamPost = Stream 'POST 200

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

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

-------------------------------------------------------------------------------
-- 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 -> IO a

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

sourceFromSourceIO :: forall m a. MonadIO m => SourceT IO a -> SourceT m a
sourceFromSourceIO src =
  SourceT $ \k ->
    k $ Effect $ liftIO $ unSourceT src (pure . 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)
    void (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))
