{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}

module Database.PostgreSQL.Opium.Connection
  ( Connection
  , ConnectionError
  , unsafeWithRawConnection
  , withRawConnection
  , connect
  , close
  ) where

import Control.Concurrent.MVar (MVar, newMVar, modifyMVar_, withMVar)
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Database.PostgreSQL.LibPQ (ConnStatus (..))
import GHC.Stack (HasCallStack)

import qualified Data.Text.Encoding as Encoding
import qualified Database.PostgreSQL.LibPQ as LibPQ

newtype Connection = Connection
  { rawConnection :: MVar (Maybe LibPQ.Connection)
  }

newtype ConnectionError = ConnectionError Text
  deriving (Show)

withRawConnection
  :: HasCallStack
  => (Maybe LibPQ.Connection -> IO a)
  -> Connection
  -> IO a
withRawConnection f conn = withMVar conn.rawConnection f

unsafeWithRawConnection
  :: HasCallStack
  => (LibPQ.Connection -> IO a)
  -> Connection
  -> IO a
unsafeWithRawConnection f = withRawConnection $ \case
  Nothing -> error "raw connection is missing! perhaps the connection was already closed."
  Just rawConn -> f rawConn

connect :: Text -> IO (Either ConnectionError Connection)
connect connectionString = do
  -- Appending the client_encoding setting overrides any previous setting in the connection string.
  -- We set the client encoding here to make sure we can use it below for decoding connection
  -- error messages.
  rawConn <- LibPQ.connectdb $ Encoding.encodeUtf8 $ connectionString <> " client_encoding=UTF8"
  status <- LibPQ.status rawConn
  if status == ConnectionOk then
    Right . Connection <$> newMVar (Just rawConn)
  else do
    rawError <- fromMaybe "" <$> LibPQ.errorMessage rawConn
    pure $ Left $ ConnectionError $ Encoding.decodeUtf8Lenient rawError

close :: Connection -> IO ()
close conn = modifyMVar_ conn.rawConnection $ \case
  Just rawConn -> do
    LibPQ.finish rawConn
    pure Nothing
  Nothing ->
    pure Nothing