-- | Validate native queries against postgres-like flavors.
module Hasura.Backends.Postgres.Instances.NativeQueries
  ( validateNativeQuery,
    nativeQueryToPreparedStatement,
  )
where

import Control.Exception.Lifted (finally)
import Data.Aeson (toJSON)
import Data.Bifunctor
import Data.HashMap.Strict qualified as HashMap
import Data.HashMap.Strict.InsOrd qualified as InsOrdHashMap
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import Data.Text.Extended (commaSeparated, dquoteList, toTxt)
import Data.Tuple (swap)
import Database.PG.Query qualified as PG
import Database.PostgreSQL.LibPQ qualified as PQ
import Hasura.Backends.Postgres.Execute.Types (PGSourceConfig, runPgSourceReadTx)
import Hasura.Backends.Postgres.Instances.Types ()
import Hasura.Backends.Postgres.SQL.Types (PGScalarType (..), pgScalarTypeToText)
import Hasura.Base.Error
import Hasura.LogicalModel.Cache (LogicalModelInfo (..))
import Hasura.LogicalModel.Common (columnsFromFields)
import Hasura.NativeQuery.InterpolatedQuery (trimQueryEnd)
import Hasura.NativeQuery.Metadata
  ( ArgumentName,
    InterpolatedItem (..),
    InterpolatedQuery (..),
    NativeQueryMetadata (..),
  )
import Hasura.NativeQuery.Types (NullableScalarType (nstType))
import Hasura.NativeQuery.Validation (DisableNativeQueryValidation (..), validateArgumentDeclaration)
import Hasura.Prelude
import Hasura.RQL.Types.BackendType

-- | Prepare a native query query against a postgres-like database to validate it.
validateNativeQuery ::
  forall m pgKind.
  (MonadIO m, MonadError QErr m) =>
  InsOrdHashMap.InsOrdHashMap PGScalarType PQ.Oid ->
  DisableNativeQueryValidation ->
  PGSourceConfig ->
  LogicalModelInfo ('Postgres pgKind) ->
  NativeQueryMetadata ('Postgres pgKind) ->
  m (InterpolatedQuery ArgumentName)
validateNativeQuery pgTypeOidMapping disableNativeQueryValidation sourceConfig logicalModel nq = do
  validateArgumentDeclaration nq
  let nqmCode = trimQueryEnd (_nqmCode nq)
      model = nq {_nqmCode = nqmCode}
  case disableNativeQueryValidation of
    NeverValidateNativeQueries ->
      pure ()
    AlwaysValidateNativeQueries -> do
      (prepname, preparedQuery) <- nativeQueryToPreparedStatement logicalModel model
      description <- runCheck prepname (PG.fromText preparedQuery)
      let returnColumns = bimap toTxt nstType <$> InsOrdHashMap.toList (columnsFromFields $ _lmiFields logicalModel)
      for_ (toList returnColumns) (matchTypes description)

  pure nqmCode
  where
    -- Run stuff against the database.
    --
    -- Prepares this query as a statement, fetches the prepared description,
    -- and deallocates the statement.
    runCheck :: Text -> PG.Query -> m (PG.PreparedDescription PQ.Oid)
    runCheck prepname stmt =
      liftEither
        =<< liftIO
          ( runPgSourceReadTx
              sourceConfig
              ( ( do
                    -- prepare statement
                    PG.rawQE @_ @()
                      ( \e ->
                          (err400 ValidationFailed "Failed to validate query")
                            { qeInternal = Just $ ExtraInternal $ toJSON e
                            }
                      )
                      stmt
                      []
                      False
                    -- extract description
                    PG.describePreparedStatement
                      ( \e ->
                          (err400 ValidationFailed "Failed to validate query")
                            { qeInternal = Just $ ExtraInternal $ toJSON e
                            }
                      )
                      (Text.encodeUtf8 prepname)
                )
                  `finally` do
                    -- deallocate prepare statement
                    PG.rawQE @_ @()
                      ( \e ->
                          (err400 ValidationFailed "Failed to validate query")
                            { qeInternal = Just $ ExtraInternal $ toJSON e
                            }
                      )
                      (PG.fromText $ "DEALLOCATE PREPARE " <> prepname)
                      []
                      False
              )
          )

    -- Look for the type for a particular column in the prepared statement description
    --   and compare them.
    --   fail if not found, try to provide a good error message if you can.
    matchTypes :: PG.PreparedDescription PQ.Oid -> (Text, PGScalarType) -> m ()
    matchTypes description (name, expectedType) =
      case lookup (Just (Text.encodeUtf8 name)) (PG.pd_fname_ftype description) of
        Nothing ->
          throwError
            (err400 ValidationFailed "Failed to validate query")
              { qeInternal =
                  Just
                    $ ExtraInternal
                    $ toJSON @Text
                    $ "Column named '"
                    <> toTxt name
                    <> "' is not returned from the query."
              }
        Just actualOid
          | Just expectedOid <- InsOrdHashMap.lookup expectedType pgTypeOidMapping,
            expectedOid /= actualOid ->
              throwError
                (err400 ValidationFailed "Failed to validate query")
                  { qeInternal =
                      Just
                        $ ExtraInternal
                        $ toJSON @Text
                        $ Text.unwords
                        $ [ "Return column '" <> name <> "' has a type mismatch.",
                            "The expected type is '" <> toTxt expectedType <> "',"
                          ]
                        <> case Map.lookup actualOid (invertPgTypeOidMap pgTypeOidMapping) of
                          Just t ->
                            ["but the actual type is '" <> toTxt t <> "'."]
                          Nothing ->
                            [ "and has the " <> tshow expectedOid <> ",",
                              "but the actual type has the " <> tshow actualOid <> "."
                            ]
                  }
        Just {} -> pure ()

-- | Invert the type/oid mapping.
invertPgTypeOidMap :: InsOrdHashMap PGScalarType PQ.Oid -> Map PQ.Oid PGScalarType
invertPgTypeOidMap = Map.fromList . map swap . InsOrdHashMap.toList

---------------------------------------

-- | The environment and fresh-name generator used by 'renameIQ'.
data RenamingState = RenamingState
  { rsNextFree :: Int,
    rsBoundVars :: Map ArgumentName Int
  }

-- | 'Rename' an 'InterpolatedQuery' expression with 'ArgumentName' variables
-- into one which uses ordinal arguments instead of named arguments, suitable
-- for a prepared query.
renameIQ ::
  InterpolatedQuery ArgumentName ->
  ( InterpolatedQuery Int,
    Map Int ArgumentName
  )
renameIQ = runRenaming . fmap InterpolatedQuery . mapM renameII . getInterpolatedQuery
  where
    runRenaming :: forall a. State RenamingState a -> (a, Map Int ArgumentName)
    runRenaming action =
      let (res, st) = runState action (RenamingState 1 mempty)
       in (res, inverseMap $ rsBoundVars st)

    drawFree :: State RenamingState Int
    drawFree = do
      i <- gets rsNextFree
      modify (\s -> s {rsNextFree = i + 1})
      return i

    -- Rename a variable, assigning a fresh argument index when encounting new
    -- variables and reusing the previously assigned indices when encountering a
    -- previously treated variable accordingly.
    renameII ::
      InterpolatedItem ArgumentName ->
      State RenamingState (InterpolatedItem Int)
    renameII = traverse \v -> do
      env <- gets rsBoundVars
      (Map.lookup v env)
        `onNothing` ( do
                        i <- drawFree
                        modify \s -> s {rsBoundVars = Map.insert v i (rsBoundVars s)}
                        return i
                    )

    -- When renaming from the named representation to the ordinal representation
    -- it is convenient for the variable renaming environment to be keyed by the
    -- names.
    -- When subsequently rendering the prepared statement definition however, it
    -- is more convenient to inspect the environment by index.
    -- Therefore we invert the map as part of renaming.
    inverseMap :: (Ord b) => Map a b -> Map b a
    inverseMap = Map.fromList . map swap . Map.toList

-- | Pretty print an interpolated query with numbered parameters.
renderIQ :: InterpolatedQuery Int -> Text
renderIQ (InterpolatedQuery items) = foldMap printItem items
  where
    printItem :: InterpolatedItem Int -> Text
    printItem (IIText t) = t
    printItem (IIVariable i) = "$" <> tshow i

-----------------------------------------

-- | Convert a native query to a prepared statement to be validate.
--
-- Used by 'validateNativeQuery'. Exported for testing.
nativeQueryToPreparedStatement ::
  forall m pgKind.
  (MonadError QErr m) =>
  LogicalModelInfo ('Postgres pgKind) ->
  NativeQueryMetadata ('Postgres pgKind) ->
  m (Text, Text)
nativeQueryToPreparedStatement logicalModel model = do
  let (preparedIQ, argumentMapping) = renameIQ $ _nqmCode model
      logimoCode :: Text
      logimoCode = renderIQ preparedIQ
      prepname = "_logimo_vali_"

      occurringArguments, declaredArguments, undeclaredArguments :: Set ArgumentName
      occurringArguments = Set.fromList (Map.elems argumentMapping)
      declaredArguments = Set.fromList $ HashMap.keys (_nqmArguments model)
      undeclaredArguments = occurringArguments `Set.difference` declaredArguments

      argumentTypes :: Map Int PGScalarType
      argumentTypes = nstType <$> Map.fromList (HashMap.toList $ _nqmArguments model) `Map.compose` argumentMapping

      argumentSignature
        | argumentTypes /= mempty = "(" <> commaSeparated (pgScalarTypeToText <$> Map.elems argumentTypes) <> ")"
        | otherwise = ""

      returnedColumnNames :: Text
      returnedColumnNames =
        dquoteList $ InsOrdHashMap.keys (columnsFromFields $ _lmiFields logicalModel)

      wrapInCTE :: Text -> Text
      wrapInCTE query =
        Text.intercalate
          "\n"
          [ "WITH " <> ctename <> " AS (",
            query,
            ")",
            "SELECT " <> returnedColumnNames,
            "FROM " <> ctename
          ]
        where
          ctename = "_cte" <> prepname

      preparedQuery = "PREPARE " <> prepname <> argumentSignature <> " AS " <> wrapInCTE logimoCode

  when (Set.empty /= undeclaredArguments)
    $ throwError
    $ err400 ValidationFailed
    $ "Undeclared arguments: "
    <> commaSeparated (map tshow $ Set.toList undeclaredArguments)

  pure (prepname, preparedQuery)
