{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}

module Antiope.S3
( s3ObjectSource
, putFile, putFile', putContent , putContent'
, copySingle
, fromS3Uri
, toS3Uri
, lsBucketResponseStream
, lsBucketStream
, lsPrefix
, deleteFiles
, deleteFilesExcept
, fileExists
, Region(..)
, BucketName(..)
, ObjectKey(..)
, ETag(..)
, S3Uri(..)
) where

import Antiope.Core.Error           (handle404ToNone)
import Antiope.S3.Internal
import Antiope.S3.Types             (S3Uri (S3Uri, objectKey))
import Conduit
import Control.Lens
import Control.Monad
import Control.Monad.Trans.AWS      hiding (send)
import Control.Monad.Trans.Resource
import Data.Conduit.List            (unfoldM)
import Data.Maybe                   (catMaybes, isJust)
import Data.Monoid                  ((<>))
import Data.Text                    as T (Text, pack, unpack)
import Network.AWS                  (MonadAWS)
import Network.AWS.Data.Body        (_streamBody)
import Network.AWS.Data.Text        (toText)
import Network.AWS.S3
import Network.URI                  (URI (..), URIAuth (..), parseURI, unEscapeString)

import qualified Data.ByteString      as BS
import qualified Data.ByteString.Lazy as LBS
import qualified Data.List            as List
import qualified Network.AWS          as AWS

chunkSize :: ChunkSize
chunkSize = ChunkSize (1024 * 1024)

type Prefix = Text

fromS3Uri :: Text -> Maybe S3Uri
fromS3Uri uri = do
  puri <- parseURI (unpack uri)
  auth <- puri & uriAuthority
  let b = pack $ auth & uriRegName       -- URI lib is pretty weird
  let k = pack $ unEscapeString $ drop 1 $ puri & uriPath
  pure $ S3Uri (BucketName b) (ObjectKey k)

s3ObjectSource :: (MonadAWS m, MonadResource m)
  => BucketName
  -> ObjectKey
  -> m (ConduitT () BS.ByteString m ())
s3ObjectSource bkt obj = do
  resp <- AWS.send $ getObject bkt obj
  return $ transPipe liftResourceT $ _streamBody $ resp ^. gorsBody

-- | Puts file into a specified S3 bucket
putFile :: MonadAWS m
  => BucketName       -- ^ Target bucket
  -> ObjectKey        -- ^ File name on S3
  -> FilePath         -- ^ Source file path
  -> m (Maybe ETag)   -- ^ Etag when the operation is successful
putFile b k f = do
  req <- chunkedFile chunkSize f
  view porsETag <$> AWS.send (putObject b k req)

putFile' :: MonadAWS m
  => S3Uri            -- ^ S3 URI
  -> FilePath         -- ^ Source file path
  -> m (Maybe ETag)   -- ^ Etag when the operation is successful
putFile' (S3Uri b k) f = do
  req <- chunkedFile chunkSize f
  view porsETag <$> AWS.send (putObject b k req)

putContent :: MonadAWS m
  => BucketName
  -> ObjectKey
  -> LBS.ByteString
  -> m (Maybe ETag)
putContent b k c = view porsETag <$> AWS.send (putObject b k (toBody c))

putContent' :: MonadAWS m
  => S3Uri
  -> LBS.ByteString
  -> m (Maybe ETag)
putContent' (S3Uri b k) = putContent b k

-- | Copies a single object within S3
copySingle :: MonadAWS m
  => BucketName          -- ^ Source bucket name
  -> ObjectKey           -- ^ Source key
  -> BucketName          -- ^ Target bucket name
  -> ObjectKey           -- ^ Target key
  -> m ()
copySingle sb sk tb tk = void . AWS.send $ copyObject tb (toText sb <> "/" <> toText sk) tk
     & coMetadataDirective ?~ MDCopy

-- Private --

-- Builds the request for the next page of a NextObjectsV2 request,
-- based on the original request and the most recent response.
nextPageReq :: ListObjectsV2 -> ListObjectsV2Response -> ListObjectsV2
nextPageReq initial resp =
  initial & lovContinuationToken .~ resp ^. lovrsNextContinuationToken

-- The type signature is like this so that it can be used with `unfoldM`
lsBucketPage :: MonadAWS m
  => Maybe ListObjectsV2
  -> m (Maybe (ListObjectsV2Response, Maybe ListObjectsV2))
lsBucketPage Nothing    = pure Nothing
lsBucketPage (Just req) = do
  resp <- AWS.send req
  pure . Just . (resp, ) $
    case resp ^. lovrsIsTruncated of
      Just True -> Just $ nextPageReq req resp
      _         -> Nothing

-- | Streams all pages of the result (ListObjectsV2Responses) of a ListObjectsV2
-- request from S3.
-- lsBucketResponseStream :: MonadAWS m => ListObjectsV2 -> ConduitT i ListObjectsV2Response m ()
lsBucketResponseStream :: MonadAWS m
  => ListObjectsV2
  -> ConduitM a ListObjectsV2Response m ()
lsBucketResponseStream bar = unfoldM lsBucketPage (Just bar)

-- | Streams all Objects from all pages of the result of a ListObjectsV2
-- request from S3.
-- lsBucketStream :: MonadAWS m => ListObjectsV2 -> ConduitT i Object m ()
lsBucketStream :: MonadAWS m
  => ListObjectsV2
  -> ConduitM a Object m ()
lsBucketStream bar = lsBucketResponseStream bar .| concatMapC (^. lovrsContents)

-- | Lists the specified prefix in a bucket.
lsPrefix :: MonadAWS m
  => BucketName
  -> Prefix
  -> m [S3Uri]
lsPrefix b p =
  runConduit $
    lsBucketStream (listObjectsV2 b & lovPrefix ?~ p)
    .| mapC (S3Uri b . view oKey)
    .| sinkList

-- | Deletes specified keys in a bucket.
-- Returns a list of keys that were successfully deleted.
--
-- Will fail monadically (using 'fail') if the response indicates any errors.
deleteFiles :: MonadAWS m
  => BucketName
  -> [ObjectKey]
  -> m [S3Uri]
deleteFiles b ks = do
  let dObjs = delete' & dObjects .~ (objectIdentifier <$> ks)
  resp <- AWS.send (deleteObjects b dObjs)
  unless (List.null $ resp ^. drsErrors) $
    fail (resp ^. drsErrors & show)
  let deleted = resp ^.. drsDeleted . each . dKey & catMaybes <&> S3Uri b
  pure deleted

-- | Deletes all the keys in a specified prefix EXCEPT the specified ones.
-- Returns a list of objects that were successfully deleted.
deleteFilesExcept :: MonadAWS m
  => BucketName
  -> Prefix
  -> [ObjectKey]
  -> m [S3Uri]
deleteFilesExcept b p uris = do
  existing <- lsPrefix b p
  case (objectKey <$> existing) List.\\ uris of
    [] -> pure []
    xs -> deleteFiles b xs

-- | Checks if the file exists on S3
fileExists :: MonadAWS m
  => S3Uri
  -> m Bool
fileExists (S3Uri b k) =
  isJust <$> handle404ToNone (AWS.send (headObject b k))