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

module Handler.Admin.Test.Download
  ( testDownload
  , generateDownload', Random.seedNew
  ) where

import Import hiding (Builder)
import Handler.Utils

import qualified Crypto.Random as Random

import qualified Data.ByteString.Base64.URL as Base64

import qualified Data.Conduit.Combinators as C

import qualified Data.Binary as Binary

import Data.Binary.Builder (Builder)

import Control.Monad.Random.Lazy (evalRandT, liftRandT)

import qualified Database.Esqueleto.Legacy as E hiding (random_)
import qualified Database.Esqueleto.PostgreSQL as E


data TestDownloadMode
  = TestDownloadDirect
  | TestDownloadInTransaction
  | TestDownloadFromDatabase
  deriving (Eq, Ord, Enum, Bounded, Read, Show, Generic)
  deriving anyclass (Universe, Finite, Binary)

nullaryPathPiece ''TestDownloadMode $ camelToPathPiece' 2
embedRenderMessage ''UniWorX ''TestDownloadMode id

data TestDownloadOptions = TestDownloadOptions
  { dlSeed :: Random.Seed
  , dlMaxSize, dlChunkSize :: Int
  , dlMode :: TestDownloadMode
  } deriving (Generic)
    deriving anyclass (Binary)

testDownloadForm :: Form TestDownloadOptions
testDownloadForm = identifyForm FIDTestDownload . renderWForm FormStandard $ do
  randomSeed <- liftIO Random.seedNew

  maxSizeRes <- wreq intField (fslI MsgTestDownloadMaxSize) . Just $ 2 * 2^30
  modeRes <- wpopt (selectField optionsFinite) (fslI MsgTestDownloadMode) $ Just TestDownloadDirect

  return $ TestDownloadOptions
        randomSeed
    <$> maxSizeRes
    <*> pure (2^20)
    <*> modeRes


generateDownload :: Monad m => TestDownloadOptions -> ConduitT i ByteString m ()
generateDownload TestDownloadOptions{..}
  = C.unfold genChunk dlMaxSize .| generateDownload' dlSeed
  where genChunk remaining
          | remaining >= dlChunkSize = Just (dlChunkSize, remaining - dlChunkSize)
          | remaining <= 0 = Nothing
          | otherwise = Just (remaining, 0)

generateDownload' :: Monad m => Random.Seed -> ConduitT Int ByteString m ()
generateDownload' seed = transPipe (evalRandT ?? Random.drgNewSeed seed) $
     C.mapM (liftRandT . (return .) . Random.randomBytesGenerate)


testDownload :: Handler Widget
testDownload = do
  ((dlRes, dlForm), dlEnctype) <- runFormPost testDownloadForm

  formResult dlRes $ \opts@TestDownloadOptions{..} -> do
    addHeader "Test-Download-Options" . decodeUtf8 . Base64.encode . toStrict $ Binary.encode opts
    setContentDisposition ContentAttachment $ Just "u2w-test-download.bin"
    let
      sendDownload :: forall m. Monad m => ConduitT ByteString (Flush Builder) m ()
      sendDownload = awaitForever $ \bs -> sendChunkBS bs >> sendFlush

      sourceDBChunks :: ConduitT () Int DB ()
      sourceDBChunks = forever sourceDBFiles
                    .| C.mapM (\x -> x <$ $logDebugS "testDownload.sourceDBChunks" (tshow $ entityKey x))
                    .| C.map ((length $!!) . fileContentChunkContent . entityVal)
                    .| takeLimit dlMaxSize
        where
          sourceDBFiles = E.selectSource . E.from $ \fileContentChunk -> do
            E.orderBy [E.asc $ E.random_ @Int64]
            return fileContentChunk

          takeLimit n | n <= 0 = return ()
          takeLimit n = do
            c <- await
            $logDebugS "testDownload.takeLimit" $ tshow c
            case c of
              Nothing -> return ()
              Just c' -> yield c' >> takeLimit (n - c')
    (sendResponse =<<) $ case dlMode of
      TestDownloadDirect        -> respondSource   typeOctet $ generateDownload opts .| sendDownload
      TestDownloadInTransaction -> respondSourceDB typeOctet $ generateDownload opts .| sendDownload
      TestDownloadFromDatabase  -> respondSourceDB typeOctet $ sourceDBChunks .| generateDownload' dlSeed .| sendDownload

  return $
    wrapForm dlForm def
      { formEncoding = dlEnctype
      }
