{-# LANGUAGE NamedFieldPuns #-}

module Festung.Vault.Persistence
    ( openVault
    , openVault'
    , closeVault
    , Error(..)
    , VaultHandle
    , QueryResult(..)
    , executeQuery
    , executeParameterizedQuery
    , Value(..)
    , Header
    , ColumnName
    , ColumnType
    , Password
    , VaultParameters(..)
    ) where


import qualified Codec.Binary.UTF8.String as UTF8
import           Control.Exception.Base (assert)
import           Control.Monad
import           Control.Monad.Catch (finally)
import           Control.Monad.Trans.Either (runEitherT, left, right, EitherT(..))
import           Control.Monad.Error.Class (throwError, catchError, MonadError)
import           Control.Monad.Trans (liftIO, lift)
import           Data.Word
import           Data.Either as E
import           Foreign.C
import           Foreign.Ptr
import           Foreign.ForeignPtr
import           Foreign.Marshal.Alloc
import           Foreign.Marshal.Array
import qualified Foreign.Concurrent as Conc
import           Foreign.Storable
import           Data.Functor ((<$>), void)
import           Data.Int (Int64)
import           Database.SQLCipher.Base
import           Database.SQLCipher.Types
import           Text.Printf (printf)

import           Festung.Utils (hoistMEither, whenJust)


{- This module basically interracts with raw C calls to SQLCipher. This
 - belongs in haskell-sqlcipher, however I (Antoine) wasn't sure what should go
 - there, and how the Haskell API should be structured.
 -
 - For this reason, I decided to put these C wrapping functions here, and see
 - where it goes. The plan is to upstream this in a near future™.
 -
 - A lot of these functions have been inspired by haskell-sqlcipher (which is based
 - on Database.SQLite)
 -}

newtype VaultHandle = VaultHandle (ForeignPtr ())

instance Show VaultHandle where
    show _ = "VaultHandle"

data Error = CouldNotOpenVault        Status !String
           | MultipleStatements
           | WrongParametrization     { got :: Int, expected :: Int }

           | NotSupportedError        String
           | InternalError            String
           | IntegrityError           Status !String
           | OperationalError         Status !String
           | DatabaseError            Status !String
           | DataError                Status !String
           | ProgrammingError         Status !String
           deriving (Show)

data Value = StringValue String
           | IntValue    Int64
           | FloatValue  Double
           | BlobValue   [Word8]
           | NullValue
           deriving (Show, Eq)

type Row = [Value]

-- FIXME(Antoine): Later on this should be an Enum. (Nothing should be "DYNAMIC")
type ColumnType = Maybe String
type ColumnName = String
type Header = (ColumnName, ColumnType)

data QueryResult = QueryResult
    { rows        :: [Row]
    , lastRowId   :: Int
    , rowsChanged :: Maybe Int
    , headers     :: [Header]
    } deriving (Show)


-- | Binary password to a vault (this password will be escaped)
type Password = [Word8]

newtype VaultParameters = VaultParameters { kdfIter :: Maybe Integer }
  deriving (Eq, Show)


head' :: [a] -> Maybe a
head' []    = Nothing
head' (h:_) = Just h


valueString :: Value -> String
valueString (StringValue s) = s
valueString (IntValue    i) = show i
valueString (FloatValue  f) = show f
valueString (BlobValue   _) = undefined
valueString NullValue       = "NULL"


-- |@'encodeCString' encode a Haskell string for it to be used in C.
--
-- C functions only manipulates bytes because they use @char *@, this encodes
-- the string in UTF8, and returns a list of byte.
--
-- This utility function is only for internal use in Festung.Vault.Persistence.
encodeCString :: String -> [CChar]
encodeCString = map fromIntegral <$> UTF8.encode

-- |@'decodeCString' reverses @'encodeCString'
--
-- @decodeCString . encodeCString == id@
--
-- This utility function is only for internal use in Festung.Vault.Persistence.
decodeCString :: [CChar] -> String
decodeCString = UTF8.decode <$> map fromIntegral


withUTF8CStringLen :: String -> (CStringLen -> IO a) -> IO a
withUTF8CStringLen str action =
    let cStr = encodeCString str
     in withArray0 0 cStr $ \ptr -> action (ptr, length cStr)


-- |Convert a Haskell String to a @char *@ string, the time to execute an IO action.
withUTF8CString :: String -> (CString -> IO a) -> IO a
withUTF8CString str = 
    let cStr = encodeCString str
     in withArray0 0 cStr


-- |peek @char *@ string into a Haskell String.
peekUTF8CString :: CString -> IO String
peekUTF8CString cStr = decodeCString <$> peekArray0 0 cStr


-- |Get the error string from a errorneous status code
sqlCipherErrorString :: Status -> IO String
sqlCipherErrorString = peekCString . sqlite3_errstr


sqlCipherErrorMessage :: SQLite -> IO String
sqlCipherErrorMessage db = peekUTF8CString =<< sqlite3_errmsg db


-- |Get the @'Error' object for a status code
sqlCipherError :: Status -> (Status -> String -> Error) -> IO Error
sqlCipherError status constructor = constructor status <$> sqlCipherErrorString status


sqlCipherErrorFromDb :: Status -> SQLite -> (Status -> String -> Error) -> IO Error
sqlCipherErrorFromDb status db constructor = constructor status <$> sqlCipherErrorMessage db


toError :: Status -> SQLite -> IO Error
toError status db =
    let constructor = case status of
            _ | status `elem` 
                    [ sQLITE_ERROR
                    , sQLITE_PERM
                    , sQLITE_ABORT
                    , sQLITE_BUSY
                    , sQLITE_LOCKED
                    , sQLITE_READONLY
                    , sQLITE_INTERRUPT
                    , sQLITE_IOERR
                    , sQLITE_FULL
                    , sQLITE_CANTOPEN
                    , sQLITE_PROTOCOL
                    , sQLITE_EMPTY
                    , sQLITE_SCHEMA
                    ]                    -> OperationalError
              | status `elem` 
                    [ sQLITE_CONSTRAINT
                    , sQLITE_MISMATCH
                    ]                    -> IntegrityError
              | status == sQLITE_CORRUPT -> DatabaseError
              | status == sQLITE_TOOBIG  -> DataError
              | status == sQLITE_MISUSE  -> ProgrammingError
              | otherwise                -> DatabaseError
     in sqlCipherErrorFromDb status db constructor


-- |Pass @[Word8]@ to a c function with the signature @int func(void *data, int dataLen)@
withData :: [Word8] -> (Int -> Ptr () -> IO a) -> IO a
withData data_ action =
    let cData = map fromIntegral data_ :: [CChar] 
     in withArrayLen cData $ \len dataPtr -> action len (castPtr dataPtr)


-- |Read @void *data@
peekData :: Ptr () -> Int -> IO [Word8]
peekData ptr len =
    let ptr' = castPtr ptr :: Ptr CChar
     in map fromIntegral <$> peekArray len ptr'


withPrim :: VaultHandle -> (SQLite -> IO a) -> IO a
withPrim (VaultHandle ptr) action = withForeignPtr ptr (action . SQLite)


liftMEither :: Monad m => m (Either e a) -> EitherT e m a
liftMEither m = lift m >>= E.either left right


-- |A computation to run, in case an error occurs.
--
-- This functions doesn't restore the state, it keeps the errored state.
onError :: (MonadError e m) => m () -> m a -> m a
onError errorAction action = catchError action $ \err -> errorAction >> throwError err


-- |Just runs @sqlite3_open@.
--
-- This functions is for internal use only.
openSQLiteDB :: FilePath -> EitherT Error IO VaultHandle
openSQLiteDB filename =
    let newVaultHandle h@(SQLite ptr) = VaultHandle <$> Conc.newForeignPtr ptr sqlite3_close'
            where sqlite3_close' = void $ sqlite3_close h
     in liftMEither . alloca $ \dbPtr -> do
            status <- withUTF8CString filename $ flip sqlite3_open dbPtr
            if status == sQLITE_OK
                then Right <$> (newVaultHandle =<< peek dbPtr)
                else Left  <$> sqlCipherError status OperationalError


-- |@'usePassword' uses SQLCipher's @sqlite3_key@ special function to send the database password.
--
-- This utility function is only for internal use in Festung.Vault.Persistence,
-- more specificaly it should only be used in @'openVault'.
usePassword :: VaultHandle -> Password -> EitherT Error IO ()
usePassword handle password = liftMEither $ withPrim handle $ \db -> do
    status <- withData password $ \len passwordData ->
        let cLen = fromIntegral len 
         in sqlite3_key db passwordData cLen
    if status == sQLITE_OK
       then return $ Right ()
       else Left <$> toError status db


-- |@'setKdfIter' run the @PRAGMA kdf_iter@
--
-- This utility function is for internal use only
setKdfIter :: VaultHandle -> Integer -> EitherT Error IO ()
setKdfIter handle kdfIter =
    let query = printf "PRAGMA kdf_iter = '%d'" kdfIter
     in void $ hoistMEither $ executeQuery handle query


ensureIntegrity :: VaultHandle -> EitherT Error IO ()
ensureIntegrity handle = do
    result <- liftMEither $ executeQuery handle "PRAGMA quick_check(1)"
    let integrity_check = head . head . rows $ isIntegrityCheck result

    unless (isOk integrity_check) $ left $ IntegrityError sQLITE_MISMATCH (valueString integrity_check)
        where isOk (StringValue v) | v == "ok" = True
                                   | otherwise = False
              isIntegrityCheck r@QueryResult { headers } =
                  assert (fmap fst (head' headers) == Just "integrity_check") r


-- |@'openVault' opens the vault and returns an handle
openVault :: FilePath -> Password -> IO (Either Error VaultHandle)
openVault filename password = openVault' filename password noParams
    where noParams = VaultParameters { kdfIter = Nothing }


openVault' :: FilePath -> Password -> VaultParameters -> IO (Either Error VaultHandle)
openVault' filename password VaultParameters{ kdfIter } = runEitherT $ do
    handle <- openSQLiteDB filename
    onError (liftIO $ closeVault handle) $ do
        usePassword handle password
        whenJust kdfIter $ setKdfIter handle
        ensureIntegrity handle
        return handle


-- |@'withPreparedStatement' allocate a prepared statement
--
-- This function will finalize and dealocate the statement.
--
-- This helper function is inteded
withPreparedStatement :: SQLite -> String -> (SQLiteStmt -> IO a) -> IO (Either Error a)
withPreparedStatement db query action =
    alloca $ \ppStmt ->
    alloca $ \pzTail ->
    withArrayLen cData $ \nByte zSql -> do
        let nByte' = fromIntegral nByte
        status <- sqlite3_prepare db zSql nByte' ppStmt pzTail
        pStmt  <- peek ppStmt
        runEitherT $ do
            -- FIXME(Antoine): Too many liftIO
            when (status /= sQLITE_OK) $
                left =<< liftIO (toError status db)

            when (isNullStmt pStmt) $
                left $ InternalError "Prepared statement was null."

            -- FIXME(Antoine): The double use of finalize is not great.
            let finalize = sqlite3_finalize pStmt

            -- TODO: Multiple statements
            -- when (pzTail /= nullPtr) (liftIO finalize >> left MultipleStatements)
            liftIO $ finally (action pStmt) finalize
    where cData = encodeCString query


handleBindStatus :: Status -> SQLite -> EitherT Error IO ()
handleBindStatus status db =
    if status == sQLITE_OK
        then right ()
        else left =<< liftIO (toError status db)


-- |@'bindParameter' binds one paramater to a statmeent.
--
-- This function is indented for internal use only.
bindParameter :: SQLite -> SQLiteStmt -> CInt -> Value -> EitherT Error IO ()
bindParameter db stmt idx (StringValue str) = do
    status <- liftIO $ withUTF8CStringLen str $ \(cStr, cStrLen) ->
        let cStrLen' = fromIntegral cStrLen
         in sqlite3_bind_text64 stmt idx cStr cStrLen' sqlite3_transient_destructor sQLITE_UTF8
    handleBindStatus status db

bindParameter db stmt idx (IntValue int) = do
    status <- liftIO $ sqlite3_bind_int64 stmt idx (fromIntegral int)
    handleBindStatus status db

bindParameter db stmt idx (FloatValue float) = do
    status <- liftIO $ sqlite3_bind_double stmt idx float
    handleBindStatus status db

bindParameter _db _stmt _idx (BlobValue _data) = left $ NotSupportedError "Binding blob is not supported yet."
bindParameter db stmt idx NullValue = do
    status <- liftIO $ sqlite3_bind_null stmt idx
    handleBindStatus status db


fetchColumn :: SQLiteStmt -> Int -> EitherT Error IO Value
fetchColumn stmt pos = do
    let pos' = fromIntegral pos -- XXX(Antoine): This is copy pasted everywhere
    ct <- liftIO $ sqlite3_column_type stmt pos'
    case ct of
        _ | ct == sQLITE_INTEGER ->
              liftIO $ IntValue . fromIntegral <$> sqlite3_column_int64 stmt pos'
          | ct == sQLITE_FLOAT   ->
              liftIO $ FloatValue <$> sqlite3_column_double stmt pos'
          | ct == sQLITE_NULL    -> return NullValue
          | ct == sQLITE_TEXT    -> liftIO $ do
              -- TODO: Use sqlite3_column_bytes AFTER
              -- TODO: Handle null pointer sqilte3_column_text
              cStr <- sqlite3_column_text stmt pos'
              StringValue <$> peekUTF8CString cStr
          | ct == sQLITE_BLOB    -> liftIO $ do
              -- TODO: Handle null pointer
              blob    <- sqlite3_column_blob  stmt pos'
              blobLen <- sqlite3_column_bytes stmt pos'
              BlobValue <$> peekData blob (fromIntegral blobLen)
          | otherwise            -> left $ ProgrammingError sQLITE_MISUSE $ "Unknown column type: " ++ show ct


fetchColumnName :: SQLiteStmt -> Int -> EitherT Error IO ColumnName
fetchColumnName stmt pos = do
    let pos' = fromIntegral pos -- XXX(Antoine): This is copy pasted everywhere
    columnName <- liftIO $ sqlite3_column_name stmt pos'
    when (columnName == nullPtr) $ left $ InternalError "Couldn't fetch column name."
    liftIO $ peekUTF8CString columnName


fetchColumnType :: SQLiteStmt -> Int -> EitherT Error IO ColumnType
fetchColumnType stmt pos = do
    let pos' = fromIntegral pos -- XXX(Antoine): This is copy pasted everywhere
    columnName <- liftIO $ sqlite3_column_decltype stmt pos'
    if columnName == nullPtr
        then return Nothing
        else liftIO $ Just <$> peekUTF8CString columnName


fetchColumnHeader :: SQLiteStmt -> Int -> EitherT Error IO Header
fetchColumnHeader stmt pos =
    pure (,)                  <*>
    fetchColumnName stmt pos  <*>
    fetchColumnType stmt pos


mapColumns :: (SQLiteStmt -> Int -> EitherT Error IO a) -> SQLiteStmt -> EitherT Error IO [a]
mapColumns f stmt = do
    nCol <- liftIO $ fromIntegral <$> sqlite3_column_count stmt
    mapM (f stmt) $ take nCol [0..]


fetchRow :: SQLiteStmt -> EitherT Error IO Row
fetchRow = mapColumns fetchColumn


fetchHeaders :: SQLiteStmt -> EitherT Error IO [Header]
fetchHeaders = mapColumns fetchColumnHeader


fetchResults :: SQLite -> SQLiteStmt -> EitherT Error IO [Row]
fetchResults db stmt =
    let go acc = do
            status <- liftIO $ sqlite3_step stmt
            case status of
                _ | status == sQLITE_DONE -> return $ reverse acc
                  | status == sQLITE_ROW  -> do
                        row <- fetchRow stmt
                        go (row:acc)
                  | otherwise             -> left =<< liftIO (toError status db)
     in go []


bindManyParameters :: SQLite -> SQLiteStmt -> [Value] -> EitherT Error IO ()
bindManyParameters db stmt params = do
    expects <- liftIO $ fromIntegral <$> sqlite3_bind_parameter_count stmt

    let got = length params
    unless (expects == got) $ left $ ProgrammingError sQLITE_MISUSE (
        "Wrong param count. Got: " ++ show got ++ " Expected: " ++ show expects)

    let indices = take (length params) [1..]
        binders = map (bindParameter db stmt) indices

    zipWithM_ ($) binders params


executeQuery :: VaultHandle -> String -> IO (Either Error QueryResult)
{-# INLINE executeQuery #-}
executeQuery handle query = executeParameterizedQuery handle query []


stmtReadOnly :: SQLiteStmt -> IO Bool
stmtReadOnly stmt = (/= 0) <$> sqlite3_stmt_readonly stmt


rowsAffected :: SQLite -> SQLiteStmt -> IO (Maybe Int)
rowsAffected db stmt = do
    readOnly <- stmtReadOnly stmt
    if not readOnly
       then Just . fromIntegral <$> sqlite3_changes db
       else return Nothing


-- |@'fetchResults' fetch all results from a prepared


-- |@'executeParameterizedQuery' executes a SQL query, with *un-named* parameters.
--
-- @executeParameterizedQuery "SELECT * FROM table WHERE value = ?" [NullValue]
--
executeParameterizedQuery :: VaultHandle -> String -> [Value] -> IO (Either Error QueryResult)
executeParameterizedQuery handle query params = withPrim handle $ \db ->
    fmap join <$> withPreparedStatement db query $ \pStmt -> runEitherT $ do
        bindManyParameters db pStmt params
        rows        <- fetchResults db pStmt
        headers     <- fetchHeaders pStmt
        rowsChanged <- liftIO $ rowsAffected db pStmt
        lastRowId   <- liftIO $ lastInsertRowId db
        return QueryResult { rows = rows
                           , headers = headers
                           , lastRowId  = lastRowId
                           , rowsChanged = rowsChanged
                           }


lastInsertRowId :: SQLite -> IO Int
lastInsertRowId = fmap fromIntegral . sqlite3_last_insert_rowid


-- |@'closeVault' closes a vault handle
closeVault :: VaultHandle -> IO ()
closeVault (VaultHandle ptr) = finalizeForeignPtr ptr