{-# language DeriveFunctor #-}
{-# language DerivingVia #-}
{-# language FlexibleContexts #-}
{-# language GADTs #-}
{-# language LambdaCase #-}
{-# language NamedFieldPuns #-}
{-# language RankNTypes #-}
{-# language RecordWildCards #-}
{-# language ScopedTypeVariables #-}
{-# language StandaloneKindSignatures #-}
{-# language TypeApplications #-}

module Rel8.Statement
  ( Statement
  , statementReturning
  , statementNoReturning
  , ppDecodeStatement
  )
where

-- base
import Control.Applicative (liftA2)
import Control.Monad (ap, liftM2)
import Data.Foldable (fold, toList)
import Data.Int (Int64)
import Data.Kind (Type)
import Data.List.NonEmpty (NonEmpty, intersperse)
import Data.Monoid (Endo (Endo))
import Data.String (fromString)
import Prelude

-- hasql
import qualified Hasql.Decoders as Hasql

-- opaleye
import qualified Opaleye.Internal.Tag as Opaleye

-- pretty
import Text.PrettyPrint
  ( Doc
  , (<+>)
  , ($$)
  , comma
  , hcat
  , parens
  , punctuate
  , text
  , vcat
  )

-- rel8
import Rel8.Expr (Expr)
import Rel8.Expr.Bool (false)
import Rel8.Query (Query)
import Rel8.Query.Aggregate (countRows)
import Rel8.Query.Each (each)
import Rel8.Schema.Escape (escape)
import Rel8.Schema.Table (TableSchema (..))
import Rel8.Statement.Rows (Rows (..))
import Rel8.Table (Table)
import Rel8.Table.Cols (fromCols)
import Rel8.Table.Name (namesFromLabelsWithA, showNames)
import Rel8.Table.Serialize (parse)

-- semigroupoids
import Data.Functor.Apply (Apply, WrappedApplicative (..))
import Data.Functor.Bind (Bind, (>>-))

-- transformers
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.State.Strict (State, evalState)
import Control.Monad.Trans.Writer.CPS (WriterT, runWriterT, tell)


type Binding :: Type
data Binding = Binding
  { relation :: !String
  , columns :: !(Maybe (NonEmpty String))
  , doc :: !Doc
  , returning :: !Returning
  }


type Result :: Type -> Type
data Result a = Unmodified !a | Modified !a


instance Functor Result where
  fmap f = \case
    Unmodified a -> Modified (f a)
    Modified a -> Modified (f a)


getResult :: Result a -> a
getResult = \case
  Unmodified a -> a
  Modified a -> a


type Returning :: Type
data Returning where
  NoReturning :: Returning
  Returning :: Query (Expr Int64) -> Returning 


-- | 'Statement' represents a single PostgreSQL statement. Most commonly,
-- this is constructed using 'Rel8.select', 'Rel8.insert', 'Rel8.update'
-- or 'Rel8.delete'.
--
-- However, in addition to @SELECT@, @INSERT@, @UPDATE@ and @DELETE@,
-- PostgreSQL also supports compositions thereof via its statement-level
-- @WITH@ syntax (with some caveats). Each such \"sub-statement\" can
-- reference the results of previous sub-statements. 'Statement' provides a
-- 'Monad' instance that captures this \"binding\" pattern.
--
-- The caveat with this is that the [side-effects of these sub-statements
-- are not visible to other sub-statements](https://www.postgresql.org/docs/current/queries-with.html#QUERIES-WITH-MODIFYING);
-- only the explicit results of previous sub-statements (from @SELECT@s or
-- @RETURNING@ clauses) are visible. So, for example, an @INSERT@ into a table
-- followed immediately by a @SELECT@ therefrom will not return the inserted
-- rows. However, it is possible to return the inserted rows using
-- @RETURNING@, 'Rel8.unionAll'ing this with the result of a @SELECT@
-- from the same table will produce the desired result.
--
-- An example of where this can be useful is if you want to delete rows from
-- a table and simultaneously log their deletion in a log table.
--
-- @
-- deleteFoo :: (Foo Expr -> Expr Bool) -> Statement ()
-- deleteFoo predicate = do
--   foos <-
--     delete Delete
--       { from = fooSchema
--       , using = pure ()
--       , deleteWhere = \_ -> predicate
--       , returning = Returning id
--       }
--   insert Insert
--     { into = deletedFooSchema
--     , rows = do
--         Foo {..} <- foos
--         let
--           deletedAt = 'Rel8.Expr.Time.now'
--         pure DeletedFoo {..}
--     , onConflict = Abort
--     , returning = NoReturning
--     }
-- @
newtype Statement a =
  Statement (WriterT (Endo [Binding]) (State Opaleye.Tag) (Result a))
  deriving stock (Functor)
  deriving (Apply) via WrappedApplicative Statement


instance Applicative Statement where
  pure = Statement . pure . Modified
  (<*>) = ap
  liftA2 = liftM2


instance Bind Statement where
  Statement m >>- f = Statement $ do
    result <- m
    case f (getResult result) of
      Statement m' -> m'


instance Monad Statement where
  (>>=) = (>>-)


statementNoReturning :: State Opaleye.Tag Doc -> Statement ()
statementNoReturning pp = Statement $ do
  binding <- lift $ do
    doc <- pp
    tag <- Opaleye.fresh
    let
      relation = Opaleye.tagWith tag "statement"
      columns = Nothing
      returning = NoReturning
      binding = Binding {..}
    pure binding
  tell (Endo (binding :))
  pure $ Unmodified ()


statementReturning :: Table Expr a 
  => State Opaleye.Tag Doc -> Statement (Query a)
statementReturning pp = Statement $ do
  (binding, query) <- lift $ do
    doc <- pp
    tag <- Opaleye.fresh
    let
      relation = Opaleye.tagWith tag "statement"
      symbol labels = do
        subtag <- Opaleye.fresh
        let
          suffix = Opaleye.tagWith tag (Opaleye.tagWith subtag "")
        pure $ take (63 - length suffix) label ++ suffix
        where
          label = fold (intersperse "/" labels)
      names = namesFromLabelsWithA symbol `evalState` Opaleye.start
      columns = Just $ showNames names
      query =
        fromCols <$> each
          TableSchema
            { name = fromString relation
            , columns = names
            }
      returning = Returning (countRows query)
      binding = Binding {..}
    pure (binding, query)
  tell (Endo (binding :))
  pure $ Unmodified query


ppDecodeStatement :: ()
  => (forall x. Table Expr x => Query x -> State Opaleye.Tag Doc)
  -> Rows exprs a -> Statement exprs -> (Doc, Hasql.Result a)
ppDecodeStatement ppSelect rows (Statement m) = evalState go Opaleye.start
  where
    go = do
      (result, Endo dlist) <- runWriterT m
      let
        bindings' = dlist []
      case unsnoc bindings' of
        Nothing -> case rows of
          Void -> do
            doc <- ppSelect (pure false)
            pure (doc, Hasql.noResult)
          RowsAffected -> do
            doc <- ppSelect (pure false)
            pure (doc, 0 <$ Hasql.noResult)
          Single @exprs @a -> do
            doc <- ppSelect (getResult result)
            pure (doc, Hasql.singleRow (parse @exprs @a))
          Maybe @exprs @a -> do
            doc <- ppSelect (getResult result)
            pure (doc, Hasql.rowMaybe (parse @exprs @a))
          List @exprs @a -> do
            doc <- ppSelect (getResult result)
            pure (doc, Hasql.rowList (parse @exprs @a))
          Vector @exprs @a -> do
            doc <- ppSelect (getResult result)
            pure (doc, Hasql.rowVector (parse @exprs @a))
        Just (bindings, binding@Binding {doc = after}) -> case rows of
          Void -> pure (doc, Hasql.noResult)
            where
              doc = ppWith bindings after
          RowsAffected -> do
            case result of
              Unmodified _ -> pure (doc, Hasql.rowsAffected)
                where
                  doc = ppWith bindings after
              Modified _ -> case returning binding of
                NoReturning -> pure (doc, Hasql.rowsAffected)
                  where
                    doc = ppWith bindings after 
                Returning query -> do
                  doc <- ppWith bindings' <$> ppSelect query
                  pure (doc, Hasql.singleRow parse)
          Single @exprs @a -> do
            case result of
              Unmodified _ -> pure (doc, Hasql.singleRow (parse @exprs @a))
                where
                  doc = ppWith bindings after
              Modified query -> do
                doc <- ppWith bindings' <$> ppSelect query
                pure (doc, Hasql.singleRow (parse @exprs @a))
          Maybe @exprs @a -> do
            case result of
              Unmodified _ -> pure (doc, Hasql.rowMaybe (parse @exprs @a))
                where
                  doc = ppWith bindings after
              Modified query -> do
                doc <- ppWith bindings' <$> ppSelect query
                pure (doc, Hasql.rowMaybe (parse @exprs @a))
          List @exprs @a -> do
            case result of
              Unmodified _ -> pure (doc, Hasql.rowList (parse @exprs @a))
                where
                  doc = ppWith bindings after
              Modified query -> do
                doc <- ppWith bindings' <$> ppSelect query
                pure (doc, Hasql.rowList (parse @exprs @a))
          Vector @exprs @a -> do
            case result of
              Unmodified _ -> pure (doc, Hasql.rowVector (parse @exprs @a))
                where
                  doc = ppWith bindings after
              Modified query -> do
                doc <- ppWith bindings' <$> ppSelect query
                pure (doc, Hasql.rowVector (parse @exprs @a))


ppWith :: [Binding] -> Doc -> Doc
ppWith bindings after = pre $$ after
  where
    pre = case bindings of
      [] -> mempty
      _ ->
        text "WITH" <+>
        vcat (punctuate comma (map go bindings))
    go binding@Binding {doc = before} =
      ppAlias binding $$
      text "AS" <+>
      parens before


ppAlias :: Binding -> Doc
ppAlias Binding {relation, columns = mcolumns} = case mcolumns of
  Nothing -> escape relation
  Just columns -> 
    escape relation <+>
    parens (hcat (punctuate comma (escape <$> toList columns)))


unsnoc :: [a] -> Maybe ([a], a)
unsnoc = foldr (\x -> Just . maybe ([], x) (\(~(a, b)) -> (x : a, b))) Nothing