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

{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fno-warn-deprecations #-}

module Load
  ( main
  ) where

import "uniworx" Import hiding (Option(..), Normal, responseBody, responseStatus)
import Utils.Form (FormIdentifier(..))
import Handler.Admin.Test.Download (generateDownload', seedNew)

import System.Console.GetOpt

import qualified Data.Text as Text

import qualified Data.Map.Strict as Map

import Data.Ratio ((%))
import Data.Random.Normal
import qualified Control.Monad.Random.Class as Random
import System.Random (RandomGen)

import System.Exit (exitWith, ExitCode(..))
import System.IO (hPutStrLn)

import UnliftIO.Concurrent (threadDelay)

import System.Clock (getTime, Clock(Monotonic))
import qualified System.Clock as Clock

import Network.URI

import qualified Data.ByteString.Lazy as Lazy (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as CBS
import qualified Data.Char as Char (isSpace)

import Network.Wreq
import Network.Wreq.Types (FormValue(..))
import Network.Wreq.Session (Session)
import qualified Network.Wreq.Session as Session
import Network.HTTP.Client.MultipartFormData (partFileRequestBody)
import Network.HTTP.Client.TLS (tlsManagerSettings)

import qualified Text.HTML.Scalpel as Scalpel

import qualified Data.Conduit.Combinators as C
import Data.List (genericLength)

import qualified Control.Retry as Retry


data Normal k = Normal
  { dAvg :: k
  , dRelDev :: Centi
  } deriving (Eq, Ord, Read, Show, Generic)

sampleN :: (Random.MonadSplit g m, RandomGen g) => (k -> Centi -> k) -> Normal k -> m k
sampleN scale Normal{..}
  | dRelDev == 0 = return dAvg
  | otherwise = do
      gen <- Random.getSplit
      let (realToFrac -> r, _) = normal' (1, realToFrac dRelDev :: Double) gen
      return $ dAvg `scale` r

instance PathPiece k => PathPiece (Normal k) where
  toPathPiece Normal{dRelDev = MkFixed perc, dAvg}
    | perc == 0 = toPathPiece dAvg
    | otherwise = toPathPiece dAvg <> ";" <> toPathPiece perc <> "%"
  fromPathPiece t
    | (avg, relDev') <- Text.breakOn ";" t
    , Just relDev <- Text.stripSuffix "%" =<< Text.stripPrefix ";" relDev'
    = Normal <$> fromPathPiece avg <*> (MkFixed <$> fromPathPiece relDev)
    | otherwise
    = Normal <$> fromPathPiece t   <*> pure 0

scaleDiffTime :: DiffTime -> Centi -> DiffTime
scaleDiffTime (diffTimeToPicoseconds -> ps) s = picosecondsToDiffTime . round $ s * fromIntegral ps

sampleNDiffTime :: (Random.MonadSplit g m, RandomGen g) => Normal DiffTime -> m DiffTime
sampleNDiffTime = sampleN scaleDiffTime


scaleIntegral :: Integral n => n -> Centi -> n
scaleIntegral n s = round $ toRational n * toRational s

sampleIntegral :: (Random.MonadSplit g m, RandomGen g, Integral n) => Normal n -> m n
sampleIntegral = sampleN scaleIntegral


instance PathPiece DiffTime where
  toPathPiece = (toPathPiece :: Pico -> Text) . MkFixed . diffTimeToPicoseconds
  fromPathPiece t = fromPathPiece t <&> \(MkFixed ps :: Pico) -> picosecondsToDiffTime ps
  

data LoadSimulation
  = LoadSheetDownload
  | LoadSheetSubmission
  deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic)
  deriving anyclass (Universe, Finite)

nullaryPathPiece ''LoadSimulation $ camelToPathPiece' 1

data LoadOptions = LoadOptions
  { loadSimulations :: Map LoadSimulation SimulationOptions
  , loadBaseURI :: URI
  , loadToken :: Maybe Jwt
  , loadTerm :: TermId, loadSchool :: SchoolId, loadCourse :: CourseShorthand, loadSheet :: SheetName
  , loadUploadChunks :: Normal Natural, loadUploadChunkSize :: Normal Natural
  } deriving (Eq, Ord, Show, Generic)

instance Default LoadOptions where
  def = LoadOptions
    { loadSimulations = Map.empty
    , loadBaseURI = error "No BaseURI given"
    , loadToken = Nothing
    , loadTerm = error "No term given", loadSchool = error "No school given", loadCourse = error "No course given", loadSheet = error "No sheet given"
    , loadUploadChunks = Normal 48 0.11
    , loadUploadChunkSize = Normal (2^16) 0
    }

data SimulationOptions = SimulationOptions
  { simParallel :: Natural
  , simDelay, simDuration :: Normal DiffTime
  } deriving (Eq, Ord, Show, Generic)

instance Default SimulationOptions where
  def = SimulationOptions
    { simParallel = 1
    , simDelay = Normal 0 0
    , simDuration = Normal 10 0
    }


data SimulationContext = SimulationContext
  { loadOptions :: LoadOptions
  , simulationOptions :: SimulationOptions
  , targetDuration :: DiffTime
  , runtime :: forall m. MonadIO m => m DiffTime
  }


makeLenses_ ''LoadOptions
makeLenses_ ''SimulationOptions
makeLenses_ ''SimulationContext


_MapF :: (Finite k, Ord k) => Iso' (Map k v) (k -> Maybe v)
_MapF = iso (flip Map.lookup) (\f -> Map.fromList $ mapMaybe (\k -> (k, ) <$> f k) universeF)


argsDescr :: [OptDescr (Kleisli IO LoadOptions LoadOptions)]
argsDescr
  = [ Option ['n', 'p'] ["number", "parallel"] (ReqArg (\(splitArg -> (cloneIndexedTraversal -> f, arg)) -> Kleisli $ return . over f (set _simParallel arg)) "NATURAL") "Number of simulations to run in parallel"
    , Option ['r'] ["run"] (ReqArg (\(ppArg -> sim) -> Kleisli $ return . over (_loadSimulations . at sim) (<|> Just def)) "SIMULATION") "Run the given Simulation"
    , Option ['d'] ["duration"] (ReqArg (\(splitArg -> (cloneIndexedTraversal -> f, arg)) -> Kleisli $ return . over f (set _simDuration arg)) "DURATION") "Try to run each simulation to take up the given duration"
    , Option ['w', 's'] ["wait", "delay", "stagger"] (ReqArg (\(splitArg -> (cloneIndexedTraversal -> f, arg)) -> Kleisli $ return . over f (set _simDelay arg)) "DURATION") "Wait the given time before starting each simulation"
    , Option ['b', 'u'] ["base", "uri"] (ReqArg (\uriStr -> let uri = fromMaybe (error $ "Could not parse URI: " <> uriStr) $ parseURI uriStr in Kleisli $ return . set _loadBaseURI uri ) "URI") "Base URI"
    , Option ['t'] ["token"] (ReqArg (Kleisli . loadTokenFile) "FILE") "File containing bearer token"
    , Option [] ["tid", "term"] (ReqArg (\(ppArg -> tid) -> Kleisli $ return . set _loadTerm tid) "TERM") "TermId"
    , Option [] ["ssh", "school"] (ReqArg (\(ppArg -> ssh) -> Kleisli $ return . set _loadSchool ssh) "SCHOOL") "SchoolId"
    , Option [] ["csh", "course"] (ReqArg (\(ppArg -> csh) -> Kleisli $ return . set _loadCourse csh) "COURSE") "CourseName"
    , Option [] ["shn", "sheet"] (ReqArg (\(ppArg -> shn) -> Kleisli $ return . set _loadSheet shn) "SHEET") "SheetName"
    , Option [] ["chunks"] (ReqArg (\(ppArg -> cs) -> Kleisli $ return . set _loadUploadChunks cs) "NATURAL") "Number of chunks to upload"
    , Option [] ["chunk-size"] (ReqArg (\(ppArg -> cs) -> Kleisli $ return . set _loadUploadChunkSize cs) "NATURAL") "Size of chunks to upload"
    ]
  where
    splitArg :: PathPiece p => String -> (AnIndexedTraversal' LoadSimulation LoadOptions SimulationOptions, p)
    splitArg (Text.pack -> t)
      | (ref, arg) <- Text.breakOn ":" t
      , let refs = Text.splitOn "," ref
            sArg = Text.stripPrefix ":" arg
      , Just refs' <- if | is _Just sArg -> mapM fromPathPiece refs
                         | otherwise     -> Just []
      , Just arg' <- fromPathPiece $ fromMaybe ref sArg
      = (, arg') $ if
          | null refs' -> _loadSimulations . itraversed
          | otherwise  -> _loadSimulations . _MapF . itraversed . indices (`elem` refs') . iplens (fromMaybe def) (const Just)
      | otherwise
      = terror $ "Invalid option argument: " <> t

    ppArg :: PathPiece p => String -> p
    ppArg (Text.pack -> a) = fromMaybe (terror $ "Invalid option argument: " <> a) $ fromPathPiece a

    loadTokenFile :: FilePath -> LoadOptions -> IO LoadOptions
    loadTokenFile fp pOpts = do
      token <- Jwt . CBS.filter (not . Char.isSpace) <$> BS.readFile fp
      return $ pOpts & _loadToken ?~ token

main :: IO ()
main = do
  args <- map unpack <$> getArgs
  case getOpt Permute argsDescr args of
    (kl, [], []) -> do
      cfg <- over (mapped . _loadSimulations) (Map.filter $ (> 0) . simParallel) . (`runKleisli` def) . getDual $ foldMap Dual kl
      if | not . Map.null $ loadSimulations cfg
           -> imapM_ (\sim simOpts -> runReaderT (runSimulation sim) (cfg & _loadSimulations . at sim .~ Nothing, simOpts)) $ loadSimulations cfg
         | otherwise -> do
             hPutStrLn stderr $ usageInfo "uniworxload" argsDescr
             exitWith $ ExitFailure 2
    (_, _, errs) -> do
      forM_ errs $ hPutStrLn stderr
      hPutStrLn stderr $ usageInfo "uniworxload" argsDescr
      exitWith $ ExitFailure 2

runSimulation :: LoadSimulation -> ReaderT (LoadOptions, SimulationOptions) IO ()
runSimulation sim = do
  p <- view $ _2 . _simParallel
  delays <- replicateM (fromIntegral p) $ do
    d <- view $ _2 . _simDelay
    sampleNDiffTime d
   
  forConcurrently_ ([1..p] `zip` sort delays) $ \(n, d') -> do
    begin <- liftIO getCurrentTime

    dur <- view $ _2 . _simDuration
    tDuration <- sampleNDiffTime dur
    
    let MkFixed us = realToFrac d' :: Micro
    threadDelay $ fromInteger us
    start <- liftIO getCurrentTime
    print ("start", n, diffUTCTime start begin, utctDayTime start)

    cTime <- liftIO $ getTime Monotonic
    let running :: forall m. MonadIO m => m DiffTime
        running = do
          cTime' <- liftIO $ getTime Monotonic
          let diff = MkFixed . Clock.toNanoSecs $ cTime' - cTime :: Nano
              MkFixed ps = realToFrac diff :: Pico
          return $ picosecondsToDiffTime ps

    withReaderT (\(lO, sO) -> SimulationContext lO sO tDuration running) $ runSimulation' sim
    end <- liftIO getCurrentTime

    print ("end", n, diffUTCTime start begin, diffUTCTime end start)


delayRemaining :: (MonadReader SimulationContext m, MonadIO m, RealFrac r) => r -> m ()
delayRemaining p = do
  total <- asks targetDuration
  cTime <- join $ asks runtime
  let remaining = MkFixed . diffTimeToPicoseconds $ total - cTime :: Pico
      MkFixed us = realToFrac $ realToFrac remaining * p :: Micro
  threadDelay $ fromInteger us


runSimulation' :: LoadSimulation -> ReaderT SimulationContext IO ()
runSimulation' LoadSheetDownload = do
  session <- newLoadSession
  uri <- sheetZipURI
  resp <- liftIO . Session.get session $ uriToString id uri mempty
  void . evaluate $! resp
  -- print . length $ resp ^. responseBody
runSimulation' LoadSheetSubmission = do
  LoadOptions{..} <- asks loadOptions
  session <- newLoadSession
  let formURI = formURI' `relativeTo` loadBaseURI
        where formURI' = nullURI { uriPath = unpack . Text.intercalate "/" $ "." : formPath }
              (formPath, _) = renderRoute $ CSheetR loadTerm loadSchool loadCourse loadSheet SubmissionNewR
  resp <- liftIO . httpRetry . Session.get session $ uriToString id formURI mempty
  void . evaluate $! resp
  procStart <- join $ asks runtime
  -- Just formData <- return . getFormData FIDsubmission $ resp ^. responseBody
  -- Just addButtonData <- return . flip (runFormScraper FIDsubmission) (resp ^. responseBody) $ do
  --   let btnSel = "button" Scalpel.@: [Scalpel.hasClass "btn-mass-input-add"]
    
  --   name <- Scalpel.attr "name" btnSel
  --   value <- Scalpel.attr "value" btnSel
  --   guard $ value == "add__0__0"
  --   return $ toStrict name := value
  -- let miData = addButtonData : map addEmail formData
  --       where addEmail dat@(name := _)
  --               | "__add__0__fields__emails" `isSuffixOf` name = name := ("loadtest@example.invalid" :: Text)
  --               | otherwise = dat
  -- resp2 <- liftIO $ Session.post session (uriToString id formURI mempty) miData
  -- Just formData2 <- return . getFormData FIDsubmission $ resp2 ^. responseBody
  Just formData2 <- return . getFormData FIDsubmission $ resp ^. responseBody
  uploadSeed <- liftIO seedNew
  chunkCount <- sampleIntegral loadUploadChunks
  chunks <- replicateM (fromIntegral chunkCount) $ sampleIntegral loadUploadChunkSize
  simCtx <- ask
  let fileUploadPart = requestBodySourceChunked $
           yieldMany (zip [0..] chunks)
        .| runReaderC simCtx
             ( C.mapM $ \(ci, cs) ->
                 fromIntegral cs <$ delayRemaining ((1 % max 1 (genericLength chunks - ci)) :: Rational)
             )
        .| generateDownload' uploadSeed
  -- print $ ala Sum foldMap chunks
  Just fileData <- return . flip (runFormScraper FIDsubmission) (resp ^. responseBody) $ do
    let fileSel = "input" Scalpel.@: ["type" Scalpel.@= "file"]

    name <- Scalpel.attr "name" fileSel
    return $ partFileRequestBody (decodeUtf8 $ toStrict name) "loadtest.bin" fileUploadPart

  let subData = (:) fileData $ formData2 >>= \(name := (renderFormValue -> value)) -> do
        guard $ name /= encodeUtf8 (fileData ^. partName)
        return $ partBS (decodeUtf8 name) value
  void . evaluate $! subData
  procEnd <- join $ asks runtime

  print ("proc", procEnd - procStart)
  
  resp3 <- liftIO . httpRetry $ Session.post session (uriToString id formURI mempty) subData
  void . evaluate $! resp3
  where
    httpRetry act = Retry.recovering policy handlers $ \Retry.RetryStatus{..} -> do
      putStrLn $ "httpRetry; rsIterNumber = " <> tshow rsIterNumber
      act
      where policy = Retry.fullJitterBackoff 1e3 & Retry.limitRetriesByCumulativeDelay 10e6
            handlers = Retry.skipAsyncExceptions `snoc` Retry.logRetries suggestRetry logRetry

            suggestRetry :: forall m. Monad m => SomeException -> m Bool
            suggestRetry _ = return True

            logRetry :: forall e m.
                        ( Exception e
                        , MonadIO m
                        )
                       => Bool -- ^ Will retry
                       -> e
                       -> Retry.RetryStatus
                       -> m ()
            logRetry shouldRetry err status = liftIO . putStrLn . pack $ Retry.defaultLogMsg shouldRetry err status

    
-- runSimulation' other = terror $ "Not implemented: " <> tshow other

runFormScraper :: FormIdentifier -> Scalpel.Scraper Lazy.ByteString a -> Lazy.ByteString -> Maybe a
runFormScraper fid innerS = fmap join . flip Scalpel.scrapeStringLike $ 
  fmap listToMaybe . Scalpel.chroots "form" $ do
    fid' <- Scalpel.attr "value" $ "input" Scalpel.@: ["name" Scalpel.@= "form-identifier"]
    guard $ fid' == encodeUtf8 (fromStrict $ toPathPiece fid)

    innerS

getFormData :: FormIdentifier -> Lazy.ByteString -> Maybe [FormParam]
getFormData = flip runFormScraper $
  Scalpel.chroots ("input") $ do
    name <- Scalpel.attr "name" Scalpel.anySelector
    value <- Scalpel.attr "value" Scalpel.anySelector <|> pure ""
    return $ toStrict name := value
      

newLoadSession :: ReaderT SimulationContext IO Session
newLoadSession = do
  LoadOptions{..} <- asks loadOptions

  let withToken = case loadToken of
        Nothing -> id
        Just (Jwt bs) -> (:) (hAuthorization, "Bearer " <> bs) . filter ((/= hAuthorization) . fst)
      

  liftIO . Session.newSessionControl (Just mempty) $ tlsManagerSettings
    { managerModifyRequest = \req -> return $ req { requestHeaders = withToken $ requestHeaders req }
    , managerResponseTimeout = responseTimeoutNone
    }

sheetZipURI :: ReaderT SimulationContext IO URI
sheetZipURI = do
  LoadOptions{..} <- asks loadOptions
  let zipURI = nullURI { uriPath = unpack . Text.intercalate "/" $ "." : zipPath }
        where (zipPath, _) = renderRoute . CSheetR loadTerm loadSchool loadCourse loadSheet $ SZipR SheetExercise -- FIXME: Broken with ApprootUserGenerated
  return $ zipURI `relativeTo` loadBaseURI