-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

{-# LANGUAGE UndecidableInstances #-}

{-|
Description: Support for partial and conditional http requests (Range, ETag, If-Match, ...)
-}

module Utils.HttpConditional
  ( ByteRangesSpecifier(..), ByteRangeSpecification(..)
  , ByteContentRangeSpecification(..), ByteRangeResponseSpecification(..)
  , IsRangeUnit(..)
  , ETag(..)
  , RepresentationConditionalInformation(..)
  , mkResponseConditional
  , respondSourceConditional, respondSourceDBConditional
  ) where

import ClassyPrelude hiding (Builder)
import Yesod.Core
import Yesod.Persist.Core
import Data.Conduit
import qualified Data.Conduit.Combinators as C
import Data.Binary.Builder (Builder)

import Web.HttpApiData

import qualified Data.Attoparsec.Text as A

import Data.Char (chr, ord)

import qualified Data.Set as Set
import Data.CaseInsensitive (CI)
import qualified Data.CaseInsensitive as CI

import Data.Time

import Network.HTTP.Types
import Network.HTTP.Types.Header

import Control.Lens
import Control.Lens.Extras

import Data.Kind (Type)
import Data.Coerce
import Data.Proxy

import Control.Monad.Trans.Maybe (MaybeT(..))
import Control.Monad.Fail (MonadFail(..))
import Control.Monad.Trans.Resource (ResourceT)

import Network.Wai

import Control.Monad.Random.Class

import qualified Data.ByteString.Base64.URL as Base64
import qualified Data.ByteString as BS

import Data.List.NonEmpty (NonEmpty(..))


ows :: A.Parser ()
ows = A.skipMany $ A.satisfy (`elem` [chr 0x20, chr 0x09])

httpList :: A.Parser a -> A.Parser [a]
httpList itemParser = do
  let sep = A.many1 $ ows *> A.char ',' <* ows
  A.skipMany sep
  xs <- itemParser `A.sepBy1` sep
  A.skipMany sep
  return xs

parseUrlPiece' :: A.Parser a -> (Text -> Either Text a)
parseUrlPiece' p = first pack . A.parseOnly (p <* A.endOfInput)


newtype ByteRangesSpecifier = ByteRangesSpecifier (NonNull (Set ByteRangeSpecification))
  deriving (Eq, Ord, Read, Show, Generic)
data ByteRangeSpecification
  = ByteRangeSpecification
      { byteRangeSpecFirstPosition :: Word64
      , byteRangeSpecLastPosition :: Maybe Word64
      }
  | ByteRangeSuffixSpecification
      { byteRangeSpecSuffixLength :: Word64
      }
  deriving (Eq, Ord, Read, Show, Generic)

instance FromHttpApiData ByteRangesSpecifier where
  parseUrlPiece = parseUrlPiece' parser
    where parser :: A.Parser ByteRangesSpecifier
          parser = do
            ranges <- httpList brSpecParser
            ByteRangesSpecifier <$> maybe (fail "Parser definition error: empty list of ByteRangeSpecifications") return (fromNullable $ Set.fromList ranges)
          brSpecParser :: A.Parser ByteRangeSpecification
          brSpecParser = brSpecParser' <|> brSuffixParser
            where brSpecParser' = do
                    byteRangeSpecFirstPosition <- A.decimal
                    void $ A.char '-'
                    byteRangeSpecLastPosition <- optional A.decimal
                    return ByteRangeSpecification{..}
                  brSuffixParser = do
                    void $ A.char '-'
                    byteRangeSpecSuffixLength <- A.decimal
                    return ByteRangeSuffixSpecification{..}

data ByteContentRangeSpecification
  = ByteContentRangeSpecification
      { byteRangeResponse :: Maybe ByteRangeResponseSpecification
      , byteRangeInstanceLength :: Maybe Word64
      } deriving (Eq, Ord, Read, Show, Generic)
data ByteRangeResponseSpecification
  = ByteRangeResponseSpecification
      { byteRangeResponseSpecFirstPosition
      , byteRangeResponseSpecLastPosition :: Word64
      } deriving (Eq, Ord, Read, Show, Generic)

instance ToHttpApiData ByteContentRangeSpecification where
  toUrlPiece ByteContentRangeSpecification{..} = maybe "*" encByteRangeResponse byteRangeResponse <> "/" <> maybe "*" encByteRangeInstanceLength byteRangeInstanceLength
    where
      encByteRangeInstanceLength = toUrlPiece
      encByteRangeResponse ByteRangeResponseSpecification{..} = toUrlPiece byteRangeResponseSpecFirstPosition <> "-" <> toUrlPiece byteRangeResponseSpecLastPosition


class (FromHttpApiData req, ToHttpApiData resp, Ord (SingularRangeSpecification req), Show resp) => IsRangeUnit req resp | req -> resp, resp -> req where
  type SingularRangeSpecification req :: Type
  rangeUnit :: forall p1 p2. p1 req -> p2 resp -> Text
  rangeRequestAll :: forall p. p req -> SingularRangeSpecification req
  _RangeSpecifications :: Iso' req (NonNull (Set (SingularRangeSpecification req)))
  default _RangeSpecifications :: Coercible req (NonNull (Set (SingularRangeSpecification req)))
                               => Iso' req (NonNull (Set (SingularRangeSpecification req)))
  _RangeSpecifications = coerced
  rangeInstanceLength :: resp -> Maybe Word64
  rangeInstanceLength _ = Nothing

instance IsRangeUnit ByteRangesSpecifier ByteContentRangeSpecification where
  type SingularRangeSpecification ByteRangesSpecifier = ByteRangeSpecification
  rangeUnit _ _ = "bytes"
  rangeRequestAll _ = ByteRangeSpecification 0 Nothing
  rangeInstanceLength = byteRangeInstanceLength


data ETag = WeakETag { unETag :: Text } | StrongETag { unETag :: Text }
  deriving (Eq, Ord, Read, Show, Generic)

parseETag :: A.Parser ETag
parseETag = do
  isWeak <- is _Just <$> optional (A.string "W/")
  void $ A.char '"'
  tag <- pack <$> many (A.satisfy isETagChar)
  void $ A.char '"'
  return $ bool StrongETag WeakETag isWeak tag
  where
    isETagChar c = c == '!'
                || (0x23 <= ord c && ord c <= 0x7e)
                || (0x80 <= ord c && ord c <= 0xff)

instance FromHttpApiData ETag where
  parseUrlPiece = parseUrlPiece' parseETag
instance ToHttpApiData ETag where
  toUrlPiece (WeakETag t) = "W/\"" <> t <> "\"" 
  toUrlPiece (StrongETag t) = "\"" <> t <> "\"" 

strongETagEq, weakETagEq :: ETag -> ETag -> Bool
strongETagEq (StrongETag a) (StrongETag b) = a == b
strongETagEq _ _ = False
weakETagEq = (==) `on` unETag
  
data RepresentationConditionalInformation = RepresentationConditionalInformation
  { representationETag :: Maybe ETag
  , representationLastModified :: Maybe UTCTime
  , representationExists :: Bool
  , requestedActionAlreadySucceeded :: Maybe Status
  } deriving (Eq, Ord, Show, Generic)

newtype ETagMatch = ETagMatch (Set ETag)
  deriving (Eq, Ord, Read, Show, Generic)
  deriving newtype (Semigroup, Monoid)

instance FromHttpApiData ETagMatch where
  parseUrlPiece = parseUrlPiece' parseIfMatch
    where parseIfMatch :: A.Parser ETagMatch
          parseIfMatch = parseEmptyIfMatch <|> parseNonEmptyIfMatch
          parseEmptyIfMatch = mempty <* A.char '*'
          parseNonEmptyIfMatch = ETagMatch . Set.fromList <$> httpList parseETag

parseHTTPTime :: A.Parser UTCTime
parseHTTPTime = do
  inpT <- A.takeText
  maybe (fail "Could not parse time specification") return . parseTimeM True defaultTimeLocale "%a, %d %b %Y %X %Z" $ unpack inpT

newtype ModifiedMatch = ModifiedMatch UTCTime
  deriving (Eq, Ord, Read, Show, Generic)

instance FromHttpApiData ModifiedMatch where
  parseUrlPiece = parseUrlPiece' $ ModifiedMatch <$> parseHTTPTime

data IfRange = IfRangeETag ETag | IfRangeModified UTCTime
  deriving (Eq, Ord, Read, Show, Generic)

instance FromHttpApiData IfRange where
  parseUrlPiece = parseUrlPiece' parseIfRange
    where parseIfRange = parseIfRangeETag <|> parseIfRangeModified
          parseIfRangeETag = IfRangeETag <$> parseETag
          parseIfRangeModified = IfRangeModified <$> parseHTTPTime

newtype RangeRequest req = RangeRequest { unRangeRequest :: req }
  deriving (Generic)
  deriving newtype (Eq, Ord, Read, Show)

instance IsRangeUnit req resp => FromHttpApiData (RangeRequest req) where
  parseUrlPiece = parseUrlPiece' parseRangeRequest
    where parseRangeRequest :: A.Parser (RangeRequest req)
          parseRangeRequest = do
            void . A.string $ rangeUnit (Proxy @req) (Proxy @resp)
            void $ A.char '='
            t <- A.takeText
            either (fail . unpack) (return . RangeRequest) $ parseUrlPiece t

newtype RangeResponse resp = RangeResponse resp
  deriving (Generic)
  deriving newtype (Eq, Ord, Read, Show)

instance IsRangeUnit req resp => ToHttpApiData (RangeResponse resp) where
  toUrlPiece (RangeResponse r) = rangeUnit (Proxy @req) (Proxy @resp) <> " " <> toUrlPiece r

newtype MultipartBoundary = MultipartBoundary ByteString
  deriving (Eq, Ord, Read, Show, Generic)

instance ToHttpApiData MultipartBoundary where
  toUrlPiece (MultipartBoundary bs) = decodeUtf8 $ Base64.encodeUnpadded bs
  
mkResponseConditional :: forall rangeReq rangeResp builder m m'.
                         ( MonadHandler m, Monad m'
                         , IsRangeUnit rangeReq rangeResp
                         , ToFlushBuilder builder
                         )
                      => RepresentationConditionalInformation
                      -> ContentType
                      -> Either (ConduitT () builder m' ()) (SingularRangeSpecification rangeReq -> (ConduitT () builder m' (), rangeResp))
                      -> m (Status, ContentType, ConduitT () (Flush Builder) m' ())
-- ^ Implementes https://tools.ietf.org/html/rfc7232#section-6
--
-- Assumes we are the origin server
mkResponseConditional RepresentationConditionalInformation{..} cType cont = liftHandler $ do
  isSafeMethod <- (`elem` safeMethods) . requestMethod <$> waiRequest

  for_ representationETag $ \etag ->
    replaceOrAddHeader (decodeUtf8 $ CI.original hETag) . decodeUtf8 $ toHeader etag
  for_ representationLastModified $ \lModified ->
    replaceOrAddHeader (decodeUtf8 $ CI.original hLastModified) $ formatRFC1123 lModified
  
  ifMatch <- lookupHeader' hIfMatch
  for_ ifMatch $ \(ETagMatch ps) -> if
    | null ps, representationExists -> return ()
    | Just etag <- representationETag
    , any (`strongETagEq` etag) ps -> return ()
    | Just retCode <- requestedActionAlreadySucceeded -> sendResponseStatus retCode ()
    | otherwise -> preconditionFailed

  ifUnmodifiedSince <- lookupHeader' hIfUnmodifiedSince
  for_ (guard (is _Nothing ifMatch) *> ifUnmodifiedSince) $ \(ModifiedMatch ts) -> if
    | Just lModified <- representationLastModified
    , lModified < addUTCTime (-precision) ts -> return ()
    | Just retCode <- requestedActionAlreadySucceeded -> sendResponseStatus retCode ()
    | otherwise -> preconditionFailed
    
  ifNoneMatch <- lookupHeader' hIfNoneMatch
  for_ ifNoneMatch $ \(ETagMatch ps) -> if
    | null ps, representationExists -> bool preconditionFailed notModified isSafeMethod
    | Just etag <- representationETag
    , any (`weakETagEq` etag) ps -> bool preconditionFailed notModified isSafeMethod
    | otherwise -> return ()

  ifModifiedSince <- lookupHeader' hIfModifiedSince
  for_ (guard (isSafeMethod && is _Nothing ifNoneMatch) *> ifModifiedSince) $ \(ModifiedMatch ts) -> if
    | Just lModified <- representationLastModified
    , lModified <= addUTCTime precision ts -> notModified
    | otherwise -> return ()

  case cont of
    Left evalNoRanges -> do
      replaceOrAddHeader (decodeUtf8 $ CI.original hAcceptRanges) "none"
      return (ok200, cType, evalNoRanges .| C.map toFlushBuilder)
    Right evalRange -> do
      replaceOrAddHeader (decodeUtf8 $ CI.original hAcceptRanges) $ rangeUnit (Proxy @rangeReq) (Proxy @rangeResp)

      mRanges <- do
        ifRange <- lookupHeader' hIfRange
        range <- lookupHeader' @(RangeRequest rangeReq) hRange
        case ifRange of
          Just (IfRangeETag p)
            | Just etag <- representationETag
            , p `strongETagEq` etag -> return range
          Just (IfRangeModified ts)
            | Just lModified <- representationLastModified
            , lModified < addUTCTime (-precision) ts -> return range
          Just _ -> return Nothing
          Nothing -> return range

      let ranges = maybe (rangeRequestAll (Proxy @rangeReq) :| []) (toNonEmpty . view _RangeSpecifications . unRangeRequest) mRanges

      when (length ranges > 5) $ do
        invalidArgs ["Too many ranges"]

      case ranges of
        r :| [] -> do
          let (respSrc, rResp) = evalRange r
          when (is _Just mRanges) $
            replaceOrAddHeader (decodeUtf8 $ CI.original hContentRange) . decodeUtf8 . toHeader $ RangeResponse rResp
          return (bool partialContent206 ok200 $ r == rangeRequestAll (Proxy @rangeReq), cType, respSrc .| C.map toFlushBuilder)
        (toList -> rs) -> do
          boundary <- liftIO $ MultipartBoundary . BS.pack <$> replicateM 12 getRandom
          let cType' = "multipart/byteranges; boundary=" <> toHeader boundary
              bodySrc = do
                forM_ rs $ \r -> do
                  let (respSrc, rResp) = evalRange r
                  sendChunkBS $ "--" <> toHeader boundary <> "\r\n"
                  sendChunkBS $ CI.original hContentType <> ": " <> cType <> "\r\n"
                  sendChunkBS $ CI.original hContentRange <> ": " <> toHeader (RangeResponse rResp) <> "\r\n"
                  sendChunkBS   "\r\n"
                  respSrc .| C.map toFlushBuilder
                  sendChunkBS   "\r\n"
                  sendFlush
                sendChunkBS $ "--" <> toHeader boundary <> "--\r\n"
          return (partialContent206, cType', bodySrc)
  
  where
    lookupHeader' :: forall hdr n. (MonadHandler n, FromHttpApiData hdr) => CI ByteString -> n (Maybe hdr)
    lookupHeader' hdrName = liftHandler . runMaybeT $ do
      hdrBS <- MaybeT $ lookupHeader hdrName
      case parseHeader hdrBS of
        Left errMsg -> do
          $logInfoS "lookupHeader'" $ "Could not parse value for request header “" <> decodeUtf8 (CI.original hdrName) <> "”, “" <> tshow hdrBS <> "”: " <> errMsg
          mzero
        Right val -> return val

    precision :: NominalDiffTime
    precision = 1

    safeMethods = [ methodGet, methodHead, methodOptions ]

    preconditionFailed = sendResponseStatus preconditionFailed412 ()

respondSourceConditional :: forall rangeReq rangeResp builder m a.
                            ( MonadHandler m
                            , IsRangeUnit rangeReq rangeResp
                            , ToFlushBuilder builder
                            )
                         => RepresentationConditionalInformation
                         -> ContentType
                         -> Either (ConduitT () builder (HandlerFor (HandlerSite m)) ()) (SingularRangeSpecification rangeReq -> (ConduitT () builder (HandlerFor (HandlerSite m)) (), rangeResp))
                         -> m a
respondSourceConditional cInfo cType cont = liftHandler $ do
  (rStatus, cType', cont') <- mkResponseConditional cInfo cType cont
  UnliftIO{..} <- askUnliftIO
  sendResponseStatus rStatus ( cType'
                             , toContent $
                                 transPipe (lift @ResourceT . unliftIO) cont'
                             )
  
respondSourceDBConditional :: forall rangeReq rangeResp builder m a.
                              ( MonadHandler m, YesodPersistRunner (HandlerSite m)
                              , IsRangeUnit rangeReq rangeResp
                              , ToFlushBuilder builder
                              )
                           => RepresentationConditionalInformation
                           -> ContentType
                           -> Either (ConduitT () builder (YesodDB (HandlerSite m)) ()) (SingularRangeSpecification rangeReq -> (ConduitT () builder (YesodDB (HandlerSite m)) (), rangeResp))
                           -> m a
respondSourceDBConditional cInfo cType cont = liftHandler $ do
  (rStatus, cType', cont') <- mkResponseConditional cInfo cType cont
  UnliftIO{..} <- askUnliftIO
  sendResponseStatus rStatus ( cType'
                             , toContent . transPipe (lift @ResourceT . unliftIO) $ runDBSource cont'
                             )