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

module Database
  ( main
  , truncateDb
  , module Database.Fill
  ) where

import "uniworx" Import hiding (Option(..), getArgs)
import "uniworx" Application (db', getAppSettings)

import Database.Persist.Postgresql
import Database.Persist.SqlBackend.Internal ( connEscapeFieldName )
import Control.Monad.Logger

import System.Console.GetOpt
import System.Exit (exitWith, ExitCode(..))
import System.IO (hPutStrLn)
import System.Environment (getArgs, withArgs)

import Database.Persist.Sql.Raw.QQ

import Database.Fill (fillDb)

import qualified Utils.Pool as Custom


data DBAction = DBClear
              | DBTruncate
              | DBMigrate
              | DBFill


argsDescr :: [OptDescr DBAction]
argsDescr =
  [ Option ['c'] ["clear"]    (NoArg DBClear)    "Delete everything accessable by the current database user"
  , Option ['t'] ["truncate"] (NoArg DBTruncate) "Truncate all tables mentioned in the current schema (This cannot be run concurrently with any other activity accessing the database)"
  , Option ['m'] ["migrate"]  (NoArg DBMigrate)  "Perform database migration"
  , Option ['f'] ["fill"]     (NoArg DBFill)     "Fill database with example data"
  ]


main :: IO ()
main = do
  args <- map unpack <$> getArgs
  case getOpt' Permute argsDescr args of
    (acts@(_:_), nonOpts, unrecOpts, []) -> withArgs (unrecOpts ++ nonOpts) . forM_ acts $ \case
      DBClear   -> runStderrLoggingT $ do -- We don't use `db` here, since we do /not/ want any migrations to run, yet
        settings <- liftIO getAppSettings
        withPostgresqlConn (pgConnStr $ appDatabaseConf settings) . runSqlConn $ do
          [executeQQ|drop owned by current_user|] :: ReaderT SqlBackend _ ()
      DBTruncate -> db' $ do
        foundation <- getYesod
        Custom.purgePool $ appConnPool foundation
        truncateDb
      DBMigrate -> db' $ return ()
      DBFill    -> db' $ fillDb
    (_, _, _, errs) -> do
      forM_ errs $ hPutStrLn stderr
      hPutStrLn stderr $ usageInfo "uniworxdb" argsDescr
      exitWith $ ExitFailure 2

truncateDb :: MonadIO m => ReaderT SqlBackend m ()
truncateDb = do
  tables <- map unSingle <$> [sqlQQ|SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'|]
  sqlBackend <- ask

  let escapedTables = map (connEscapeFieldName sqlBackend . FieldNameDB) $ filter (not . (`elem` protected)) tables -- ugh. We assume `connEscapeFieldName` behaves identically to `connEscapeTableName`
      query = "TRUNCATE TABLE " ++ intercalate ", " escapedTables ++ " RESTART IDENTITY"
      protected = ["applied_migration"]
  rawExecute query []