-- | Postgres SQL IdentifierUniqueness
--
-- Bypass the Postgres limitation of truncating identifiers to 63 characters long
-- by prepending unique numbers.
--
-- See Note [Postgres identifier length limitations]
module Hasura.Backends.Postgres.SQL.IdentifierUniqueness
  ( prefixNumToAliases,
    prefixNumToAliasesSelectWith,
  )
where

import Data.HashMap.Strict qualified as Map
import Hasura.Backends.Postgres.SQL.DML qualified as S
import Hasura.Backends.Postgres.SQL.Types (Identifier (..))
import Hasura.Prelude

{- Note [Postgres identifier length limitations]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Postgres truncates identifiers to a maximum of 63 characters by default (see
https://www.postgresql.org/docs/12/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS).
-}

-- Prefix an Int to all aliases to preserve the uniqueness of identifiers.
-- See Note [Postgres identifier length limitations].
prefixNumToAliases :: S.Select -> S.Select
prefixNumToAliases s =
  uSelect s `evalState` UniqSt 0 Map.empty

prefixNumToAliasesSelectWith ::
  S.SelectWithG S.Select -> S.SelectWithG S.Select
prefixNumToAliasesSelectWith s =
  uSelectWith s `evalState` UniqSt 0 Map.empty

type Rewrite a = State a

data UniqSt = UniqSt
  { _uqVar :: !Int,
    _uqIdens :: !(Map.HashMap Identifier Int)
  }
  deriving (Show, Eq)

type Uniq = Rewrite UniqSt

withNumPfx :: Identifier -> Int -> Identifier
withNumPfx iden i =
  Identifier pfx <> iden
  where
    pfx = "_" <> tshow i <> "_"

addAlias :: S.Alias -> Uniq S.Alias
addAlias (S.Alias iden) = do
  UniqSt var idens <- get
  put $ UniqSt (var + 1) $ Map.insert iden var idens
  return $ S.Alias $ withNumPfx iden var

getIdentifier :: Identifier -> Uniq Identifier
getIdentifier iden = do
  UniqSt _ idens <- get
  let varNumM = Map.lookup iden idens
  return $ maybe iden (withNumPfx iden) varNumM

restoringIdens :: Uniq a -> Uniq a
restoringIdens action = do
  UniqSt _ idens <- get
  res <- action
  -- restore the idens to before the action
  modify' $ \s -> s {_uqIdens = idens}
  return res

uSelectWith :: S.SelectWithG S.Select -> Uniq (S.SelectWithG S.Select)
uSelectWith (S.SelectWith ctes baseSelect) =
  S.SelectWith
    <$> forM ctes (\(als, sel) -> (als,) <$> restoringIdens (uSelect sel))
    <*> uSelect baseSelect

uSelect :: S.Select -> Uniq S.Select
uSelect sel = do
  -- this has to be the first thing to process
  newFromM <- mapM uFromExp fromM
  newCTEs <- for ctes $ \(alias, cte) -> (,) <$> addAlias alias <*> uSelect cte

  newWhereM <- forM whereM $
    \(S.WhereFrag be) -> S.WhereFrag <$> uBoolExp be
  newGrpM <- forM grpM $
    \(S.GroupByExp l) -> S.GroupByExp <$> mapM uSqlExp l
  newHavnM <- forM havnM $
    \(S.HavingExp be) -> S.HavingExp <$> uBoolExp be
  newOrdM <- mapM uOrderBy ordByM
  newDistM <- mapM uDistinct distM
  newExtrs <- mapM uExtractor extrs
  return $
    S.Select
      newCTEs
      newDistM
      newExtrs
      newFromM
      newWhereM
      newGrpM
      newHavnM
      newOrdM
      limitM
      offM
  where
    S.Select ctes distM extrs fromM whereM grpM havnM ordByM limitM offM = sel
    uDistinct = \case
      S.DistinctSimple -> return S.DistinctSimple
      S.DistinctOn l -> S.DistinctOn <$> mapM uSqlExp l
    uExtractor (S.Extractor e alM) =
      S.Extractor <$> uSqlExp e <*> return alM

uFromExp :: S.FromExp -> Uniq S.FromExp
uFromExp (S.FromExp fromItems) =
  S.FromExp <$> mapM uFromItem fromItems

uFunctionArgs :: S.FunctionArgs -> Uniq S.FunctionArgs
uFunctionArgs (S.FunctionArgs positional named) =
  S.FunctionArgs <$> mapM uSqlExp positional <*> mapM uSqlExp named

uFunctionAlias :: S.FunctionAlias -> Uniq S.FunctionAlias
uFunctionAlias (S.FunctionAlias alias definitionList) =
  S.FunctionAlias <$> addAlias alias <*> pure definitionList

uFunctionExp :: S.FunctionExp -> Uniq S.FunctionExp
uFunctionExp (S.FunctionExp qf args alM) =
  S.FunctionExp qf <$> uFunctionArgs args <*> mapM uFunctionAlias alM

uFromItem :: S.FromItem -> Uniq S.FromItem
uFromItem fromItem = case fromItem of
  S.FISimple t alM ->
    S.FISimple t <$> mapM addAlias alM
  S.FIIdentifier iden ->
    S.FIIdentifier <$> return iden
  S.FIFunc funcExp ->
    S.FIFunc <$> uFunctionExp funcExp
  S.FIUnnest args als cols ->
    S.FIUnnest <$> mapM uSqlExp args <*> addAlias als <*> mapM uSqlExp cols
  S.FISelect isLateral sel al -> do
    -- we are kind of ignoring if we have to reset
    -- idens to empty based on correlation
    -- unless isLateral $ modify' $ \s -> s { _uqIdens = Map.empty}
    newSel <- restoringIdens $ uSelect sel
    newAls <- addAlias al
    return $ S.FISelect isLateral newSel newAls
  S.FISelectWith isLateral selectWith al -> do
    newSelectWith <- uSelectWith selectWith
    newAls <- addAlias al
    return $ S.FISelectWith isLateral newSelectWith newAls
  S.FIValues (S.ValuesExp tups) als mCols -> do
    newValExp <- fmap S.ValuesExp $
      forM tups $ \(S.TupleExp ts) ->
        S.TupleExp <$> mapM uSqlExp ts
    return $ S.FIValues newValExp als mCols
  S.FIJoin joinExp ->
    S.FIJoin <$> uJoinExp joinExp

uJoinExp :: S.JoinExpr -> Uniq S.JoinExpr
uJoinExp (S.JoinExpr left ty right joinCond) = do
  leftN <- uFromItem left
  rightN <- uFromItem right
  S.JoinExpr leftN ty rightN <$> uJoinCond joinCond

uJoinCond :: S.JoinCond -> Uniq S.JoinCond
uJoinCond joinCond = case joinCond of
  S.JoinOn be -> S.JoinOn <$> uBoolExp be
  S.JoinUsing cols -> return $ S.JoinUsing cols

uBoolExp :: S.BoolExp -> Uniq S.BoolExp
uBoolExp =
  restoringIdens . \case
    S.BELit b -> return $ S.BELit b
    S.BEBin op left right ->
      S.BEBin <$> return op <*> uBoolExp left <*> uBoolExp right
    S.BENot b -> S.BENot <$> uBoolExp b
    S.BECompare op left right ->
      S.BECompare <$> return op <*> uSqlExp left <*> uSqlExp right
    S.BECompareAny op left right ->
      S.BECompareAny <$> return op <*> uSqlExp left <*> uSqlExp right
    S.BENull e -> S.BENull <$> uSqlExp e
    S.BENotNull e -> S.BENotNull <$> uSqlExp e
    S.BEExists sel -> S.BEExists <$> uSelect sel
    S.BEIN left exps -> S.BEIN <$> uSqlExp left <*> mapM uSqlExp exps
    S.BEExp e -> S.BEExp <$> uSqlExp e

uOrderBy :: S.OrderByExp -> Uniq S.OrderByExp
uOrderBy (S.OrderByExp ordByItems) =
  S.OrderByExp <$> mapM uOrderByItem ordByItems
  where
    uOrderByItem (S.OrderByItem e ordTyM nullsOrdM) =
      S.OrderByItem
        <$> uSqlExp e
        <*> return ordTyM
        <*> return nullsOrdM

uSqlExp :: S.SQLExp -> Uniq S.SQLExp
uSqlExp =
  restoringIdens . \case
    S.SEPrep i -> return $ S.SEPrep i
    S.SENull -> return S.SENull
    S.SELit t -> return $ S.SELit t
    S.SEUnsafe t -> return $ S.SEUnsafe t
    S.SESelect s -> S.SESelect <$> uSelect s
    S.SEStar qual -> S.SEStar <$> traverse uQual qual
    -- this is for row expressions
    -- todo: check if this is always okay
    S.SEIdentifier iden -> return $ S.SEIdentifier iden
    S.SERowIdentifier iden -> S.SERowIdentifier <$> getIdentifier iden
    S.SEQIdentifier (S.QIdentifier qual iden) -> do
      newQual <- uQual qual
      return $ S.SEQIdentifier $ S.QIdentifier newQual iden
    S.SEFnApp fn args ordByM ->
      S.SEFnApp
        <$> return fn
        <*> mapM uSqlExp args
        <*> mapM uOrderBy ordByM
    S.SEOpApp op args ->
      S.SEOpApp op
        <$> mapM uSqlExp args
    S.SETyAnn e ty ->
      S.SETyAnn
        <$> uSqlExp e
        <*> return ty
    S.SECond be onTrue onFalse ->
      S.SECond
        <$> uBoolExp be
        <*> uSqlExp onTrue
        <*> uSqlExp onFalse
    S.SEBool be ->
      S.SEBool <$> uBoolExp be
    S.SEExcluded t ->
      S.SEExcluded <$> return t
    S.SEArray l ->
      S.SEArray <$> mapM uSqlExp l
    S.SEArrayIndex arrayExp indexExp ->
      S.SEArrayIndex <$> uSqlExp arrayExp <*> uSqlExp indexExp
    S.SETuple (S.TupleExp l) ->
      S.SETuple . S.TupleExp <$> mapM uSqlExp l
    S.SECount cty -> return $ S.SECount cty
    S.SENamedArg arg val -> S.SENamedArg arg <$> uSqlExp val
    S.SEFunction funcExp -> S.SEFunction <$> uFunctionExp funcExp
  where
    uQual = \case
      S.QualifiedIdentifier iden ty -> S.QualifiedIdentifier <$> getIdentifier iden <*> pure ty
      S.QualTable t -> return $ S.QualTable t
      S.QualVar t -> return $ S.QualVar t