{-# LANGUAGE QuasiQuotes          #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE TupleSections        #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}

{-| Tests for the transactional DB API -}

module Test.Database.Transactions (
     tests
  ) where

import Control.Concurrent.Async (forConcurrently)
import Control.Exception.Safe
import Control.Exception.Safe qualified as Safe
import Control.Monad.Reader
import Control.Monad.Trans.Control
import Data.List.NonEmpty qualified as NE
import Data.Pool
import Data.Profunctor.Product.TH (makeAdaptorAndInstance)
import Data.String
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Database.PostgreSQL.Simple qualified as PG
import Database.PostgreSQL.Simple.FromField
import Database.PostgreSQL.Simple.FromRow
import Database.PostgreSQL.Simple.Options qualified as Client
import Database.PostgreSQL.Simple.SqlQQ (sql)
import Database.PostgreSQL.Simple.ToField
import Database.Postgres.Temp qualified as Tmp
import Gargantext.Core.Types.Individu
import Gargantext.Database.Query.Table.User
import Gargantext.Database.Schema.Prelude (Table (..))
import Gargantext.Database.Transactional
import Gargantext.Prelude hiding (throwIO, catch)
import Opaleye (selectTable, requiredTableField, SqlInt4)
import Opaleye qualified as O
import Prelude qualified
import Shelly as SH
import System.Random.Stateful
import Test.API.Setup (setupEnvironment)
import Test.Database.Setup
import Test.Database.Types hiding (Counter)
import Test.Hspec
import Test.Tasty.HUnit hiding (assert)
import Text.RawString.QQ
import Gargantext.Database.Action.User
import Gargantext.Database.Query.Table.Node.Error

--
-- For these tests we do not want to test the normal GGTX database queries, but rather
-- the foundational approach for the DBTx monad. Therefore we don't use the usual
-- 'withTestDB' code, but we rely on something very simple, a single table representing
-- counters with IDs, like so:
--
-- | ID | Counter_value |
-- | 1  | 0
-- | 2  | ...
--

newtype CounterId = CounterId { _CounterId :: Int }
  deriving (Show, Eq, ToField, FromField)

data Counter' id value = Counter
  { counterId    :: !id
  , counterValue :: value
  }
  deriving (Show, Eq)

type Counter = Counter' CounterId Int

$(makeAdaptorAndInstance "pCounter" ''Counter')

type CounterOpa = Counter' (O.Field SqlInt4) (O.Field SqlInt4)

countersTable :: Table CounterOpa CounterOpa
countersTable  =
  Table "ggtx_test_counter_table"
         ( pCounter
           Counter { counterId = requiredTableField "id"
                   , counterValue = requiredTableField "counter_value"
                   }
         )


newtype TestDBTxMonad a = TestDBTxMonad { _TestDBTxMonad :: TestMonadM DBHandle IOException a }
  deriving ( Functor, Applicative, Monad
           , MonadReader DBHandle, MonadError IOException
           , MonadBase IO
           , MonadBaseControl IO
           , MonadFail
           , MonadIO
           , MonadMask
           , MonadCatch
           , MonadThrow
           )

runTestDBTxMonad :: DBHandle -> TestMonadM DBHandle IOException a -> IO a
runTestDBTxMonad env m = do
  res <- flip runReaderT env . runExceptT . _TestMonad $ m
  case res of
    Left err -> throwIO $ Prelude.userError ("runTestDBTxMonad: " <> displayException err)
    Right x  -> pure x

setup :: IO DBHandle
setup = do
  res <- Tmp.startConfig tmpPgConfig
  case res of
    Left err -> Prelude.fail $ show err
    Right db -> do
      let idleTime = 60.0
      let maxResources = 2
      let poolConfig = defaultPoolConfig (PG.connectPostgreSQL (Tmp.toConnectionString db))
                                         PG.close
                                         idleTime
                                         maxResources
      pool <- newPool (setNumStripes (Just 2) poolConfig)
      bootstrapCounterDB db pool
      pure $ DBHandle pool db
  where
    tmpPgConfig :: Tmp.Config
    tmpPgConfig = Tmp.defaultConfig <>
      Tmp.optionsToDefaultConfig mempty
        { Client.dbname   = pure dbName
        , Client.user     = pure dbUser
        , Client.password = pure dbPassword
        }

dbUser, dbPassword, dbName, dbTable :: String
dbUser = "ggtx_test_counter_db_user"
dbPassword = "ggtx_test_counter_db_pwd"
dbName = "ggtx_test_counter_db"
dbTable = "public.ggtx_test_counter_table"

bootstrapCounterDB :: Tmp.DB -> Pool PG.Connection -> IO ()
bootstrapCounterDB tmpDB pool = withResource pool $ \conn -> do
  void $ PG.execute_ conn (fromString $ "ALTER USER \"" <> dbUser <> "\" with PASSWORD '" <> dbPassword <> "'")
  let schemaContent = counterDBSchema
  let connString = Tmp.toConnectionString tmpDB
  (res,ec) <- shelly $ silently $ escaping False $ do
    withTmpDir $ \tdir -> do
      let schemaPath = tdir <> "/schema.sql"
      writefile schemaPath (T.pack schemaContent)
      result <- SH.run "psql" ["-d", "\"" <> TE.decodeUtf8 connString <> "\"", "<", fromString schemaPath]
      (result,) <$> lastExitCode
  unless (ec == 0) $ Safe.throwIO (Prelude.userError $ show ec <> ": " <> T.unpack res)

counterDBSchema :: String
counterDBSchema = [r|
  CREATE TABLE |] <> dbTable <> [r| (
      id SERIAL,
      counter_value INT NOT NULL DEFAULT 0,
      PRIMARY KEY (id)
  );
  ALTER TABLE public.ggtx_test_counter_table OWNER TO |] <> dbUser <> ";" <> [r|
  INSERT INTO public.ggtx_test_counter_table(counter_value) VALUES(42);
|]

withTestCounterDB :: (DBHandle -> IO ()) -> IO ()
withTestCounterDB = Safe.bracket setup teardown

teardown :: DBHandle -> IO ()
teardown test_db = do
  destroyAllResources $ _DBHandle test_db
  Tmp.stop $ _DBTmp test_db

--
-- Helpers and transactions to work with counters
--

instance PG.FromRow Counter where
  fromRow = Counter <$> field <*> field

getCounterById :: CounterId -> DBQuery IOException r Counter
getCounterById (CounterId cid) = do
  xs <- mkPGQuery [sql| SELECT * FROM public.ggtx_test_counter_table WHERE id = ?; |] (PG.Only cid)
  case xs of
    [c] -> pure c
    rst -> dbFail $ Prelude.userError $ "getCounterId returned more than one result: " <> show rst

insertCounter :: DBUpdate IOException Counter
insertCounter = do
  mkPGUpdateReturningOne [sql| INSERT INTO public.ggtx_test_counter_table(counter_value) VALUES(0) RETURNING id, counter_value|] ()

updateCounter :: CounterId -> Int -> DBUpdate IOException Counter
updateCounter cid x = do
  mkPGUpdateReturningOne [sql| UPDATE public.ggtx_test_counter_table SET counter_value = ? WHERE id = ? RETURNING *|] (x, cid)

-- | We deliberately write this as a composite operation.
stepCounter :: CounterId -> DBUpdate IOException Counter
stepCounter cid = do
  Counter{..} <- getCounterById cid
  mkPGUpdateReturningOne [sql| UPDATE public.ggtx_test_counter_table SET counter_value = ? WHERE id = ? RETURNING *|] (counterValue + 1, cid)

--
-- MAIN TESTS
--

tests :: Spec
tests = describe "Database Transactions" $ do
  counterDBTests
  ggtxDBTests

-- | Testing the transactional behaviour outside the classic GGTX operations.
-- We test that throwing exceptions in IO leads to rollbacks.
counterDBTests :: Spec
counterDBTests = parallel $ around withTestCounterDB $
  describe "Counter Transactions" $ do
    describe "Opaleye count queries" $ do
      it "Supports counting rows" opaCountQueries
    describe "Pure PG Queries" $ do
      it "Simple query works" simplePGQueryWorks
    describe "Pure PG Inserts" $ do
      it "Simple insert works" simplePGInsertWorks
    describe "Pure PG Updates" $ do
      it "Simple updates works" simplePGUpdateWorks
    describe "PG Queries and Updates" $ do
      it "Supports mixing queries and updates" mixQueriesAndUpdates
    describe "Rollback support" $ do
      it "can rollback in case of errors" testRollback
    describe "Read/Write Consistency" $ do
      it "should return a consistent state to different actors" testConsistency

-- | Testing the transactional behaviour inside the classic GGTX operations.
-- We test that throwing something like a 'NodeError' results in a proper rollback.
ggtxDBTests :: Spec
ggtxDBTests = parallel $ around withTestDB $ beforeWith (\ctx -> setupEnvironment ctx >>= (const $ pure ctx)) $
  describe "GGTX Transactions" $ do
    describe "Rollback support" $ do
      it "can rollback if a ggtx error gets thrown" testGGTXErrorRollback

simplePGQueryWorks :: DBHandle -> Assertion
simplePGQueryWorks env = runTestDBTxMonad env $ do
  x <- runDBQuery $ getCounterById (CounterId 1)
  liftIO $ counterValue x `shouldBe` 42

simplePGInsertWorks :: DBHandle -> Assertion
simplePGInsertWorks env = runTestDBTxMonad env $ do
  x <- runDBTx $ insertCounter
  liftIO $ x `shouldBe` (Counter (CounterId 2) 0)

simplePGUpdateWorks :: DBHandle -> Assertion
simplePGUpdateWorks env = runTestDBTxMonad env $ do
  x <- runDBTx $ updateCounter (CounterId 1) 99
  liftIO $ x `shouldBe` (Counter (CounterId 1) 99)

mixQueriesAndUpdates :: DBHandle -> Assertion
mixQueriesAndUpdates env = runTestDBTxMonad env $ do
  (final_1, final_2) <- runDBTx $ do
    c1 <- insertCounter
    c2 <- insertCounter
    c1' <- getCounterById (counterId c1)
    c2' <- stepCounter (counterId c2)
    pure (c1', c2')
  liftIO $ do
    final_1 `shouldBe` (Counter (CounterId 2) 0)
    final_2 `shouldBe` (Counter (CounterId 3) 1)

testRollback :: DBHandle -> Assertion
testRollback env = runTestDBTxMonad env $ do
  initialCounter <- runDBTx $ insertCounter >>= stepCounter . counterId
  liftIO $ counterValue initialCounter `shouldBe` 1
  -- Let's do another transaction where at the very last instruction we
  -- fail.
  handleDBTxError (\(_ :: IOException) -> pure ()) $ runDBTx $ do
    _x' <- stepCounter (counterId initialCounter)
    dbFail $ Prelude.userError "urgh"

  -- Let's check that the second 'stepCounter' didn't actually modified the counter's value.
  finalCounter <- runDBTx $ getCounterById (counterId initialCounter)
  liftIO $ counterValue finalCounter `shouldBe` 1

-- | In this test we create concurrent actors all writing to the /same/ counter.
-- Each one should observe only the state it's updating.
testConsistency :: DBHandle -> Assertion
testConsistency env = do

  let competing_actors = 10
  initialCounter <- runTestDBTxMonad env $ runDBTx insertCounter

  results <- forConcurrently [ 1 .. competing_actors ] $ \x -> runTestDBTxMonad env $ do
    -- random delay
    liftIO $ do
      delay_us <- uniformRM (100, 2_000_000) globalStdGen
      threadDelay delay_us
    runDBTx $ do
      _ <- updateCounter (counterId initialCounter) x
      getCounterById (counterId initialCounter)

  -- Each actor should observe a consistent state.
  liftIO $ results `shouldBe` map (Counter (CounterId 2)) [ 1 .. competing_actors ]

opaCountQueries :: DBHandle -> Assertion
opaCountQueries env = runTestDBTxMonad env $ do
  num0 <- runDBTx $ mkOpaCountQuery (selectTable countersTable)
  liftIO $ num0 @?= 1 -- Returns the master counter created alongside the schema.

  num <- runDBTx $ do
    _ <- insertCounter
    _ <- insertCounter
    mkOpaCountQuery (selectTable countersTable)
  liftIO $ num @?= 3

-- | In this simple test we create a user node in GGTX, we try
-- to update it, and check that if we throw an error in the update
-- transaction, the changes are not propagated
testGGTXErrorRollback :: TestEnv -> Assertion
testGGTXErrorRollback env = runTestMonadM @NodeError env $ do
  let ur = NewUser "alfredo" "alfredo@foo.com" (GargPassword "mypass")
  let newUsers = ur NE.:| []
  hashed <- liftIO $ mapM toUserHash newUsers
  void $ runDBTx $ insertNewUsers hashed
  -- Retrieve the user, check the details
  insertedUr <- runDBQuery $ getUserLightDB (UserName "alfredo")
  liftIO $ userLight_username insertedUr `shouldBe` "alfredo"
  -- CRUCIAL bit: try to update the email, throw an exception in the same tx block
  void $ (runDBTx $ do
    void $ updateUserEmail (insertedUr { userLight_email = "alfredo@bar.com" })
    nodeError $ NoRootFound -- it doesn't matter which exception
    ) `catchDBTxError` \(_e :: NodeError) -> pure () -- swallow it.

  -- let's check that the email hasn't been changed.
  insertedUr' <- runDBQuery $ getUserLightDB (UserName "alfredo")
  liftIO $ userLight_email insertedUr' `shouldBe` "alfredo@foo.com"

