content: {-# LANGUAGE TypeFamilies, FlexibleInstances, ScopedTypeVariables #-}
  {-# LANGUAGE MultiParamTypeClasses, FlexibleContexts #-}
  {-# LANGUAGE UndecidableInstances #-}
  -- | Building and executing prepared statements.
  module Database.Selda.Prepared (Preparable, Prepare, prepared) where
  import Database.Selda.Backend.Internal
      ( Lit(LCustom),
        SqlType(sqlType),
        SqlTypeRep,
        Param(..),
        MonadSelda(Backend, withConnection),
        SeldaBackend(ppConfig, runPrepared, backendId, prepareStmt),
        SeldaConnection(connBackend, connStmts),
        SeldaStmt(SeldaStmt, stmtHandle, stmtParams, stmtText),
        StmtID(..),
        BackendID,
        freshStmtId,
        withBackend )
  import Database.Selda.Column ( Exp(Lit), Col(..) )
  import Database.Selda.Compile
      ( Result, Res, compileWith, buildResult )
  import Database.Selda.Query.Type ( Query )
  import Database.Selda.SQL (param, paramType)
  import Control.Exception ( Exception, try, throw, mask )
  import Control.Monad.IO.Class ( MonadIO(liftIO) )
  import qualified Data.IntMap as M
  import Data.IORef
      ( IORef, atomicModifyIORef', newIORef, readIORef, writeIORef )
  import Data.Proxy ( Proxy(..) )
  import Data.Text (Text)
  import Data.Typeable ( Typeable )
  import System.IO.Unsafe ( unsafePerformIO )
  
  data Placeholder = Placeholder Int
    deriving Show
  instance Exception Placeholder
  
  -- | Index of first argument parameter to a query.
  firstParamIx :: Int
  firstParamIx = 0
  
  -- | Result type of a monadic computation.
  type family ResultT f where
    ResultT (a -> b) = ResultT b
    ResultT (m a)    = a
  
  type family Equiv q f where
    Equiv (Col s a -> q) (a -> f) = Equiv q f
    Equiv (Query s a)    (m [b])  = (Res a ~ b, Backend m ~ s)
  
  type CompResult = (Text, [Either Int Param], [SqlTypeRep])
  
  class Preparable q where
    -- | Prepare the query and parameter list.
    mkQuery :: MonadSelda m
            => Int -- ^ Next argument index.
            -> q   -- ^ The query.
            -> [SqlTypeRep] -- ^ The list of param types so far.
            -> m CompResult
  
  -- | Some parameterized query @q@ that can be prepared into a function @f@
  --   in some @MonadSelda@.
  class Prepare q f where
    -- | Build the function that prepares and execute the query.
    mkFun :: Preparable q
          => IORef (Maybe (BackendID, CompResult))
          -> StmtID
          -> q
          -> [Param]
          -> f
  
  instance (SqlType a, Prepare q b) => Prepare q (a -> b) where
    mkFun ref sid qry ps x = mkFun ref sid qry (param x : ps)
  
  instance (Typeable a, MonadSelda m, a ~ Res (ResultT q), Result (ResultT q)) =>
           Prepare q (m [a]) where
    -- This function uses read/writeIORef instead of atomicModifyIORef.
    -- For once, this is actually safe: the IORef points to a single compiled
    -- statement, so the only consequence of a race between the read and the write
    -- is that the statement gets compiled (note: NOT prepared) twice.
    mkFun ref (StmtID sid) qry arguments = withConnection $ \conn -> do
      let backend = connBackend conn
          args = reverse arguments
      stmts <- liftIO $ readIORef (connStmts conn)
      case M.lookup sid stmts of
        Just stm -> do
          -- Statement already prepared for this connection; just execute it.
          liftIO $ runQuery conn stm args
        _ -> do
          -- Statement wasn't prepared for this connection; check if it was at
          -- least previously compiled for this backend.
          compiled <- liftIO $ readIORef ref
          (q, params, reps) <- case compiled of
            Just (bid, comp) | bid == backendId backend -> do
              return comp
            _ -> do
              comp <- mkQuery firstParamIx qry []
              liftIO $ writeIORef ref (Just (backendId backend, comp))
              return comp
  
          -- Prepare and execute
          liftIO $ mask $ \restore -> do
            hdl <- prepareStmt backend (StmtID sid) reps q
            let stm = SeldaStmt
                  { stmtHandle = hdl
                  , stmtParams = params
                  , stmtText = q
                  }
            atomicModifyIORef' (connStmts conn) $ \m -> (M.insert sid stm m, ())
            restore $ runQuery conn stm args
      where
        runQuery conn stm args = do
          let ps = replaceParams (stmtParams stm) args
              hdl = stmtHandle stm
          res <- runPrepared (connBackend conn) hdl ps
          return $ map (buildResult (Proxy :: Proxy (ResultT q))) (snd res)
  
  instance (SqlType a, Preparable b) => Preparable (Col s a -> b) where
    mkQuery n f ts = mkQuery (n+1) (f x) (t : ts)
      where
        t = sqlType (Proxy :: Proxy a)
        x = One $ Lit $ LCustom t (throw (Placeholder n) :: Lit a)
  
  instance Result a => Preparable (Query s a) where
    mkQuery _ q types = withBackend $ \b -> do
      case compileWith (ppConfig b) q of
        (q', ps) -> do
          (ps', types') <- liftIO $ inspectParams (reverse types) ps
          return (q', ps', types')
  
  -- | Create a prepared Selda function. A prepared function has zero or more
  --   arguments, and will get compiled into a prepared statement by the first
  --   backend to execute it. Any subsequent calls to the function for the duration
  --   of the connection to the database will reuse the prepared statement.
  --
  --   Preparable functions are of the form
  --   @(SqlType a, SqlType b, ...) => Col s a -> Col s b -> ... -> Query s r@.
  --   The resulting prepared function will be of the form
  --   @MonadSelda m => a -> b -> ... -> m [Res r]@.
  --   Note, however, that when using @prepared@, you must give a concrete type
  --   for @m@ due to how Haskell's type class resolution works.
  --
  --   Prepared functions rely on memoization for just-in-time preparation and
  --   caching. This means that if GHC accidentally inlines your prepared function,
  --   it may get prepared twice.
  --   While this does not affect the correctness of your program, and is
  --   fairly unlikely to happen, if you want to be absolutely sure that your
  --   queries aren't re-prepared more than absolutely necessary,
  --   consider adding a @NOINLINE@ annotation to each prepared function.
  --
  --   Note that when using a constrained backend type variable (i.e.
  --   @foo :: Bar b => SeldaM b [Int]@), optimizations must be enabled for
  --   prepared statements to be effective.
  --
  --   A usage example:
  --
  -- > persons :: Table (Text, Int)
  -- > (persons, name :*: age) = tableWithSelectors "ages" [name :- primary]
  -- >
  -- > {-# NOINLINE ageOf #-}
  -- > ageOf :: Text -> SeldaM [Int]
  -- > ageOf = prepared $ \n -> do
  -- >   person <- select ages
  -- >   restrict $ (person!name .== n)
  -- >   return age
  {-# NOINLINE prepared #-}
  prepared :: (Preparable q, Prepare q f, Equiv q f) => q -> f
  prepared q = unsafePerformIO $ do
    ref <- newIORef Nothing
    sid <- freshStmtId
    return $ mkFun ref sid q []
  
  -- | Replace every indexed parameter with the corresponding provided parameter.
  --   Keep all non-indexed parameters in place.
  replaceParams :: [Either Int Param] -> [Param] -> [Param]
  replaceParams params = map fromRight . go firstParamIx params
    where
      go n ps (x:xs) = go (n+1) (map (subst n x) ps) xs
      go _ ps _      = ps
  
      subst n x (Left n') | n == n' = Right x
      subst _ _ old                 = old
  
      fromRight (Right x) = x
      fromRight _         = error "BUG: query parameter not substituted!"
  
  -- | Inspect a list of parameters, denoting each parameter with either a
  --   placeholder index or a literal parameter.
  inspectParams :: [SqlTypeRep] -> [Param] -> IO ([Either Int Param], [SqlTypeRep])
  inspectParams ts (x:xs) = do
    res <- try $ pure $! forceParam x
    let (x', t) = case res of
          Right p               -> (Right p, paramType p)
          Left (Placeholder ix) -> (Left ix, ts !! ix)
    (xs', ts') <- inspectParams ts xs
    return (x' : xs', t : ts')
  inspectParams _ [] = do
    return ([], [])
  
  -- | Force a parameter deep enough to determine whether it is a placeholder.
  forceParam :: Param -> Param
  forceParam p@(Param (LCustom _ x)) | x `seq` True = p
  forceParam p                                      = p
  