{-# LANGUAGE QuasiQuotes #-}

module TempMessageDb
  ( withDatabaseUrl,
    withConnection,
  )
where

import Control.Exception (Exception, IOException)
import Control.Monad (void)
import Control.Monad.Catch (Handler (Handler))
import qualified Control.Retry as Retry
import Data.ByteString (ByteString)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import Data.Monoid (getLast)
import qualified Database.PostgreSQL.Simple as Postgres
import qualified Database.PostgreSQL.Simple as Simple
import Database.PostgreSQL.Simple.Options (Options)
import qualified Database.PostgreSQL.Simple.Options as PostgresOptions
import Database.PostgreSQL.Simple.SqlQQ (sql)
import Database.Postgres.Temp (Accum (Merge))
import qualified Database.Postgres.Temp as PostgresTemp
import System.Environment (getEnvironment)
import qualified System.Process.Typed as Process
import qualified UnliftIO
import UnliftIO.Exception (bracket)


migrate :: Options -> IO ()
migrate options = do
  hostEnv <- getEnvironment

  let fromLast =
        fromMaybe "" . getLast

      processEnv =
        hostEnv
          <> [
               ( "PGHOST"
               , fromLast $ PostgresOptions.host options
               )
             ,
               ( "PGDATABASE"
               , fromLast $ PostgresOptions.dbname options
               )
             ,
               ( "PGPORT"
               , fromLast . fmap show $ PostgresOptions.port options
               )
             ,
               ( "PGUSER"
               , fromLast $ PostgresOptions.user options
               )
             ]
      command = Process.proc "test/message-db-init-scripts/database/install.sh" []
   in void . Process.readProcess_ $ Process.setEnv processEnv command

  let url = PostgresOptions.toConnectionString options
   in bracket (Simple.connectPostgreSQL url) Simple.close $ \connection -> do
        let query =
              [sql|
                CREATE ROLE test_user 
                WITH LOGIN PASSWORD 'password' 
                IN ROLE message_store;
              |]
         in void $ Simple.execute_ connection query

        let query =
              [sql|
                ALTER ROLE test_user 
                SET search_path TO message_store,public;
              |]
         in void $ Simple.execute_ connection query


withDatabaseUrl :: (ByteString -> IO a) -> IO a
withDatabaseUrl use = do
  let tempConfig =
        mempty
          { PostgresTemp.connectionOptions =
              mempty
                { PostgresOptions.user = pure "postgres"
                }
          , PostgresTemp.initDbConfig =
              Merge $
                mempty
                  { PostgresTemp.commandLine =
                      mempty{PostgresTemp.keyBased = Map.fromList [("--username=", Just "postgres")]}
                  }
          , PostgresTemp.postgresConfigFile =
              [ ("message_store.sql_condition", "on")
              ]
          }

  let retryPolicy =
          Retry.limitRetriesByCumulativeDelay 5_000_000 $
            Retry.exponentialBackoff 1_000

      exceptionHandlers =
        let restartFor :: forall e a. Exception e => a -> Handler IO Bool
            restartFor _ = Handler @_ @_ @e $ \_ -> pure True
         in [ restartFor @PostgresTemp.StartError
            , restartFor @IOException
            ]

  Retry.recovering retryPolicy exceptionHandlers $ \_ -> do
    result <- PostgresTemp.withConfig tempConfig $ \db ->
      let options = PostgresTemp.toConnectionOptions db
          dbUrl =
            PostgresOptions.toConnectionString $
              options
                { PostgresOptions.user = pure "test_user"
                , PostgresOptions.password = pure "password"
                , PostgresOptions.dbname = pure "message_store"
                }
       in migrate options *> use dbUrl

    case result of
      Left err -> UnliftIO.throwIO err
      Right value -> pure value


withConnection :: (Postgres.Connection -> IO a) -> IO a
withConnection use =
  withDatabaseUrl $ \url ->
    bracket (Simple.connectPostgreSQL url) Simple.close use