content: {-# LANGUAGE BangPatterns      #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-| This module abstracts over the details of persisting the value. Journaling is
    also handled here, if enabled. -}
  module Icepeak.Server.Persistence
    ( PersistentValue
    , PersistenceConfig (..)
    , getDataFile
    , getValue
    , apply
    , loadFromBackend
    , setupStorageBackend
    , syncToBackend
    ) where
  
  import           Control.Concurrent.STM
  import           Control.Exception
  import           Control.Monad.Except
  import           Control.Monad              (when)
  import           Control.Monad.Trans.Class  (lift)
  import           Control.Monad.IO.Class     (liftIO)
  import qualified Data.Aeson                 as Aeson
  import qualified Data.Aeson.Types           as Aeson
  import qualified Data.ByteString            as SBS
  import qualified Data.ByteString.Char8      as SBS8
  import qualified Data.ByteString.Lazy       as LBS
  import qualified Data.ByteString.Lazy.Char8 as LBS8
  import           Data.Foldable
  import           Data.Maybe                 (isJust)
  import           Data.Text                  (Text)
  import qualified Data.Text                  as Text
  import           Data.Traversable
  import           Database.SQLite.Simple     (FromRow (..), NamedParam (..), Only (..), execute,
                                               execute_, executeNamed, field, open, query_)
  import qualified System.Clock               as Clock
  import           System.Directory           (getFileSize, renameFile)
  import           System.Exit                (die)
  import           System.IO
  import           System.IO.Error (tryIOError, isDoesNotExistError, isPermissionError)
  import           Icepeak.Server.Logger      (Logger, LogLevel(..))
  import qualified Icepeak.Server.Logger      as Logger
  import qualified Icepeak.Server.Metrics     as Metrics
  import qualified Icepeak.Server.Store       as Store
  
  import Icepeak.Server.Config (StorageBackend (..))
  
  data PersistentValue = PersistentValue
    { pvConfig       :: PersistenceConfig
    , pvValue        :: TVar Store.Value
      -- ^ contains the state of the whole database
    , pvIsDirty      :: TVar Bool
      -- ^ flag indicating whether the current value of 'pvValue' has not yet been persisted to disk
    , pvJournal      :: Maybe Handle
    , pvJournalStats :: TVar JournalStatistics
      -- ^ Structure keeping track of the duration of journaling operations.
    }
  
  data JournalStatistics = JournalStatistics
    { journalCount :: !Int
      -- ^ The number of times a command has been journaled.
    , journalTimeTotalMicros  :: !Int
      -- ^ The total (sum) time journaling took.
    , journalTimeMaxMicros   :: !Int
      -- ^ The maximum time a journaling operation took.
    , journalTimeMinMicros   :: !Int
      -- ^ The minimum time a journaling operation took.
    }
  
  data PersistenceConfig = PersistenceConfig
    { pcDataFile    :: FilePath
    , pcJournalFile :: Maybe FilePath
    , pcLogger      :: Logger
    , pcMetrics     :: Maybe Metrics.IcepeakMetrics
    , pcLogSync     :: Bool
    }
  
  emptyJournalStatistics :: JournalStatistics
  emptyJournalStatistics = JournalStatistics
    { journalCount           = 0
    , journalTimeTotalMicros = 0
    , journalTimeMaxMicros   = 0
    , journalTimeMinMicros   = maxBound
    }
  
  -- | Get the actual value
  getValue :: PersistentValue -> STM Store.Value
  getValue = readTVar . pvValue
  
  -- | Apply a modification, and write it to the journal if enabled.
  apply :: Store.Modification -> PersistentValue -> IO ()
  apply op val = do
    -- append to journal if enabled
    for_ (pvJournal val) $ \journalHandle -> do
      start <- Clock.getTime Clock.Monotonic
      let entry = Aeson.encode op
      LBS8.hPutStrLn journalHandle entry
      end <- Clock.getTime Clock.Monotonic
      let time = fromInteger $ Clock.toNanoSecs (Clock.diffTimeSpec end start) `div` 1000
      -- Update the structure containing the statistics on journaling.
      atomically $ do
        modifyTVar' (pvJournalStats val) $
          \s -> s { journalCount           = journalCount s + 1
                  , journalTimeTotalMicros = journalTimeTotalMicros s + time
                  , journalTimeMaxMicros   = max (journalTimeMaxMicros s) time
                  , journalTimeMinMicros   = min (journalTimeMinMicros s) time
                  }
      for_ (pcMetrics . pvConfig $ val) $ \metrics -> do
        journalPos <- hTell journalHandle
        _ <- Metrics.incrementJournalWritten (LBS8.length entry) metrics
        Metrics.setJournalSize journalPos metrics
    -- update value
    atomically $ do
      modifyTVar (pvValue val) (Store.applyModification op)
      writeTVar (pvIsDirty val) True
  
  
  -- If no --data-file was supplied we default to either icepeak.json, for the file backend,
  -- or icepeak.db, for the Sqlite backend
  getDataFile :: StorageBackend -> Maybe FilePath -> FilePath
  getDataFile _ (Just filePath) = filePath
  getDataFile File _ = "icepeak.json"
  getDataFile Sqlite _ = "icepeak.db"
  
  -- * IO
  
  -- | Ensure that we can access the chosen storage file, and validate that it has the right file format
  setupStorageBackend :: StorageBackend -> FilePath -> IO ()
  setupStorageBackend File filePath = do
    eitherEncodedValue <- tryIOError $ SBS.readFile filePath
    case eitherEncodedValue of
      Left e | isDoesNotExistError e -> do
          -- If there is no icepeak.json file yet, we create an empty one instead.
          let message = "WARNING: Could not read data from " ++ filePath ++
                        " because the file does not exist yet. Created an empty database instead."
  
          -- if this fails, we want the whole program to crash since something is wrong
          SBS.writeFile filePath "{}"
          putStrLn message
  
      Left e | isPermissionError e ->
          die $ "File " ++ filePath ++ " cannot be read due to a permission error. Please check the file permissions."
      -- other errors should also lead to program termination
      Left e -> die (show e)
  
      -- in case the data-file is empty we write the empty object "{}" to it and return it
      Right "" -> do
          let message = "WARNING: The provided --data-file " ++ filePath ++
                        " is empty. Will write a default database of {} to this file."
          putStrLn message
          SBS.writeFile filePath "{}"
      Right encodedValue -> case Aeson.eitherDecodeStrict encodedValue of
          Left msg  -> die $ "Failed to decode the initial data in " ++ filePath ++ ": " ++ show msg
          Right (_value :: Aeson.Value) -> pure ()
  setupStorageBackend Sqlite filePath = do
    -- read the data from SQLite
    conn <- liftIO $ open filePath
    liftIO $ execute_ conn "CREATE TABLE IF NOT EXISTS icepeak (value BLOB)"
  
    jsonRows <- liftIO $ (query_ conn "SELECT * from icepeak" :: IO [JsonRow])
    case jsonRows of
      -- ensure that there is one row in the table, so that we can UPDATE it later
      [] -> liftIO $ execute conn "INSERT INTO icepeak (value) VALUES (?)" (Only $ Aeson.encode Aeson.emptyObject)
      _ -> pure()
  
  loadFromBackend :: StorageBackend -> PersistenceConfig -> IO (Either String PersistentValue)
  loadFromBackend backend config = runExceptT $ do
    let metrics = pcMetrics config
        dataFilePath = pcDataFile config
  
    -- We immediately set the dataSize metric, so that Prometheus can start scraping it
    liftIO $ forM_ metrics $ \metric -> do
      size <- getFileSize dataFilePath
      Metrics.setDataSize size metric
  
    -- Read the data from disk and parse it as an Aeson.Value
    value <- case backend of
      File -> readData dataFilePath
      Sqlite -> readSqliteData dataFilePath
  
    valueVar     <- lift $ newTVarIO value
    dirtyVar     <- lift $ newTVarIO False
    journalStats <- lift $ newTVarIO emptyJournalStatistics
    journal      <- for (pcJournalFile config) openJournal
  
    let val = PersistentValue
          { pvConfig       = config
          , pvValue        = valueVar
          , pvIsDirty      = dirtyVar
          , pvJournal      = journal
          , pvJournalStats = journalStats
          }
    recoverJournal val
    return val
  
  logJournalStats :: PersistentValue -> IO ()
  logJournalStats pv = do
    jStats <- atomically $ do
      stats <- readTVar (pvJournalStats pv)
      writeTVar (pvJournalStats pv) emptyJournalStatistics
      return stats
    let jCount = journalCount jStats
    when (jCount /= 0) $ do
      let jAvg   = journalTimeTotalMicros jStats `div` jCount
          jMin   = journalTimeMinMicros jStats
          jMax   = journalTimeMaxMicros jStats
          jLog   = Text.concat ["Journaling duration (avg/min/max/count) in microseconds since last Sync: "
                               , Text.intercalate "/" $ map (Text.pack . show) [jAvg, jMin, jMax, jCount]]
      logMessage pv jLog
  
  syncToBackend :: StorageBackend -> PersistentValue -> IO ()
  syncToBackend File pv = do
      start <- Clock.getTime Clock.Monotonic
      syncFile pv
      end <- Clock.getTime Clock.Monotonic
      let time = Clock.toNanoSecs (Clock.diffTimeSpec end start) `div` 1000000
      when (pcLogSync $ pvConfig pv) $ do
        logMessage pv $ Text.concat ["It took ", Text.pack $ show time, " ms to synchronize Icepeak on disk."]
        when (isJust $ pcJournalFile $ pvConfig pv) $ logJournalStats pv
  syncToBackend Sqlite pv = syncSqliteFile pv
  
  -- * SQLite loading and syncing
  
  -- | There is currently just one row and it contains only one column of type SBS.ByteString.
  -- This single field holds the whole JSON value for now.
  data JsonRow = JsonRow {jsonByteString :: SBS.ByteString} deriving (Show)
  
  instance FromRow JsonRow where
    fromRow = JsonRow <$> field
  
  -- | Read and decode the Sqlite data file from disk
  readSqliteData :: FilePath -> ExceptT String IO Store.Value
  readSqliteData filePath = ExceptT $ do
    -- read the data from SQLite
    conn <- liftIO $ open filePath
    jsonRows <- liftIO $ (query_ conn "SELECT * from icepeak" :: IO [JsonRow])
  
    case jsonRows of
      -- the 'setupStorageBackend' function verifies that we can read the database and that at least one row exists
      [] -> pure $ Right Aeson.emptyObject
      row : _ ->
        case Aeson.eitherDecodeStrict (jsonByteString row) of
              Left msg  -> pure $ Left $ "Failed to decode the initial data: " ++ show msg
              Right value -> pure $ Right (value :: Store.Value)
  
  -- | Write the data to the SQLite file if it has changed.
  syncSqliteFile :: PersistentValue -> IO ()
  syncSqliteFile val = do
    (dirty, value) <- atomically $ (,) <$> readTVar (pvIsDirty val)
                                       <*> readTVar (pvValue val)
                                       <*  writeTVar (pvIsDirty val) False
    -- simple optimization: only write when something changed
    when dirty $ do
      let filePath = pcDataFile $ pvConfig val
  
      conn <- open filePath
      -- we can always UPDATE here, since we know that there will be at least one row, since
      -- we issue an INSERT when we load in an empty database
      liftIO $ executeNamed conn "UPDATE icepeak SET value = :value" [":value" := Aeson.encode value]
  
      -- the journal is idempotent, so there is no harm if icepeak crashes between
      -- the previous and the next action
      truncateJournal val
  
      -- handle metrics last
      updateMetrics val
  
  -- * File syncing
  
  -- | Write the data to disk if it has changed.
  syncFile :: PersistentValue -> IO ()
  syncFile val = do
    (dirty, value) <- atomically $ (,) <$> readTVar (pvIsDirty val)
                                       <*> readTVar (pvValue val)
                                       <*  writeTVar (pvIsDirty val) False
    -- simple optimization: only write when something changed
    when dirty $ do
      let fileName = pcDataFile $ pvConfig val
          tempFileName = fileName ++ ".new"
      -- we first write to a temporary file here and then do a rename on it
      -- because rename is atomic on Posix and a crash during writing the
      -- temporary file will thus not corrupt the datastore
      LBS.writeFile tempFileName (Aeson.encode value)
      renameFile tempFileName fileName
  
      -- the journal is idempotent, so there is no harm if icepeak crashes between
      -- the previous and the next action
      truncateJournal val
  
      -- handle metrics last
      updateMetrics val
  
  -- * Private helper functions
  
  -- Note that some of these functions are still exported in order to be usable in the test suite
  
  -- | Seek to the beginning of the journal file and set the file size to zero.
  -- This should be called after all journal entries have been replayed, and the data has been
  -- synced to disk.
  truncateJournal :: PersistentValue -> IO ()
  truncateJournal val =
      for_ (pvJournal val) $ \journalHandle -> do
        -- we must seek back to the beginning of the file *before* calling hSetFileSize, since that
        -- function does not change the file cursor, which means that the first write that follows
        -- would fill up the file with \NUL bytes up to the original cursor position.
        hSeek journalHandle AbsoluteSeek 0
        hSetFileSize journalHandle 0
  
  -- | We keep track of three metrics related to persistence:
  -- 1. The current size of the data file
  -- 2. The current size of the journal file
  -- 3. The total amount of data written to disk since the process was started
  --    (not counting journal writes)
  updateMetrics :: PersistentValue -> IO ()
  updateMetrics val = do
      let filePath = pcDataFile . pvConfig $ val
          metrics = pcMetrics . pvConfig $ val
      forM_ metrics $ \metric -> do
        size <- getFileSize filePath
        Metrics.setDataSize size metric
        Metrics.setJournalSize (0 :: Int) metric
        Metrics.incrementDataWritten size metric
  
  -- | Open or create the journal file
  openJournal :: FilePath -> ExceptT String IO Handle
  openJournal journalFile = ExceptT $ do
    eitherHandle <- try $ do
      h <- openBinaryFile journalFile ReadWriteMode
      hSetBuffering h LineBuffering
      pure h
    case eitherHandle :: Either SomeException Handle of
      Left exc -> pure $ Left $ "Failed to open journal file: " ++ show exc
      Right fileHandle -> pure $ Right fileHandle
  
  -- | Read the modifications from the journal file, apply them and sync again.
  -- This should be done when loading the database from disk.
  recoverJournal :: PersistentValue -> ExceptT String IO ()
  recoverJournal pval = for_ (pvJournal pval) $ \journalHandle -> ExceptT $ fmap formatErr $ try $ do
    initialValue <- atomically $ readTVar (pvValue pval)
    (finalValue, successful, total) <- runRecovery journalHandle initialValue
  
    when (successful > 0) $ do
      atomically $ do
        writeTVar (pvValue pval) finalValue
        writeTVar (pvIsDirty pval) True
      -- syncing takes care of cleaning the journal
      syncFile pval
  
    when (total > 0) $ do
      logMessage pval "Journal replayed"
      logMessage pval $ "  failed:     " <> Text.pack (show $ total - successful)
      logMessage pval $ "  successful: " <> Text.pack (show $ successful)
  
    where
      formatErr :: Either SomeException a -> Either String a
      formatErr (Left exc) = Left $ "Failed to read journal: " ++ show exc
      formatErr (Right x)  = Right x
  
      runRecovery journalHandle value = do
        -- read modifications from the beginning
        hSeek journalHandle AbsoluteSeek 0
        foldJournalM journalHandle replayLine (value, 0 :: Integer, 0 :: Integer)
  
      replayLine line (!value, !successful, !total) = do
        when (total == 0) $ do
          logMessage pval "Journal not empty, recovering"
        case Aeson.eitherDecodeStrict line of
          Left err -> do
            let lineNumber = total + 1
            logMessage pval $ failedRecoveryMsg err lineNumber
            pure (value, successful, total + 1)
          Right op -> pure (Store.applyModification op value, successful + 1, total + 1)
  
      failedRecoveryMsg err line = "Failed to recover journal entry "
        <> Text.pack (show line) <> ": " <> Text.pack err
  
  -- | Read and decode the data file from disk
  readData :: FilePath -> ExceptT String IO Store.Value
  readData filePath = ExceptT $ do
    eitherEncodedValue <- tryIOError $ SBS.readFile filePath
    case eitherEncodedValue of
      -- we do not expect any errors here, since we validated the file earlier already
      Left e -> pure $ Left (show e)
      Right encodedValue -> case Aeson.eitherDecodeStrict encodedValue of
        Left msg  -> pure $ Left $ "Failed to decode the initial data: " ++ show msg
        Right value -> pure $ Right value
  
  -- | Log a message in the context of a PersistentValue.
  logMessage :: PersistentValue -> Text -> IO ()
  logMessage pval msg = Logger.postLogBlocking (pcLogger $ pvConfig pval) LogInfo msg
  
  -- | Left fold over all journal entries.
  foldJournalM :: Handle -> (SBS8.ByteString -> a -> IO a) -> a -> IO a
  foldJournalM h f = go
    where
      go !x = do
        eof <- hIsEOF h
        if eof
          then pure x
          else do
            line <- SBS8.hGetLine h
            x' <- f line x
            go x'
  