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

module Preql.QuasiQuoter.Raw.TH where

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

import Data.String (IsString(..))
import Language.Haskell.TH
import Language.Haskell.TH.Quote

-- | Convert a rewritten SQL string to a ByteString, leaving width free
makeQuery :: String -> Q Exp
makeQuery :: String -> Q Exp
makeQuery String
string = [e|(fromString string :: Query $(VarT <$> (newName "n"))) |]

-- | 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
"sql " ((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

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
maxSoFar =
      case Token
token of
          NumberedParam Word
i -> Word -> Word -> Word
forall a. Ord a => a -> a -> a
max Word
i Word
maxSoFar
          Token
_               -> Word
maxSoFar

numberAntiquotes :: Word -> [Token] -> (String, [String])
numberAntiquotes :: Word -> [Token] -> (String, [String])
numberAntiquotes Word
mp [Token]
tokens = ([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]
tokens
  go :: Word -> [Token] -> ([String], [String])
go Word
_maxSoFar [] = ([], [])
  go Word
maxSoFar (Token
token : [Token]
ts) =
      case Token
token of
          HaskellParam String
name -> let
              newParam :: Word
newParam = Word
maxSoFar 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
maxSoFar [Token]
ts
          Token
_ -> let ([String]
ss, [String]
ns) = Word -> [Token] -> ([String], [String])
go Word
maxSoFar [Token]
ts in (Token -> String
unLex Token
token String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
ss, [String]
ns)