{-# LANGUAGE CPP #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE DuplicateRecordFields    #-}
{-# LANGUAGE NamedFieldPuns           #-}
{-# LANGUAGE TemplateHaskell          #-}

module Preql.QuasiQuoter.Raw.TH where

import           Preql.QuasiQuoter.Raw.Lex (Token(..), unLex, parseQuery)
import           Preql.Wire (Query(..))

import           Data.String (IsString (..))
import           Data.Word (Word)
import           Language.Haskell.TH
import           Language.Haskell.TH.Quote
import           Language.Haskell.TH.Syntax (Lift (..))

import qualified Data.Text as T

-- | A list of n Names beginning with the given character
cNames :: Char -> Int -> Q [Name]
cNames :: Char -> Int -> Q [Name]
cNames Char
c Int
n = (String -> Q Name) -> [String] -> Q [Name]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse String -> Q Name
newName (Int -> String -> [String]
forall a. Int -> a -> [a]
replicate Int
n (Char
c Char -> String -> String
forall a. a -> [a] -> [a]
: String
""))

-- | Convert a rewritten SQL string to a ByteString
makeQuery :: String -> Q Exp
makeQuery :: String -> Q Exp
makeQuery String
string = [e|(fromString string :: Query) |]

-- | Given a SQL query with ${} antiquotes, splice a pair @(Query
-- p r, p)@ or a function @\p' -> (Query p r, p)@ if the SQL
-- string includes both antiquote and positional parameters.

-- | The @sql@ Quasiquoter allows passing parameters to a query by name, inside a @${}@ antiquote.  For example:
-- @[sql| SELECT name, age FROM cats WHERE age >= ${minAge} and age < ${maxAge} |]@
-- The Haskell term within @{}@ must be a variable in scope; more complex expressions are not supported.
--
-- Antiquotes are replaced by positional (@$1, $2@) parameters supported by Postgres, and the
-- encoded values are sent with @PexecParams@
--
-- Mixed named & numbered parameters are also supported.  It is hoped that this will be useful when
-- migrating existing queries.  For example:
-- @query $ [sql| SELECT name, age FROM cats WHERE age >= ${minAge} and age < $1 |] maxAge@
-- Named parameters will be assigned numbers higher than the highest numbered paramater placeholder.
--
-- A quote with only named parameters is converted to a tuple '(Query, p)'.  For example:
-- @("SELECT name, age FROM cats WHERE age >= $1 and age < $2", (minAge, maxAge))@
-- If there are no parameters, the inner tuple is @()@, like @("SELECT * FROM cats", ())@.
-- If there are both named & numbered params, the splice is a function taking a tuple and returning
-- @(Query, p)@ where p includes both named & numbered params.  For example:
-- @\a -> ("SELECT name, age FROM cats WHERE age >= $1 and age < $2", (a, maxAge))@
sql  :: QuasiQuoter
sql :: QuasiQuoter
sql  = String -> (String -> Q Exp) -> QuasiQuoter
expressionOnly String
"aritySql " ((String -> Q Exp) -> QuasiQuoter)
-> (String -> Q Exp) -> QuasiQuoter
forall a b. (a -> b) -> a -> b
$ \String
raw -> do
    Loc
loc <- Q Loc
location
    let e_ast :: Either String [Token]
e_ast = String -> String -> Either String [Token]
parseQuery (Loc -> String
forall a. Show a => a -> String
show Loc
loc) String
raw
    case Either String [Token]
e_ast of
        Right [Token]
parsed -> do
            let
                positionalCount :: Word
positionalCount = [Token] -> Word
maxParam [Token]
parsed
                (String
rewritten, [String]
haskellExpressions) = Word -> [Token] -> (String, [String])
numberAntiquotes Word
positionalCount [Token]
parsed
                -- mkName, because we intend to capture what's in scope
                antiNames :: [Name]
antiNames = (String -> Name) -> [String] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map String -> Name
mkName [String]
haskellExpressions
            Exp
query <- String -> Q Exp
makeQuery String
rewritten
            case Word
positionalCount of
                Word
0 -> -- only antiquotes (or no params)
                    Exp -> Q Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
tupleE [Exp
query, [Name] -> Exp
tupleOrSingle [Name]
antiNames]
                Word
1 -> do -- one positional param, doesn't take a tuple
                    Name
patternName <- String -> Q Name
newName String
"c"
                    Exp -> Q Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Pat] -> Exp -> Exp
LamE [Name -> Pat
VarP Name
patternName]
                        ([Exp] -> Exp
tupleE [Exp
query, [Name] -> Exp
tupleOrSingle (Name
patternName Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: [Name]
antiNames)])
                Word
_ -> do -- at least two positional parameters
                    [Name]
patternNames <- Char -> Int -> Q [Name]
cNames Char
'q' (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
positionalCount)
                    Exp -> Q Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Pat] -> Exp -> Exp
LamE
                        [[Pat] -> Pat
TupP ((Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
patternNames)]
                        ([Exp] -> Exp
tupleE [Exp
query, [Name] -> Exp
tupleOrSingle ([Name]
patternNames [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
++ [Name]
antiNames)])
        Left String
err -> String -> Q Exp
forall a. HasCallStack => String -> a
error String
err

tupleOrSingle :: [Name] -> Exp
tupleOrSingle :: [Name] -> Exp
tupleOrSingle [Name]
names = case [Name]
names of
    [Name
name] -> Name -> Exp
VarE Name
name
    [Name]
vs -> [Exp] -> Exp
tupleE ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Name -> Exp) -> [Name] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE [Name]
vs

expressionOnly :: String -> (String -> Q Exp) -> QuasiQuoter
expressionOnly :: String -> (String -> Q Exp) -> QuasiQuoter
expressionOnly String
name String -> Q Exp
qq = QuasiQuoter :: (String -> Q Exp)
-> (String -> Q Pat)
-> (String -> Q Type)
-> (String -> Q [Dec])
-> QuasiQuoter
QuasiQuoter
    { quoteExp :: String -> Q Exp
quoteExp = String -> Q Exp
qq
    , quotePat :: String -> Q Pat
quotePat = \String
_ -> String -> Q Pat
forall a. HasCallStack => String -> a
error (String -> Q Pat) -> String -> Q Pat
forall a b. (a -> b) -> a -> b
$ String
"qq " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" cannot be used in pattern context"
    , quoteType :: String -> Q Type
quoteType = \String
_ -> String -> Q Type
forall a. HasCallStack => String -> a
error (String -> Q Type) -> String -> Q Type
forall a b. (a -> b) -> a -> b
$ String
"qq " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" cannot be used in type context"
    , quoteDec :: String -> Q [Dec]
quoteDec = \String
_ -> String -> Q [Dec]
forall a. HasCallStack => String -> a
error (String -> Q [Dec]) -> String -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ String
"qq " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" cannot be used in declaration context"
    }

maxParam :: [Token] -> Word
maxParam :: [Token] -> Word
maxParam = (Token -> Word -> Word) -> Word -> [Token] -> Word
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Token -> Word -> Word
nextParam Word
0 where
  nextParam :: Token -> Word -> Word
nextParam Token
token Word
maxParam =
      case Token
token of
          NumberedParam Word
i -> Word -> Word -> Word
forall a. Ord a => a -> a -> a
max Word
i Word
maxParam
          Token
_ -> Word
maxParam

numberAntiquotes :: Word -> [Token] -> (String, [String])
numberAntiquotes :: Word -> [Token] -> (String, [String])
numberAntiquotes Word
mp [Token]
ts = ([String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [String]
sqlStrings, [String]
variableNames) where
  ([String]
sqlStrings, [String]
variableNames) = Word -> [Token] -> ([String], [String])
go Word
mp [Token]
ts
  go :: Word -> [Token] -> ([String], [String])
go Word
_maxParam [] = ([], [])
  go Word
maxParam (Token
token : [Token]
ts) =
      case Token
token of
          HaskellParam String
name -> let
              newParam :: Word
newParam = Word
maxParam Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1
              ([String]
ss, [String]
ns) = Word -> [Token] -> ([String], [String])
go Word
newParam [Token]
ts
              in (Token -> String
unLex (Word -> Token
NumberedParam Word
newParam) String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
ss, String
name String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
ns)
          Token
EOF -> Word -> [Token] -> ([String], [String])
go Word
maxParam [Token]
ts
          Token
_ -> let ([String]
ss, [String]
ns) = Word -> [Token] -> ([String], [String])
go Word
maxParam [Token]
ts in (Token -> String
unLex Token
token String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
ss, [String]
ns)

tupleE :: [Exp] -> Exp
#if MIN_VERSION_template_haskell(2,16,0)
tupleE :: [Exp] -> Exp
tupleE = [Maybe Exp] -> Exp
TupE ([Maybe Exp] -> Exp) -> ([Exp] -> [Maybe Exp]) -> [Exp] -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp -> Maybe Exp) -> [Exp] -> [Maybe Exp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Exp -> Maybe Exp
forall a. a -> Maybe a
Just
#else
tupleE = TupE
#endif