{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-error=unused-imports #-}

-- | Server-sent events
--
-- See <https://www.w3.org/TR/2009/WD-eventsource-20090421/> for more details
-- on server-sent events (SSE).
module Servant.Client.Core.ServerSentEvents
  ( EventMessage (..)
  , EventIgnoreReason (..)
  , Event (..)
  , EventStreamT (..)
  , JsonEventStreamT (..)
  , EventMessageStreamT (..)
  )
where

import Control.Applicative (Alternative ((<|>)))
import Control.Monad.IO.Class (MonadIO)
import qualified Data.Aeson as Aeson
import qualified Data.Attoparsec.ByteString as Attoparsec
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as ByteString.Char8
import qualified Data.ByteString.Lazy as ByteString.Lazy
import Data.Char (chr)
import Data.Coerce (coerce)
import Data.Foldable (traverse_)
import Data.Functor (void)
-- For compatibility with GHC <= 8.2
import Data.Semigroup (Semigroup (..))
import qualified Data.Text as Text
import Data.Text.Encoding (encodeUtf8)
import GHC.Generics (Generic)
import Numeric.Natural (Natural)
import Servant.API.ContentTypes (EventStreamChunk (..))
import Servant.API.Stream (FromSourceIO (..))
import Servant.Types.SourceT
  ( SourceT
  , StepT (..)
  , foreachYieldStep
  , mapStepT
  , transformStepWithAtto
  )

-- | Line (or frame) of an event stream
newtype EventStreamLine = EventStreamLine
  {unEventStreamLine :: ByteString.ByteString}
  deriving (Show)

-- | Consume chunks to produce event stream lines.
eventLinesFromRawChunks
  :: Monad m
  => StepT m ByteString.ByteString
  -> StepT m EventStreamLine
eventLinesFromRawChunks =
  transformStepWithAtto eventLine

-- | Consume event stream chunks to produce event stream lines.
eventLinesFromChunks
  :: Monad m
  => StepT m EventStreamChunk
  -> StepT m EventStreamLine
eventLinesFromChunks =
  -- 'coerce' efficiently unpacks the 'EventStreamChunk'
  eventLinesFromRawChunks . fmap (coerce ByteString.Lazy.toStrict)

-- | Apply a 'Attoparsec.Parser' to each line of the event stream individually.
parseEventLines
  :: Monad m
  => Attoparsec.Parser a
  -> StepT m EventStreamLine
  -> StepT m a
parseEventLines parser =
  foreachYieldStep $ \(EventStreamLine line) next ->
    case Attoparsec.parseOnly parser line of
      Left err -> Error err
      Right value -> Yield value next

-- | A line of an event stream
eventLine :: Attoparsec.Parser EventStreamLine
eventLine = do
  Attoparsec.option () byteOrderMark -- A line may be prefixed with a byte order mark
  EventStreamLine <$> untilLineEnd <* lineEnd

-- | Byte order mark (U+FEFF) in UTF-8 representation
byteOrderMark :: Attoparsec.Parser ()
byteOrderMark =
  traverse_ Attoparsec.word8 $
    ByteString.unpack $
      encodeUtf8 $
        Text.singleton $
          chr 0xFEFF

-- | Event stream line ending
lineEnd :: Attoparsec.Parser ()
lineEnd =
  (cr >> lf) <|> cr <|> lf <|> Attoparsec.endOfInput
  where
    cr = void (Attoparsec.word8 0x0D)
    lf = void (Attoparsec.word8 0x0A)

-- | Consume all contents until the end of the line.
untilLineEnd :: Attoparsec.Parser ByteString.ByteString
untilLineEnd = Attoparsec.takeWhile (\w8 -> w8 /= 0x0D && w8 /= 0x0A)

-- | Structured variant of an event line of an event stream
data EventMessage
  = -- | Dispatch on the accumulated event.
    EventDispatch
  | -- | Set the name of the current event.
    EventSetName ByteString.ByteString
  | -- | Set the last event identifier.
    EventSetLastId ByteString.ByteString
  | -- | Append data to the event's data buffer.
    EventData ByteString.ByteString
  | -- | Set the event stream's reconnection time.
    EventRetry Natural
  | -- | Ignored
    EventIgnore EventIgnoreReason
  deriving (Eq, Ord, Show)

-- | Reason why a event line can be ignored
data EventIgnoreReason
  = EventFieldNameUnknown ByteString.ByteString
  | EventRetryNonNumeric ByteString.ByteString
  | EventComment ByteString.ByteString
  deriving (Eq, Ord, Show)

-- | Parse the event stream lines into more structured messages.
eventMessagesFromLines
  :: Monad m
  => StepT m EventStreamLine
  -> StepT m EventMessage
eventMessagesFromLines =
  ensureLastDispatch False . parseEventLines eventMessage
  where
    -- \| Make sure the last event message is a dispatch.
    ensureLastDispatch didDispatch step = case step of
      Stop ->
        if not didDispatch then Yield EventDispatch Stop else Stop
      Yield other next ->
        Yield other $ ensureLastDispatch (other == EventDispatch) next
      Skip next ->
        Skip $ ensureLastDispatch didDispatch next
      Effect eff ->
        Effect $ ensureLastDispatch didDispatch <$> eff
      err@Error{} ->
        err

-- | Event line parser for an event message.
eventMessage :: Attoparsec.Parser EventMessage
eventMessage =
  ignore <|> field <|> dispatch
  where
    ignore = do
      _ <- Attoparsec.word8 0x3A -- ':'
      EventIgnore . EventComment <$> Attoparsec.takeByteString

    dispatch = do
      Attoparsec.endOfInput
      pure EventDispatch

    field = do
      name <- Attoparsec.takeWhile1 (/= 0x3A) -- Up to ':' or the end
      value <- Attoparsec.option ByteString.empty $ do
        _ <- Attoparsec.word8 0x3A -- ':'
        _ <- Attoparsec.option 0x20 $ Attoparsec.word8 0x20 -- Optional ' '
        Attoparsec.takeByteString

      pure $ case name of
        "event" -> EventSetName value
        "data" -> EventData value
        "id" -> EventSetLastId value
        "retry" ->
          -- The retry value consist of digits.
          if ByteString.all (\w8 -> 0x30 <= w8 && w8 <= 0x39) value
            then
              EventRetry (read (ByteString.Char8.unpack value))
            else
              EventIgnore (EventRetryNonNumeric value)
        _ -> EventIgnore (EventFieldNameUnknown name)

-- | Event sent by the remote
data Event a = Event
  { eventName :: Maybe ByteString.ByteString
  , eventData :: a
  }
  deriving (Eq, Functor, Generic, Ord, Show)

-- | Accumulate event messages to build individual 'Event's.
eventsFromMessages
  :: Functor m
  => StepT m EventMessage
  -> StepT m (Event ByteString.ByteString)
eventsFromMessages =
  initGo
  where
    initGo = go Nothing ByteString.Lazy.empty

    combineData dataBuffer newData =
      if ByteString.Lazy.null dataBuffer
        then
          ByteString.Lazy.fromStrict newData
        else
          ByteString.Lazy.concat
            [ dataBuffer
            , ByteString.Lazy.singleton 0x0A -- Line feed
            , ByteString.Lazy.fromStrict newData
            ]

    go name dataBuffer step = case step of
      Stop ->
        Stop
      Skip next ->
        go name dataBuffer next
      Effect eff ->
        Effect (go name dataBuffer <$> eff)
      Error err ->
        Error err
      Yield message next -> case message of
        EventSetName newName ->
          go (Just newName) dataBuffer next
        EventData newData ->
          go name (combineData dataBuffer newData) next
        EventDispatch ->
          Yield
            (Event name (ByteString.Lazy.toStrict dataBuffer))
            (initGo next)
        _ ->
          -- We ignore other message because they don't fit into
          -- the 'Event' type. If a user needs more fine grained
          -- control, the 'EventMessage' interface is better suited.
          go name dataBuffer next

-- | Server-sent event stream (SSE)
--
-- See <https://www.w3.org/TR/2009/WD-eventsource-20090421/> for more details.
newtype EventMessageStreamT m = EventMessageStreamT
  {unEventMessageStreamT :: SourceT m EventMessage}
  deriving stock (Show)
  deriving newtype (Monoid, Semigroup)

-- | Server-sent event messages
--
-- 'EventMessage' gives you more control over the communication with the server
-- than 'Event'.
instance MonadIO m => FromSourceIO EventStreamChunk (EventMessageStreamT m) where
  fromSourceIO src =
    EventMessageStreamT
      . mapStepT (eventMessagesFromLines . eventLinesFromChunks)
      <$> fromSourceIO src

-- | Server-sent event stream (SSE)
--
-- See <https://www.w3.org/TR/2009/WD-eventsource-20090421/> for more details.
newtype EventStreamT m = EventStreamT
  {unEventStreamT :: SourceT m (Event ByteString.ByteString)}
  deriving stock (Show)
  deriving newtype (Monoid, Semigroup)

-- | Server-sent events
instance MonadIO m => FromSourceIO EventStreamChunk (EventStreamT m) where
  fromSourceIO input = do
    src :: EventMessageStreamT m <- fromSourceIO input
    pure $
      -- 'coerce' is used in place of unpacking and repacking 'EventStreamT'
      coerce (mapStepT eventsFromMessages) src

-- | Try to parse event data to JSON.
jsonEventsFromEvents
  :: (Aeson.FromJSON a, Functor m)
  => StepT m (Event ByteString.ByteString)
  -> StepT m (Event a)
jsonEventsFromEvents =
  foreachYieldStep $ \(Event name datas) next ->
    either
      Error
      (\value -> Yield (Event name value) next)
      (Aeson.eitherDecode (ByteString.Lazy.fromStrict datas))

-- | Server-sent event stream (SSE) for JSON values
newtype JsonEventStreamT m a = JsonEventStreamT
  {unJsonEventStreamT :: SourceT m (Event a)}
  deriving stock (Functor, Show)
  deriving newtype (Monoid, Semigroup)

-- | Server-sent JSON event stream
instance (Aeson.FromJSON a, MonadIO m) => FromSourceIO EventStreamChunk (JsonEventStreamT m a) where
  fromSourceIO input = do
    src :: EventStreamT m <- fromSourceIO input
    pure $
      -- The 'coerce' efficiently unwraps the 'EventStreamT' and wraps the
      -- JsonEventStreamT.
      coerce
        (mapStepT jsonEventsFromEvents)
        src
