{-# 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)

import System.Console.GetOpt

import qualified Data.Text as Text

import qualified Data.Map.Strict as Map

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 as BS
import qualified Data.ByteString.Char8 as CBS
import qualified Data.Char as Char (isSpace)

import Network.Wreq
import Network.Wreq.Session (Session)
import qualified Network.Wreq.Session as Session
import Network.HTTP.Client.TLS (tlsManagerSettings)


instance (a ~ b, Monad m) => Monoid (Kleisli m a b) where
  mempty = Kleisli return

instance (a ~ b, Monad m) => Semigroup (Kleisli m a b) where
  Kleisli f <> Kleisli g = Kleisli $ f <=< g


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

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


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

data LoadSimulation
  = LoadSheetDownload
  | LoadSheetSubmission
  deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Typeable)
  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
  } deriving (Eq, Ord, Show, Generic, Typeable)

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"
    }

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

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"
    ]
  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
  replicateConcurrently_ (fromIntegral p) $ do
    d <- view $ _2 . _simDelay
    d' <- sampleNDiffTime d

    dur <- view $ _2 . _simDuration
    tDuration <- sampleNDiffTime dur
    
    let MkFixed us = realToFrac d' :: Micro
    threadDelay $ fromInteger us

    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


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
  print . length $ resp ^. responseBody
runSimulation' other = terror $ "Not implemented: " <> tshow other


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

  let withToken = case loadToken of
        Nothing -> id
        Just (Jwt bs) -> (:) $ traceShowId (hAuthorization, "Bearer " <> bs)
      

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

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
  return $ zipURI `relativeTo` loadBaseURI