{-|
Module: Squeal.PostgreSQL.Render
Description: render functions
Copyright: (c) Eitan Chatav, 2019
Maintainer: eitan@morphism.tech
Stability: experimental

render functions
-}

{-# LANGUAGE
    AllowAmbiguousTypes
  , ConstraintKinds
  , FlexibleContexts
  , LambdaCase
  , MagicHash
  , OverloadedStrings
  , PolyKinds
  , RankNTypes
  , ScopedTypeVariables
  , TypeApplications
#-}

module Squeal.PostgreSQL.Render
  ( -- * Render
    RenderSQL (..)
  , printSQL
  , escape
  , parenthesized
  , bracketed
  , (<+>)
  , commaSeparated
  , doubleQuoted
  , singleQuotedText
  , singleQuotedUtf8
  , escapeQuotedString
  , escapeQuotedText
  , renderCommaSeparated
  , renderCommaSeparatedConstraint
  , renderCommaSeparatedMaybe
  , renderNat
  , renderSymbol
  ) where

import Control.Monad.IO.Class (MonadIO (..))
import Data.ByteString (ByteString)
import Data.Maybe (catMaybes)
import Data.Text (Text)
import Generics.SOP
import GHC.Exts
import GHC.TypeLits hiding (Text)

import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as Char8

-- | Parenthesize a `ByteString`.
parenthesized :: ByteString -> ByteString
parenthesized str = "(" <> str <> ")"

-- | Square bracket a `ByteString`
bracketed :: ByteString -> ByteString
bracketed str = "[" <> str <> "]"

-- | Concatenate two `ByteString`s with a space between.
(<+>) :: ByteString -> ByteString -> ByteString
infixr 7 <+>
str1 <+> str2 = str1 <> " " <> str2

-- | Comma separate a list of `ByteString`s.
commaSeparated :: [ByteString] -> ByteString
commaSeparated = ByteString.intercalate ", "

-- | Add double quotes around a `ByteString`.
doubleQuoted :: ByteString -> ByteString
doubleQuoted str = "\"" <> str <> "\""

-- | Add single quotes around a `Text` and escape single quotes within it.
singleQuotedText :: Text -> ByteString
singleQuotedText str =
  "'" <> Text.encodeUtf8 (Text.replace "'" "''" str) <> "'"

-- | Add single quotes around a `ByteString` and escape single quotes within it.
singleQuotedUtf8 :: ByteString -> ByteString
singleQuotedUtf8 = singleQuotedText . Text.decodeUtf8

-- | Escape quote a string.
escapeQuotedString :: String -> ByteString
escapeQuotedString x = "E\'" <> Text.encodeUtf8 (fromString (escape =<< x)) <> "\'"

-- | Escape quote a string.
escapeQuotedText :: Text -> ByteString
escapeQuotedText x =
  "E\'" <> Text.encodeUtf8 (Text.concatMap (fromString . escape) x) <> "\'"

-- | Comma separate the renderings of a heterogeneous list.
renderCommaSeparated
  :: SListI xs
  => (forall x. expression x -> ByteString)
  -> NP expression xs -> ByteString
renderCommaSeparated render
  = commaSeparated
  . hcollapse
  . hmap (K . render)

-- | Comma separate the renderings of a heterogeneous list.
renderCommaSeparatedConstraint
  :: forall c xs expression. (All c xs, SListI xs)
  => (forall x. c x => expression x -> ByteString)
  -> NP expression xs -> ByteString
renderCommaSeparatedConstraint render
  = commaSeparated
  . hcollapse
  . hcmap (Proxy @c) (K . render)

-- | Comma separate the `Maybe` renderings of a heterogeneous list, dropping
-- `Nothing`s.
renderCommaSeparatedMaybe
  :: SListI xs
  => (forall x. expression x -> Maybe ByteString)
  -> NP expression xs -> ByteString
renderCommaSeparatedMaybe render
  = commaSeparated
  . catMaybes
  . hcollapse
  . hmap (K . render)

-- | Render a promoted `Nat`.
renderNat :: forall n. KnownNat n => ByteString
renderNat = fromString (show (natVal' (proxy# :: Proxy# n)))

-- | Render a promoted `Symbol`.
renderSymbol :: forall s. KnownSymbol s => ByteString
renderSymbol = fromString (symbolVal' (proxy# :: Proxy# s))

-- | A class for rendering SQL
class RenderSQL sql where renderSQL :: sql -> ByteString

-- | Print SQL.
printSQL :: (RenderSQL sql, MonadIO io) => sql -> io ()
printSQL = liftIO . Char8.putStrLn . renderSQL

-- | `escape` a character to prevent injection
escape :: Char -> String
escape = \case
  '\NUL' -> ""
  '\'' -> "''"
  '"' -> "\\\""
  '\b' -> "\\b"
  '\n' -> "\\n"
  '\r' -> "\\r"
  '\t' -> "\\t"
  '\\' -> "\\\\"
  c -> [c]