{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}

{--| This module exposes a custom monad and functions to model database operations within Gargantext.
 The peculiarity of the custom monad is that it describe a DSL for the operations we wish to perform,
 and there is an evaluator that runs everything at the end of the transactional block. This means
 that we can \"bundle\" different SQL queries together in a single IO operation (in the evaluation) and
 we can wrap everything with \"withTransaction\", meaning that everything will run in a single database
 transaction \"by construction\", and using the same DB connection. This limits greatly the surface area
 for concurrency bugs stemming from indepedent DB transactions running against the same subset of data. -}

module Gargantext.Database.Transactional (
    DBOperation
  , DBTransactionOp -- opaque
  , DBTx -- opaque
  , DBUpdate
  , DBQuery
  , DBTxCmd
  -- * Executing queries and updates
  , runDBQuery
  , runDBTx

  -- * Smart constructors
  , mkPGQuery
  , mkPGUpdate
  , mkPGUpdateReturningOne
  , mkPGUpdateReturningMany
  , mkOpaQuery
  , mkOpaCountQuery
  , mkOpaUpdate
  , mkOpaInsert
  , mkOpaDelete

  -- * Throwing errors (which allows rollbacks)
  , dbFail
  ) where

import Control.Exception.Safe qualified as Safe
import Control.Lens
import Control.Monad.Base
import Control.Monad.Error.Class
import Control.Monad.Free
import Control.Monad.Trans.Control (MonadBaseControl, control)
import Data.Int (Int64)
import Data.Pool (withResource, Pool)
import Data.Profunctor.Product.Default
import Database.PostgreSQL.Simple qualified as PG
import Database.PostgreSQL.Simple.Transaction qualified as PG
import Gargantext.Database.Class
import Opaleye
import Prelude
import Control.Monad.Free.Church

data DBOperation = DBRead | DBWrite

-- | A functor describing a single operation on the database. Each constructor takes a continuation
-- argument which can be used to derive a sound 'Functor' instance, making this viable to be applied
-- in a monadic/free context.
--
-- /IMPORTANT/: If you read this comment, you are probably wondering about adding more operations
-- to this functor, maybe because the query/update you are writing can't be expressed directly with
-- the subset here. While adding more operations is certainly possible, it's also important to
-- understand that this monad has to contain as little constructors as possible as it's expressing
-- /THE OPERATIONS WE CAN SAFELY DO IN A DATABASE TRANSACTION/. In other terms, everything we add
-- here will have to run in a DB transaction, and the latter might at /any/ point have to rollback!
-- This means that destructive operations can't be rolled back, so they /shouldn't/ be modelled as
-- new 'DBTransactionOp' constructors.
--
-- Here is a practical example of operations we /could/ add in the future, which would still be ok:
--
-- GetTxCurrentTime :: (UTCTime -> next) -> DBTransactionOp err r next
-- LogMsg :: LogLevel -> T.Text -> (() -> next) -> DBTransactionOp err r next
--
-- Here is a practical example of operations which /ARE NOT OK/:
--
-- EmitCENotification :: CEMessage -> (() -> next) -> DBTransactionOp err r next
--
-- What's the difference between what's OK and what's not? The rollback. If we add the ability to get
-- the current time, or log a message, even if we rollback, nothing destructive has happened: sure, in
-- the worst case scenario we would be logging a message for something that eventually we aborted, but
-- we haven't compromised the internal correctness of the system.
--
-- Conversely, if we emit a CE message /and then we have to rollback due to an exception/, now it means
-- we have notified upstream about something that we eventually cancelled! That is wrong, and is a bug.
--
-- Please refer to your best judgement and add here only operations which do have some meanining in the
-- context of a DB transaction, and keep the rollback behaviour as your north-star to decided whether or
-- not that is a good idea.
-- Everything else can just be passed as an input to a 'DBTransactionOp' or simply returned as a result
-- from a 'DBTransactionOp', and later used in the concreted monad. That's what we do for CE notifications;
-- we have DB operations returning them, and we fire them \"outside\".
data DBTransactionOp err (r :: DBOperation) next where
  -- | A Postgres /read/, returning a list of results. The 'r' in the result is polymorphic
  -- so that reads can be embedded in updates transactions.
  PGQuery   :: (PG.ToRow q, PG.FromRow a) => PG.Query -> q -> ([a] -> next) -> DBTransactionOp err r next
  -- | A Postgres /write/, returning the number of affected rows. It can be used only in
  -- 'DBWrite' transactions.
  PGUpdate  :: PG.ToRow a => PG.Query -> a -> (Int64 -> next) -> DBTransactionOp err DBWrite next
  -- | Unlike a 'PGUpdate' that returns the list of affected rows, this can be used
  -- to write updates that returns a value via the \"RETURNING\" directive. It's the programmer's
  -- responsibility to ensure that the SQL fragment contains it.
  PGUpdateReturningMany  :: (PG.ToRow q, PG.FromRow a) => PG.Query -> q -> ([a] -> next) -> DBTransactionOp err DBWrite next
  -- | Ditto as above, but the contract is that the query has to return /exactly one/ result.
  PGUpdateReturningOne   :: (PG.ToRow q, PG.FromRow a) => PG.Query -> q -> (a -> next) -> DBTransactionOp err DBWrite next
  -- | An Opaleye /read/, returning a list of results. The 'r' in the result is polymorphic
  -- so that reads can be embedded in updates transactions.
  OpaQuery  :: Default FromFields fields a => Select fields -> ([a] -> next) -> DBTransactionOp err r next
  OpaCountQuery  :: Select a -> (Int -> next) -> DBTransactionOp err r next
  -- | An Opaleye /write/, returning a result depending on the input 'Insert'. It can be used only in
  -- 'DBWrite' transactions.
  OpaInsert :: Insert a -> (a -> next) -> DBTransactionOp err DBWrite next
  -- | An Opaleye /write/, returning a result depending on the input 'Update'. It can be used only in
  -- 'DBWrite' transactions.
  OpaUpdate :: Update a -> (a -> next) -> DBTransactionOp err DBWrite next
  -- | An Opaleye /delete/, returning a result depending on the input 'Delete'. It can be used only in
  -- 'DBWrite' transactions.
  OpaDelete :: Delete a -> (a -> next) -> DBTransactionOp err DBWrite next
  -- | Monadic failure for DB transactions.
  DBFail    :: err -> DBTransactionOp err r next

-- | A 'DBTx' is a free monad (using the free church-encoding 'F') using 'DBTransactionOp' as the functor.
-- In practical terms, it's just a monad where we can execute just the operations described by the
-- 'DBTransactionOp', and nothing more.
newtype DBTx err r a = DBTx { _DBTx :: F (DBTransactionOp err r) a }
  deriving (Functor, Applicative, Monad)

type DBQuery err r a    = DBTx err r a
type DBUpdate err a     = DBTx err DBWrite a
type DBReadOnly err r a = DBTx err DBRead a

-- Strict constraints to perform transactional read and writes.
-- Isomorphic to a DBCmd, but it doesn't impose a 'HasConfig' constraint, as
-- values can always be passed as parameters of a query or update.
type DBTxCmd err a = forall m env. (IsCmd env err m, HasConnectionPool env) => m a

instance Functor (DBTransactionOp err r) where
  fmap f = \case
    PGQuery   q params cont        -> PGQuery   q params (f . cont)
    PGUpdate  q a cont             -> PGUpdate  q a      (f . cont)
    PGUpdateReturningOne q a cont  -> PGUpdateReturningOne q a      (f . cont)
    PGUpdateReturningMany q a cont -> PGUpdateReturningMany q a      (f . cont)
    OpaQuery  sel cont             -> OpaQuery  sel      (f . cont)
    OpaCountQuery sel cont         -> OpaCountQuery  sel (f . cont)
    OpaInsert ins cont             -> OpaInsert ins      (f . cont)
    OpaUpdate upd cont             -> OpaUpdate upd      (f . cont)
    OpaDelete del cont             -> OpaDelete del      (f . cont)
    DBFail    err                  -> DBFail err

-- | Generalised version of 'withResource' to work over any unlifted monad.
-- For some reason 'resource-pool' removed this from version 0.3.0.0 onwards.
withResourceM :: MonadBaseControl IO m => Pool a -> (a -> m b) -> m b
withResourceM pool func = control $ \run -> withResource pool (run . func)

-- | Generalised version of 'withTransaction' to work over any unlifted monad.
withTransactionM
  :: forall m a.
     MonadBaseControl IO m
  => PG.Connection
  -> m a
  -> m a
withTransactionM conn action = control $ \runInIO -> PG.withTransaction conn $ runInIO action

-- | Executes the input action in a single PostgreSQL "read-only" transaction,
-- suitable for read-only queries.
withReadOnlyTransactionM
  :: forall m a.
     MonadBaseControl IO m
  => PG.Connection
  -> m a
  -> m a
withReadOnlyTransactionM conn action =
  control $ \runInIO ->
    PG.withTransactionMode tmode conn (runInIO action)
  where
    tmode :: PG.TransactionMode
    tmode = PG.TransactionMode PG.DefaultIsolationLevel PG.ReadOnly

-- | Run a PostgreSQL transaction, suitable for operations that mixes read and writes,
-- and actually the only choice available to run 'DBUpdate' operations.
runDBTx :: DBUpdate err a -> DBTxCmd err a
runDBTx (DBTx m) = do
  pool <- view connPool
  withResourceM pool $ \conn -> withTransactionM conn $ foldF (evalOp conn) m

-- | Runs a DB query.
-- /NOTE/ the input type is 'DBReadOnly', i.e. a transaction where /all/
-- the operations are 'DBRead'. This makes impossible to sneak in updates
-- into otherwise read-only queries.
runDBQuery :: DBReadOnly err r a -> DBTxCmd err a
runDBQuery (DBTx m) = do
  pool <- view connPool
  withResourceM pool $ \conn -> withReadOnlyTransactionM conn $ foldF (evalOp conn) m

-- | The main evaluator, turns our pure operations into side-effects that run into the
-- 'DBCmd'.
evalOp :: PG.Connection -> DBTransactionOp err r a -> DBTxCmd err a
evalOp conn = \case
  PGQuery   qr q cc             -> cc <$> liftBase (PG.query conn qr q)
  PGUpdate  qr a cc             -> cc <$> liftBase (PG.execute conn qr a)
  PGUpdateReturningOne qr a cc  -> cc <$> liftBase (queryOne conn qr a)
  PGUpdateReturningMany qr a cc -> cc <$> liftBase (PG.query conn qr a)
  OpaQuery  sel  cc             -> cc <$> liftBase (runSelect conn sel)
  OpaCountQuery  sel  cc        -> cc <$> liftBase (evalOpaCountQuery conn sel)
  OpaInsert ins  cc             -> cc <$> liftBase (runInsert conn ins)
  OpaUpdate upd  cc             -> cc <$> liftBase (runUpdate conn upd)
  OpaDelete del  cc             -> cc <$> liftBase (runDelete conn del)
  DBFail    err                 -> throwError err

evalOpaCountQuery :: PG.Connection -> Select a -> IO Int
evalOpaCountQuery conn sel = do
  counts <- runSelect conn $ countRows sel
  -- countRows is guaranteed to return a list with exactly one row so DL.head is safe here
  pure $ fromIntegral @Int64 @Int $ head counts

-- | Executes a Postgres query that /HAS/ to contain a \"RETURNING ..\" statement /AND/
-- a \"where\" predicate that indexes a /UNIQUE/ field, i.e. something like \"WHERE id = ? RETURNING id\".
-- The above SQL fragment ensures that Postgres will return us a single row with a single result, which
-- is what we are enforcing here.
-- This function will fail with an error (and the whole transaction will rollback) if the SQL query
-- violates the contract.
queryOne :: (PG.ToRow q, PG.FromRow r) => PG.Connection -> PG.Query -> q -> IO r
queryOne conn q v = do
  rs <- PG.query conn q v
  case rs of
    [x] -> pure x
    [ ] -> Safe.throwIO $ userError "queryOne: no result returned. Check your SQL!"
    _   -> Safe.throwIO $ userError "queryOne: more than one result returned. Have you used the 'RETURNING' directive?"

--
-- Smart constructors
--
-- The following functions are just mechanical wrappers around the raw data constructors, which
-- we are not exposing for information hiding purposes.
--

dbFail :: err -> DBTx err r b
dbFail = DBTx . liftF . DBFail

mkPGQuery :: (PG.ToRow q, PG.FromRow a)
           => PG.Query
           -> q
           -> DBQuery err r [a]
mkPGQuery q a = DBTx $ liftF (PGQuery q a id)

mkPGUpdate :: PG.ToRow a => PG.Query -> a -> DBUpdate err Int64
mkPGUpdate q a = DBTx $ liftF (PGUpdate q a id)

mkPGUpdateReturningOne :: (PG.ToRow q, PG.FromRow a) => PG.Query -> q -> DBUpdate err a
mkPGUpdateReturningOne q a = DBTx $ liftF (PGUpdateReturningOne q a id)

mkPGUpdateReturningMany :: (PG.ToRow q, PG.FromRow a) => PG.Query -> q -> DBUpdate err [a]
mkPGUpdateReturningMany q a = DBTx $ liftF (PGUpdateReturningMany q a id)

mkOpaQuery :: Default FromFields fields a
           => Select fields
           -> DBQuery err x [a]
mkOpaQuery s = DBTx $ liftF (OpaQuery s id)

mkOpaCountQuery :: Select fields
                -> DBQuery err x Int
mkOpaCountQuery s = DBTx $ liftF (OpaCountQuery s id)

mkOpaUpdate :: Update a -> DBUpdate err a
mkOpaUpdate a = DBTx $ liftF (OpaUpdate a id)

mkOpaInsert :: Insert a -> DBUpdate err a
mkOpaInsert a = DBTx $ liftF (OpaInsert a id)

mkOpaDelete :: Delete a -> DBUpdate err a
mkOpaDelete a = DBTx $ liftF (OpaDelete a id)
