{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-|
Module      : PostgREST.Query.QueryBuilder
Description : PostgREST SQL queries generating functions.

This module provides functions to consume data types that
represent database queries (e.g. ReadPlanTree, MutatePlan) and SqlFragment
to produce SqlQuery type outputs.
-}
module PostgREST.Query.QueryBuilder
  ( readPlanToQuery
  , mutatePlanToQuery
  , readPlanToCountQuery
  , callPlanToQuery
  , limitedQuery
  ) where

import qualified Data.ByteString.Char8           as BS
import qualified Hasql.DynamicStatements.Snippet as SQL

import Data.Maybe (fromJust)
import Data.Tree  (Tree (..))

import PostgREST.ApiRequest.Preferences   (PreferResolution (..))
import PostgREST.Config.PgVersion         (PgVersion, pgVersion110,
                                           pgVersion130)
import PostgREST.SchemaCache.Identifiers  (QualifiedIdentifier (..))
import PostgREST.SchemaCache.Relationship (Cardinality (..),
                                           Junction (..),
                                           Relationship (..))
import PostgREST.SchemaCache.Routine      (RoutineParam (..))

import PostgREST.ApiRequest.Types
import PostgREST.Plan.CallPlan
import PostgREST.Plan.MutatePlan
import PostgREST.Plan.ReadPlan
import PostgREST.Plan.Types
import PostgREST.Query.SqlFragment
import PostgREST.RangeQuery        (allRange)

import Protolude

readPlanToQuery :: ReadPlanTree -> SQL.Snippet
readPlanToQuery node@(Node ReadPlan{select,from=mainQi,fromAlias,where_=logicForest,order, range_=readRange, relToParent, relJoinConds, relSelect} forest) =
  "SELECT " <>
  intercalateSnippet ", " ((pgFmtSelectItem qi <$> (if null select && null forest then defSelect else select)) ++ joinsSelects) <> " " <>
  fromFrag <> " " <>
  intercalateSnippet " " joins <> " " <>
  (if null logicForest && null relJoinConds
    then mempty
    else "WHERE " <> intercalateSnippet " AND " (map (pgFmtLogicTree qi) logicForest ++ map pgFmtJoinCondition relJoinConds)) <> " " <>
  groupF qi select relSelect <> " " <>
  orderF qi order <> " " <>
  limitOffsetF readRange
  where
    fromFrag = fromF relToParent mainQi fromAlias
    qi = getQualifiedIdentifier relToParent mainQi fromAlias
    -- gets all the columns in case of an empty select, ignoring/obtaining these columns is done at the aggregation stage
    defSelect = [CoercibleSelectField (unknownField "*" []) Nothing Nothing Nothing Nothing]
    joins = getJoins node
    joinsSelects = getJoinSelects node

getJoinSelects :: ReadPlanTree -> [SQL.Snippet]
getJoinSelects (Node ReadPlan{relSelect} _) =
  mapMaybe relSelectToSnippet relSelect
  where
    relSelectToSnippet :: RelSelectField -> Maybe SQL.Snippet
    relSelectToSnippet fld =
      let aggAlias = pgFmtIdent $ rsAggAlias fld
      in
        case fld of
          JsonEmbed{rsEmptyEmbed = True} ->
            Nothing
          JsonEmbed{rsSelName, rsEmbedMode = JsonObject} ->
            Just $ "row_to_json(" <> aggAlias <> ".*)::jsonb AS " <> pgFmtIdent rsSelName
          JsonEmbed{rsSelName, rsEmbedMode = JsonArray} ->
            Just $ "COALESCE( " <> aggAlias <> "." <> aggAlias <> ", '[]') AS " <> pgFmtIdent rsSelName
          Spread{rsSpreadSel, rsAggAlias} ->
            Just $ intercalateSnippet ", " (pgFmtSpreadSelectItem rsAggAlias <$> rsSpreadSel)

getJoins :: ReadPlanTree -> [SQL.Snippet]
getJoins (Node _ []) = []
getJoins (Node ReadPlan{relSelect} forest) =
  map (\fld ->
         let alias = rsAggAlias fld
             matchingNode = fromJust $ find (\(Node ReadPlan{relAggAlias} _) -> alias == relAggAlias) forest
         in getJoin fld matchingNode
      ) relSelect

getJoin :: RelSelectField -> ReadPlanTree -> SQL.Snippet
getJoin fld node@(Node ReadPlan{relJoinType} _) =
  let
    correlatedSubquery sub al cond =
      (if relJoinType == Just JTInner then "INNER" else "LEFT") <> " JOIN LATERAL ( " <> sub <> " ) AS " <> al <> " ON " <> cond
    subquery = readPlanToQuery node
    aggAlias = pgFmtIdent $ rsAggAlias fld
  in
    case fld of
      JsonEmbed{rsEmbedMode = JsonObject} ->
        correlatedSubquery subquery aggAlias "TRUE"
      Spread{} ->
        correlatedSubquery subquery aggAlias "TRUE"
      JsonEmbed{rsEmbedMode = JsonArray} ->
        let
          subq = "SELECT json_agg(" <> aggAlias <> ")::jsonb AS " <> aggAlias <> " FROM (" <> subquery <> " ) AS " <> aggAlias
          condition = if relJoinType == Just JTInner then aggAlias <> " IS NOT NULL" else "TRUE"
        in correlatedSubquery subq aggAlias condition

mutatePlanToQuery :: MutatePlan -> SQL.Snippet
mutatePlanToQuery (Insert mainQi iCols body onConflct putConditions returnings _ applyDefaults) =
  "INSERT INTO " <> fromQi mainQi <> (if null iCols then " " else "(" <> cols <> ") ") <>
  fromJsonBodyF body iCols True False applyDefaults <>
  -- Only used for PUT
  (if null putConditions then mempty else "WHERE " <> addConfigPgrstInserted True <> " AND " <> intercalateSnippet " AND " (pgFmtLogicTree (QualifiedIdentifier mempty "pgrst_body") <$> putConditions)) <>
  (if null putConditions && mergeDups then "WHERE " <> addConfigPgrstInserted True else mempty) <>
  maybe mempty (\(oncDo, oncCols) ->
    if null oncCols then
      mempty
    else
      " ON CONFLICT(" <> intercalateSnippet ", " (pgFmtIdent <$> oncCols) <> ") " <> case oncDo of
      IgnoreDuplicates ->
        "DO NOTHING"
      MergeDuplicates  ->
        if null iCols
           then "DO NOTHING"
           else "DO UPDATE SET " <> intercalateSnippet ", " ((pgFmtIdent . cfName) <> const " = EXCLUDED." <> (pgFmtIdent . cfName) <$> iCols) <> (if null putConditions && not mergeDups then mempty else "WHERE " <> addConfigPgrstInserted False)
    ) onConflct <> " " <>
    returningF mainQi returnings
  where
    cols = intercalateSnippet ", " $ pgFmtIdent . cfName <$> iCols
    mergeDups = case onConflct of {Just (MergeDuplicates,_) -> True; _ -> False;}

-- An update without a limit is always filtered with a WHERE
mutatePlanToQuery (Update mainQi uCols body logicForest range ordts returnings applyDefaults)
  | null uCols =
    -- if there are no columns we cannot do UPDATE table SET {empty}, it'd be invalid syntax
    -- selecting an empty resultset from mainQi gives us the column names to prevent errors when using &select=
    -- the select has to be based on "returnings" to make computed overloaded functions not throw
    "SELECT " <> emptyBodyReturnedColumns <> " FROM " <> fromQi mainQi <> " WHERE false"

  | range == allRange =
    "UPDATE " <> mainTbl <> " SET " <> nonRangeCols <> " " <>
    fromJsonBodyF body uCols False False applyDefaults <>
    whereLogic <> " " <>
    returningF mainQi returnings

  | otherwise =
    "WITH " <>
    "pgrst_update_body AS (" <> fromJsonBodyF body uCols True True applyDefaults <> "), " <>
    "pgrst_affected_rows AS (" <>
      "SELECT " <> rangeIdF <> " FROM " <> mainTbl <>
      whereLogic <> " " <>
      orderF mainQi ordts <> " " <>
      limitOffsetF range <>
    ") " <>
    "UPDATE " <> mainTbl <> " SET " <> rangeCols <>
    "FROM pgrst_affected_rows " <>
    "WHERE " <> whereRangeIdF <> " " <>
    returningF mainQi returnings

  where
    whereLogic = if null logicForest then mempty else " WHERE " <> intercalateSnippet " AND " (pgFmtLogicTree mainQi <$> logicForest)
    mainTbl = fromQi mainQi
    emptyBodyReturnedColumns = if null returnings then "NULL" else intercalateSnippet ", " (pgFmtColumn (QualifiedIdentifier mempty $ qiName mainQi) <$> returnings)
    nonRangeCols = intercalateSnippet ", " (pgFmtIdent . cfName <> const " = " <> pgFmtColumn (QualifiedIdentifier mempty "pgrst_body") . cfName <$> uCols)
    rangeCols = intercalateSnippet ", " ((\col -> pgFmtIdent (cfName col) <> " = (SELECT " <> pgFmtIdent (cfName col) <> " FROM pgrst_update_body) ") <$> uCols)
    (whereRangeIdF, rangeIdF) = mutRangeF mainQi (cfName . coField <$> ordts)

mutatePlanToQuery (Delete mainQi logicForest range ordts returnings)
  | range == allRange =
    "DELETE FROM " <> fromQi mainQi <> " " <>
    whereLogic <> " " <>
    returningF mainQi returnings

  | otherwise =
    "WITH " <>
    "pgrst_affected_rows AS (" <>
      "SELECT " <> rangeIdF <> " FROM " <> fromQi mainQi <>
       whereLogic <> " " <>
      orderF mainQi ordts <> " " <>
      limitOffsetF range <>
    ") " <>
    "DELETE FROM " <> fromQi mainQi <> " " <>
    "USING pgrst_affected_rows " <>
    "WHERE " <> whereRangeIdF <> " " <>
    returningF mainQi returnings

  where
    whereLogic = if null logicForest then mempty else " WHERE " <> intercalateSnippet " AND " (pgFmtLogicTree mainQi <$> logicForest)
    (whereRangeIdF, rangeIdF) = mutRangeF mainQi (cfName . coField <$> ordts)

callPlanToQuery :: CallPlan -> PgVersion -> SQL.Snippet
callPlanToQuery (FunctionCall qi params args returnsScalar returnsSetOfScalar returnsCompositeAlias returnings) pgVer =
  "SELECT " <> (if returnsScalar || returnsSetOfScalar then "pgrst_call.pgrst_scalar" else returnedColumns) <> " " <>
  fromCall
  where
    fromCall = case params of
      OnePosParam prm -> "FROM " <> callIt (singleParameter args $ encodeUtf8 $ ppType prm)
      KeyParams []    -> "FROM " <> callIt mempty
      KeyParams prms  -> fromJsonBodyF args ((\p -> CoercibleField (ppName p) mempty False (ppTypeMaxLength p) Nothing Nothing) <$> prms) False True False <> ", " <>
                         "LATERAL " <> callIt (fmtParams prms)

    callIt :: SQL.Snippet -> SQL.Snippet
    callIt argument | pgVer < pgVersion130 && pgVer >= pgVersion110 && returnsCompositeAlias = "(SELECT (" <> fromQi qi <> "(" <> argument <> ")).*) pgrst_call"
                    | returnsScalar || returnsSetOfScalar                                    = "(SELECT " <> fromQi qi <> "(" <> argument <> ") pgrst_scalar) pgrst_call"
                    | otherwise                                                              = fromQi qi <> "(" <> argument <> ") pgrst_call"

    fmtParams :: [RoutineParam] -> SQL.Snippet
    fmtParams prms = intercalateSnippet ", "
      ((\a -> (if ppVar a then "VARIADIC " else mempty) <> pgFmtIdent (ppName a) <> " := pgrst_body." <> pgFmtIdent (ppName a)) <$> prms)

    returnedColumns :: SQL.Snippet
    returnedColumns
      | null returnings = "*"
      | otherwise       = intercalateSnippet ", " (pgFmtColumn (QualifiedIdentifier mempty "pgrst_call") <$> returnings)

-- | SQL query meant for COUNTing the root node of the Tree.
-- It only takes WHERE into account and doesn't include LIMIT/OFFSET because it would reduce the COUNT.
-- SELECT 1 is done instead of SELECT * to prevent doing expensive operations(like functions based on the columns)
-- inside the FROM target.
-- If the request contains INNER JOINs, then the COUNT of the root node will change.
-- For this case, we use a WHERE EXISTS instead of an INNER JOIN on the count query.
-- See https://github.com/PostgREST/postgrest/issues/2009#issuecomment-977473031
-- Only for the nodes that have an INNER JOIN linked to the root level.
readPlanToCountQuery :: ReadPlanTree -> SQL.Snippet
readPlanToCountQuery (Node ReadPlan{from=mainQi, fromAlias=tblAlias, where_=logicForest, relToParent=rel, relJoinConds} forest) =
  "SELECT 1 " <> fromFrag <>
  (if null logicForest && null relJoinConds && null subQueries
    then mempty
    else " WHERE " ) <>
  intercalateSnippet " AND " (
    map (pgFmtLogicTreeCount qi) logicForest ++
    map pgFmtJoinCondition relJoinConds ++
    subQueries
  )
  where
    qi = getQualifiedIdentifier rel mainQi tblAlias
    fromFrag = fromF rel mainQi tblAlias
    subQueries = foldr existsSubquery [] forest
    existsSubquery :: ReadPlanTree -> [SQL.Snippet] -> [SQL.Snippet]
    existsSubquery readReq@(Node ReadPlan{relJoinType=joinType} _) rest =
      if joinType == Just JTInner
        then ("EXISTS (" <> readPlanToCountQuery readReq <> " )"):rest
        else rest
    findNullEmbedRel fld = find (\(Node ReadPlan{relAggAlias} _) -> fld == relAggAlias) forest

    -- https://github.com/PostgREST/postgrest/pull/2930#discussion_r1325293698
    pgFmtLogicTreeCount :: QualifiedIdentifier -> CoercibleLogicTree -> SQL.Snippet
    pgFmtLogicTreeCount qiCount (CoercibleExpr hasNot op frst) = SQL.sql notOp <> " (" <> intercalateSnippet (opSql op) (pgFmtLogicTreeCount qiCount <$> frst) <> ")"
      where
        notOp =  if hasNot then "NOT" else mempty
        opSql And = " AND "
        opSql Or  = " OR "
    pgFmtLogicTreeCount _ (CoercibleStmnt (CoercibleFilterNullEmbed hasNot fld)) =
      maybe mempty (\x -> (if not hasNot then "NOT " else mempty) <> "EXISTS (" <> readPlanToCountQuery x <> ")") (findNullEmbedRel fld)
    pgFmtLogicTreeCount qiCount (CoercibleStmnt flt) = pgFmtFilter qiCount flt

limitedQuery :: SQL.Snippet -> Maybe Integer -> SQL.Snippet
limitedQuery query maxRows = query <> SQL.sql (maybe mempty (\x -> " LIMIT " <> BS.pack (show x)) maxRows)

-- TODO refactor so this function is uneeded and ComputedRelationship QualifiedIdentifier comes from the ReadPlan type
getQualifiedIdentifier :: Maybe Relationship -> QualifiedIdentifier -> Maybe Alias -> QualifiedIdentifier
getQualifiedIdentifier rel mainQi tblAlias = case rel of
  Just ComputedRelationship{relFunction} -> QualifiedIdentifier mempty $ fromMaybe (qiName relFunction) tblAlias
  _                                      -> maybe mainQi (QualifiedIdentifier mempty) tblAlias

-- FROM clause plus implicit joins
fromF :: Maybe Relationship -> QualifiedIdentifier -> Maybe Alias -> SQL.Snippet
fromF rel mainQi tblAlias = "FROM " <>
  (case rel of
    -- Due to the use of CTEs on RPC, we need to cast the parameter to the table name in case of function overloading.
    -- See https://github.com/PostgREST/postgrest/issues/2963#issuecomment-1736557386
    Just ComputedRelationship{relFunction,relTableAlias,relTable} -> fromQi relFunction <> "(" <> pgFmtIdent (qiName relTableAlias) <> "::" <> fromQi relTable <> ")"
    _                                                             -> fromQi mainQi) <>
  maybe mempty (\a -> " AS " <> pgFmtIdent a) tblAlias <>
  (case rel of
    Just Relationship{relCardinality=M2M Junction{junTable=jt}} -> ", " <> fromQi jt
    _                                                           -> mempty)