module Handler.Utils.Download
  ( withFileDownloadTokenMaybe', withFileDownloadToken, withFileDownloadToken'
  , sendThisFile
  , sendFileReference
  , serveOneFile
  , serveSomeFiles
  , serveSomeFiles'
  , serveZipArchive
  , serveZipArchive'
  ) where

import Import.NoFoundation
import Foundation.Type
import Foundation.Authorization

import qualified Data.HashSet as HashSet
import qualified Data.HashMap.Strict as HashMap

import qualified Data.Conduit.Combinators as C

import Handler.Utils.Zip
import Handler.Utils.ContentDisposition
import Handler.Utils.Files


data DownloadTokenRestriction
  = DownloadRestrictSingle { downloadRestrictReference :: FileContentReference }
  | DownloadRestrictMultiple
  deriving (Eq, Ord, Read, Show, Generic, Typeable)
deriveJSON defaultOptions
  { constructorTagModifier = camelToPathPiece' 2
  , fieldLabelModifier = camelToPathPiece' 2
  } ''DownloadTokenRestriction

withFileDownloadTokenMaybe' :: forall url m.
                               ( HasRoute UniWorX url
                               , MonadHandler m, HandlerSite m ~ UniWorX
                               , MonadCrypto m
                               , MonadCryptoKey m ~ CryptoIDKey
                               , YesodAuthPersist UniWorX, AuthEntity UniWorX ~ User, AuthId UniWorX ~ UserId
                               )
                            => Maybe (ConduitT () (Either FileReference DBFile) m ())
                            -> url
                            -> m (SomeRoute UniWorX)
withFileDownloadTokenMaybe' mSource route = maybeT (return $ SomeRoute route) $ do
  let rApproot = authoritiveApproot $ urlRoute route
  case rApproot of
    ApprootDefault -> mzero
    _other -> return ()
  guardM . getsYesod $ \app -> views _appRoot (is _Just . ($ rApproot)) app
                            && views _appRoot ($ ApprootDefault) app /= views _appRoot ($ rApproot) app

  uid <- MaybeT maybeAuthId
  now <- liftIO getCurrentTime
  expireOffset <- getsYesod $ view _appDownloadTokenExpire
  restr <- case mSource of
    Just source -> do
      results <- lift . runConduit $ source .| C.take 2 .| C.foldMap (pure . either Just (const Nothing))
      return $ case results of
        [Just (FileReference{ fileReferenceContent = Just ref })] -> DownloadRestrictSingle ref
        _other -> DownloadRestrictMultiple
    Nothing -> return DownloadRestrictMultiple
  bearer <- lift $ bearerRestrict (urlRoute route) restr <$> bearerToken
    (HashSet.singleton $ Right uid)
    (Just uid)
    (HashMap.singleton BearerTokenRouteAccess . HashSet.singleton $ urlRoute route)
    Nothing
    (Just . Just $ addUTCTime expireOffset now)
    (Just now)
  encodedBearer <- lift $ encodeBearer bearer

  lift . setDownload $ SomeRoute @UniWorX route
    & over (urlRouteParams $ Proxy @UniWorX) ((toPathPiece GetBearer, toPathPiece encodedBearer) :)

  where
    setDownload :: SomeRoute UniWorX -> m (SomeRoute UniWorX)
    setDownload route' = do
      wantsDownload <- downloadFiles
      return $ route'
        & over (urlRouteParams $ Proxy @UniWorX) (bool id addDownload wantsDownload)
      where
        addDownload params | anyOf (folded . _1) (== toPathPiece GetDownload) params = params
                           | otherwise = (toPathPiece GetDownload, mempty) : params

withFileDownloadToken' :: forall file url m.
                          ( HasFileReference file
                          , HasRoute UniWorX url
                          , MonadHandler m, HandlerSite m ~ UniWorX
                          , MonadCrypto m
                          , MonadCryptoKey m ~ CryptoIDKey
                          , YesodAuthPersist UniWorX, AuthEntity UniWorX ~ User, AuthId UniWorX ~ UserId
                          )
                       => ConduitT () (Either file DBFile) m ()
                       -> url
                       -> m (SomeRoute UniWorX)
withFileDownloadToken' = withFileDownloadTokenMaybe' . Just . (.| C.map (first . view $ _FileReference . _1))


withFileDownloadToken :: forall file url m.
                         ( HasFileReference file
                         , HasRoute UniWorX url
                         , MonadHandler m, HandlerSite m ~ UniWorX
                         , MonadCrypto m
                         , MonadCryptoKey m ~ CryptoIDKey
                         , BearerAuthSite UniWorX
                         )
                      => ConduitT () file m ()
                      -> url
                      -> m (SomeRoute UniWorX)
withFileDownloadToken = withFileDownloadToken' . (.| C.map Left)

-- | Simply send a `File`-Value
sendThisFile :: (YesodAuthPersist UniWorX, AuthEntity UniWorX ~ User, AuthId UniWorX ~ UserId, YesodPersistRunner UniWorX) => DBFile -> HandlerFor UniWorX TypedContent
sendThisFile File{..}
  | Just fileContent' <- fileContent = do
      setCSPSandbox
      setContentDisposition' . Just $ takeFileName fileTitle
      let cType = simpleContentType (mimeLookup $ pack fileTitle) <> "; charset=utf-8"
      respondSourceDB cType $
        fileContent' .| C.map toFlushBuilder
  | otherwise = sendResponseStatus noContent204 ()

sendFileReference :: forall file a.
                     ( HasFileReference file
                     , BearerAuthSite UniWorX
                     , YesodPersistBackend UniWorX ~ SqlBackend
                     , YesodPersistRunner UniWorX
                     )
                  => file -> HandlerFor UniWorX a
sendFileReference (view (_FileReference . _1) -> fRef@FileReference{..}) = do
  whenIsJust fileReferenceContent $ \fRef' -> do
    dlRestr <- maybeCurrentBearerRestrictions
    case dlRestr of
      Just (DownloadRestrictSingle restrRef) | restrRef == fRef' -> return ()
      _other -> setCSPSandbox
    setContentDisposition' . Just $ takeFileName fileReferenceTitle
  let cType = simpleContentType (mimeLookup $ pack fileReferenceTitle) <> "; charset=utf-8"
  join . runDB $ respondFileConditional Nothing cType fRef

-- | Serve a single file, identified through a given DB query
serveOneFile :: forall file.
                ( HasFileReference file
                , BearerAuthSite UniWorX
                , YesodPersistBackend UniWorX ~ SqlBackend
                , YesodPersistRunner UniWorX
                ) => ConduitT () file (YesodDB UniWorX) () -> HandlerFor UniWorX TypedContent
serveOneFile source = do
  results <- runDB . runConduit $ source .| C.take 2 .| C.foldMap pure -- We don't need more than two files to make a decision below
  case results of
    [file] -> sendFileReference file
    []     -> notFound
    _other -> do
      $logErrorS "SFileR" "Multiple matching files found."
      error "Multiple matching files found."

-- | Serve one file directly or a zip-archive of files, identified through a given DB query
--
-- Like `serveOneFile`, but sends a zip-archive if multiple results are returned
serveSomeFiles :: forall file.
                  ( HasFileReference file
                  , BearerAuthSite UniWorX
                  , YesodPersistBackend UniWorX ~ SqlBackend
                  , YesodPersistRunner UniWorX
                  ) => FilePath -> ConduitT () file (YesodDB UniWorX) () -> HandlerFor UniWorX TypedContent
serveSomeFiles archiveName source = serveSomeFiles' archiveName $ source .| C.map Left

serveSomeFiles' :: forall file.
                   ( HasFileReference file
                   , BearerAuthSite UniWorX
                   , YesodPersistBackend UniWorX ~ SqlBackend
                   , YesodPersistRunner UniWorX
                   ) => FilePath -> ConduitT () (Either file DBFile) (YesodDB UniWorX) () -> HandlerFor UniWorX TypedContent
serveSomeFiles' archiveName source = do
  (source', results) <- runDB $ runPeekN 2 source

  $logDebugS "serveSomeFiles" . tshow $ length results

  case results of
    []         -> notFound
    [file]     -> either sendFileReference sendThisFile file
    _moreFiles -> do
      setCSPSandbox
      setContentDisposition' $ Just archiveName
      respondSourceDB typeZip $ do
        let zipComment = encodeUtf8 $ pack archiveName
        source' .| eitherC sourceFiles' (C.map id) .| produceZip ZipInfo{..} .| C.map toFlushBuilder

-- | Serve any number of files as a zip-archive of files, identified through a given DB query
--
-- Like `serveSomeFiles`, but always sends a zip-archive, even if a single file is returned
serveZipArchive :: forall file.
                   ( HasFileReference file
                   , BearerAuthSite UniWorX
                   , YesodPersistBackend UniWorX ~ SqlBackend
                   , YesodPersistRunner UniWorX
                   ) => FilePath -> ConduitT () file (YesodDB UniWorX) () -> HandlerFor UniWorX TypedContent
serveZipArchive archiveName source = serveZipArchive' archiveName $ source .| C.map Left

serveZipArchive' :: forall file.
                    ( HasFileReference file
                    , BearerAuthSite UniWorX
                    , YesodPersistBackend UniWorX ~ SqlBackend
                    , YesodPersistRunner UniWorX
                    ) => FilePath -> ConduitT () (Either file DBFile) (YesodDB UniWorX) () -> HandlerFor UniWorX TypedContent
serveZipArchive' archiveName source = do
  (source', results) <- runDB $ runPeekN 1 source

  $logDebugS "serveZipArchive" . tshow $ length results

  case results of
    []         -> notFound
    _moreFiles -> do
      setCSPSandbox
      setContentDisposition' $ Just archiveName
      respondSourceDB typeZip $ do
        let zipComment = encodeUtf8 $ pack archiveName
        source' .| eitherC sourceFiles' (C.map id) .| produceZip ZipInfo{..} .| C.map toFlushBuilder