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

module Load
  ( main
  ) where

import "uniworx" Import hiding (Option(..), Normal)

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


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
  = 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
  } deriving (Eq, Ord, Show, Generic, Typeable)

instance Default LoadOptions where
  def = LoadOptions
    { loadSimulations = Map.empty
    }

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 (Endo LoadOptions)]
argsDescr
  = [ Option ['n', 'p'] ["number", "parallel"] (ReqArg (\(splitArg -> (cloneIndexedTraversal -> f, arg)) -> Endo . over f $ set _simParallel arg) "NATURAL") "Number of simulations to run in parallel"
    , Option ['r'] ["run"] (ReqArg (\(ppArg -> sim) -> Endo $ over (_loadSimulations . at sim) (<|> Just def)) "SIMULATION") "Run the given Simulation"
    , Option ['d'] ["duration"] (ReqArg (\(splitArg -> (cloneIndexedTraversal -> f, arg)) -> Endo . 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)) -> Endo . over f $ set _simDelay arg) "DURATION") "Wait the given time before starting each simulation"
    ]
  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

main :: IO ()
main = do
  args <- map unpack <$> getArgs
  case over _1 (over _loadSimulations (Map.filter $ (> 0) . simParallel) . (`appEndo` def) . getDual . foldMap Dual) $ getOpt Permute argsDescr args of
    (cfg, [], []) | not . Map.null $ loadSimulations cfg
      -> imapM_ (\sim simOpts -> runReaderT (runSimulation sim) (cfg & _loadSimulations . at sim .~ Nothing, simOpts)) $ loadSimulations cfg
    (_, _, 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' = liftIO . print