{-# LANGUAGE DataKinds, UndecidableInstances, FlexibleContexts #-}
{-# LANGUAGE TypeFamilies, TypeOperators, TypeApplications #-}

-- | This module provides the top level API for Souffle related operations.
--   It makes use of Haskell's powerful typesystem to make certain invalid states
--   impossible to represent. It does this with a small type level DSL for
--   describing properties of the Datalog program (see the 'Program' and 'Fact'
--   typeclasses for more information).
--
--   The Souffle operations are exposed via 2 mtl-style interfaces
--   (see `MonadSouffle` and `MonadSouffleFileIO`) that allows them to be
--   integrated with existing monad transformer stacks.
--
--   This module also contains some helper type families for additional
--   type safety and user-friendly error messages.
module Language.Souffle.Class
  ( Program(..)
  , ProgramOptions(..)
  , Fact(..)
  , FactOptions(..)
  , Marshal.Marshal(..)
  , Direction(..)
  , ContainsInputFact
  , ContainsOutputFact
  , ContainsFact
  , MonadSouffle(..)
  , MonadSouffleFileIO(..)
  ) where

import Prelude hiding ( init )

import Control.Monad.Except
import Control.Monad.RWS.Strict
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Writer
import Data.Proxy
import Data.Kind
import Data.Word
import GHC.TypeLits
import qualified Language.Souffle.Marshal as Marshal
import Type.Errors.Pretty


-- | A helper type family for checking if a specific Souffle `Program` contains
--   a certain `Fact`. Additionally, it also checks if the fact is marked as
--   either `Input` or `InputOutput`. This constraint will generate a
--   user-friendly type error if these conditions are not met.
type family ContainsInputFact prog fact :: Constraint where
  ContainsInputFact prog fact = (ContainsFact prog fact, IsInput fact (FactDirection fact))

-- | A helper type family for checking if a specific Souffle `Program` contains
--   a certain `Fact`. Additionally, it also checks if the fact is marked as
--   either `Output` or `InputOutput`. This constraint will generate a
--   user-friendly type error if these conditions are not met.
type family ContainsOutputFact prog fact :: Constraint where
  ContainsOutputFact prog fact = (ContainsFact prog fact, IsOutput fact (FactDirection fact))

type family IsInput (fact :: Type) (dir :: Direction) :: Constraint where
  IsInput _ 'Input = ()
  IsInput _ 'InputOutput = ()
  IsInput fact dir = TypeError
    ( "You tried to use an " <> FormatDirection dir <> " fact of type " <> fact <> " as an input."
    % "Possible solution: change the FactDirection of " <> fact
      <> " to either 'Input' or 'InputOutput'."
    )

type family IsOutput (fact :: Type) (dir :: Direction) :: Constraint where
  IsOutput _ 'Output = ()
  IsOutput _ 'InputOutput = ()
  IsOutput fact dir = TypeError
    ( "You tried to use an " <> FormatDirection dir <> " fact of type " <> fact <> " as an output."
    % "Possible solution: change the FactDirection of " <> fact
      <> " to either 'Output' or 'InputOutput'."
    )

type family FormatDirection (dir :: Direction) where
  FormatDirection 'Output = "output"
  FormatDirection 'Input = "input"
  FormatDirection 'Internal = "internal"

-- | A helper type family for checking if a specific Souffle `Program` contains
--   a certain `Fact`. This constraint will generate a user-friendly type error
--   if this is not the case.
type family ContainsFact prog fact :: Constraint where
  ContainsFact prog fact =
    CheckContains prog (ProgramFacts prog) fact

type family CheckContains prog facts fact :: Constraint where
  CheckContains prog '[] fact =
    TypeError ("You tried to perform an action with a fact of type '" <> fact
    <> "' for program '" <> prog <> "'."
    % "The program contains the following facts: " <> ProgramFacts prog <> "."
    % "It does not contain fact: " <> fact <> "."
    % "You can fix this error by adding the type '" <> fact
    <> "' to the ProgramFacts type in the Program instance for " <> prog <> ".")
  CheckContains _ (a ': _) a = ()
  CheckContains prog (_ ': as) b = CheckContains prog as b

-- | A typeclass for describing a datalog program.
--
-- Example usage (assuming the program was generated from path.dl
-- and contains 2 facts: Edge and Reachable):
--
-- @
-- data Path = Path  -- Handle for the datalog program
--
-- instance Program Path where
--   type ProgramFacts Path = '[Edge, Reachable]
--   programName = const "path"
-- @
class Program a where
  -- | A type level list of facts that belong to this program.
  --   This list is used to check that only known facts are added to a program.
  type ProgramFacts a :: [Type]

  -- | Function for obtaining the name of a Datalog program.
  --   This has to be the same as the name of the .dl file (minus the extension).
  programName :: a -> String

-- | A helper data type, used in combination with the DerivingVia extension to
--   automatically generate code to bind Haskell to a Souffle Datalog program.
--
-- The following is an example how to bind to a Datalog program "path"
-- (saved as path.dl / path.cpp), that uses two facts called "edge" and
-- "reachable" (represented with the Edge and Reachable types):
--
-- @
-- data Path = Path
--   deriving Souffle.Program
--   via Souffle.ProgramOptions Path "path" '[Edge, Reachable]
-- @
--
-- See also: 'FactOptions'.
newtype ProgramOptions (prog :: Type) (progName :: Symbol) (facts :: [Type])
  = ProgramOptions prog

instance KnownSymbol progName => Program (ProgramOptions prog progName facts) where
  type ProgramFacts (ProgramOptions _ _ facts) = facts

  programName = const $ symbolVal (Proxy @progName)
  {-# INLINABLE programName #-}

-- | A typeclass for data types representing a fact in datalog.
--
-- Example usage:
--
-- @
-- instance Fact Edge where
--   type FactDirection Edge = 'Input
--   factName = const "edge"
-- @
class Marshal.Marshal a => Fact a where
  -- | The direction or "mode" a fact can be used in.
  --   This is used to perform compile-time checks that a fact is only used
  --   in valid situations. For more information, see the 'Direction' type.
  type FactDirection a :: Direction

  -- | Function for obtaining the name of a fact
  --   (has to be the same as described in the Datalog program).
  --
  -- It uses a 'Proxy' to select the correct instance.
  factName :: Proxy a -> String

-- | A helper data type, used in combination with the DerivingVia extension to
--   automatically generate code to bind a Haskell datatype to a Souffle
--   Datalog fact.
--
-- The following is an example how to bind to a Datalog fact "edge"
-- that contains two symbols (strings in Haskell) that is an input (from the
-- Datalog point of view):
--
-- @
-- data Edge = Edge String String
--   deriving (Eq, Show, Generic)
--   deriving anyclass Souffle.Marshal
--   deriving Souffle.Fact
--   via Souffle.FactOptions Edge "edge" 'Souffle.Input
-- @
--
-- See also: 'ProgramOptions'.
newtype FactOptions (fact :: Type) (factName :: Symbol) (dir :: Direction)
  = FactOptions fact

instance Marshal.Marshal fact => Marshal.Marshal (FactOptions fact name dir) where
  push (FactOptions fact) = Marshal.push fact
  {-# INLINABLE push #-}
  pop = FactOptions <$> Marshal.pop
  {-# INLINABLE pop #-}

instance ( Marshal.Marshal fact
         , KnownSymbol factName
         ) => Fact (FactOptions fact factName dir) where
  type FactDirection (FactOptions _ _ dir) = dir

  factName = const $ symbolVal (Proxy @factName)
  {-# INLINABLE factName #-}


-- | A datatype describing which operations a certain fact supports.
--   The direction is from the datalog perspective, so that it
--   aligns with ".decl" statements in Souffle.
data Direction
  = Input
  -- ^ Fact can only be stored in Datalog (using `addFact`/`addFacts`).
  | Output
  -- ^ Fact can only be read from Datalog (using `getFacts`/`findFact`).
  | InputOutput
  -- ^ Fact supports both reading from / writing to Datalog.
  | Internal
  -- ^ Supports neither reading from / writing to Datalog. This is used for
  --   facts that are only visible inside Datalog itself.

-- | A mtl-style typeclass for Souffle-related actions.
class Monad m => MonadSouffle m where
  -- | Represents a handle for interacting with a Souffle program.
  --
  --   The handle is used in all other functions of this typeclass to perform
  --   Souffle-related actions.
  type Handler m :: Type -> Type

  -- | Helper associated type constraint that allows collecting facts from
  --   Souffle in a list or vector. Only used internally.
  type CollectFacts m (c :: Type -> Type) :: Constraint
  -- | Helper associated type constraint that allows submitting facts to
  --   Souffle. Only used internally.
  type SubmitFacts m (a :: Type) :: Constraint

  -- | Runs the Souffle program.
  run :: Handler m prog -> m ()

  -- | Sets the number of CPU cores this Souffle program should use.
  setNumThreads :: Handler m prog -> Word64 -> m ()

  -- | Gets the number of CPU cores this Souffle program should use.
  getNumThreads :: Handler m prog -> m Word64

  -- | Returns all facts of a program. This function makes use of type inference
  --   to select the type of fact to return.
  getFacts :: (Fact a, ContainsOutputFact prog a, CollectFacts m c)
           => Handler m prog -> m (c a)

  -- | Searches for a fact in a program.
  --   Returns 'Nothing' if no matching fact was found; otherwise 'Just' the fact.
  --
  --   Conceptually equivalent to @List.find (== fact) \<$\> getFacts prog@,
  --   but this operation can be implemented much faster.
  findFact :: (Fact a, ContainsOutputFact prog a, Eq a, SubmitFacts m a)
           => Handler m prog -> a -> m (Maybe a)

  -- | Adds a fact to the program.
  addFact :: (Fact a, ContainsInputFact prog a, SubmitFacts m a)
          => Handler m prog -> a -> m ()

  -- | Adds multiple facts to the program. This function could be implemented
  --   in terms of 'addFact', but this is done as a minor optimization.
  addFacts :: (Foldable t, Fact a, ContainsInputFact prog a, SubmitFacts m a)
           => Handler m prog -> t a -> m ()

instance MonadSouffle m => MonadSouffle (ReaderT r m) where
  type Handler (ReaderT r m) = Handler m
  type CollectFacts (ReaderT r m) c = CollectFacts m c
  type SubmitFacts (ReaderT r m) a = SubmitFacts m a

  run = lift . run
  {-# INLINABLE run #-}
  setNumThreads prog = lift . setNumThreads prog
  {-# INLINABLE setNumThreads #-}
  getNumThreads = lift . getNumThreads
  {-# INLINABLE getNumThreads #-}
  getFacts = lift . getFacts
  {-# INLINABLE getFacts #-}
  findFact prog = lift . findFact prog
  {-# INLINABLE findFact #-}
  addFact fact = lift . addFact fact
  {-# INLINABLE addFact #-}
  addFacts facts = lift . addFacts facts
  {-# INLINABLE addFacts #-}

instance (Monoid w, MonadSouffle m) => MonadSouffle (WriterT w m) where
  type Handler (WriterT w m) = Handler m
  type CollectFacts (WriterT w m) c = CollectFacts m c
  type SubmitFacts (WriterT w m) a = SubmitFacts m a

  run = lift . run
  {-# INLINABLE run #-}
  setNumThreads prog = lift . setNumThreads prog
  {-# INLINABLE setNumThreads #-}
  getNumThreads = lift . getNumThreads
  {-# INLINABLE getNumThreads #-}
  getFacts = lift . getFacts
  {-# INLINABLE getFacts #-}
  findFact prog = lift . findFact prog
  {-# INLINABLE findFact #-}
  addFact fact = lift . addFact fact
  {-# INLINABLE addFact #-}
  addFacts facts = lift . addFacts facts
  {-# INLINABLE addFacts #-}

instance MonadSouffle m => MonadSouffle (StateT s m) where
  type Handler (StateT s m) = Handler m
  type CollectFacts (StateT s m) c = CollectFacts m c
  type SubmitFacts (StateT s m) a = SubmitFacts m a

  run = lift . run
  {-# INLINABLE run #-}
  setNumThreads prog = lift . setNumThreads prog
  {-# INLINABLE setNumThreads #-}
  getNumThreads = lift . getNumThreads
  {-# INLINABLE getNumThreads #-}
  getFacts = lift . getFacts
  {-# INLINABLE getFacts #-}
  findFact prog = lift . findFact prog
  {-# INLINABLE findFact #-}
  addFact fact = lift . addFact fact
  {-# INLINABLE addFact #-}
  addFacts facts = lift . addFacts facts
  {-# INLINABLE addFacts #-}

instance (MonadSouffle m, Monoid w) => MonadSouffle (RWST r w s m) where
  type Handler (RWST r w s m) = Handler m
  type CollectFacts (RWST r w s m) c = CollectFacts m c
  type SubmitFacts (RWST r w s m) a = SubmitFacts m a

  run = lift . run
  {-# INLINABLE run #-}
  setNumThreads prog = lift . setNumThreads prog
  {-# INLINABLE setNumThreads #-}
  getNumThreads = lift . getNumThreads
  {-# INLINABLE getNumThreads #-}
  getFacts = lift . getFacts
  {-# INLINABLE getFacts #-}
  findFact prog = lift . findFact prog
  {-# INLINABLE findFact #-}
  addFact fact = lift . addFact fact
  {-# INLINABLE addFact #-}
  addFacts facts = lift . addFacts facts
  {-# INLINABLE addFacts #-}

instance MonadSouffle m => MonadSouffle (ExceptT e m) where
  type Handler (ExceptT e m) = Handler m
  type CollectFacts (ExceptT e m) c = CollectFacts m c
  type SubmitFacts (ExceptT e m) a = SubmitFacts m a

  run = lift . run
  {-# INLINABLE run #-}
  setNumThreads prog = lift . setNumThreads prog
  {-# INLINABLE setNumThreads #-}
  getNumThreads = lift . getNumThreads
  {-# INLINABLE getNumThreads #-}
  getFacts = lift . getFacts
  {-# INLINABLE getFacts #-}
  findFact prog = lift . findFact prog
  {-# INLINABLE findFact #-}
  addFact fact = lift . addFact fact
  {-# INLINABLE addFact #-}
  addFacts facts = lift . addFacts facts
  {-# INLINABLE addFacts #-}

-- | A mtl-style typeclass for Souffle-related actions that involve file IO.
class MonadSouffle m => MonadSouffleFileIO m where
  -- | Load all facts from files in a certain directory.
  loadFiles :: Handler m prog -> FilePath -> m ()

  -- | Write out all facts of the program to CSV files in a certain directory
  --   (as defined in the Souffle program).
  writeFiles :: Handler m prog -> FilePath -> m ()

instance MonadSouffleFileIO m => MonadSouffleFileIO (ReaderT r m) where
  loadFiles prog = lift . loadFiles prog
  {-# INLINABLE loadFiles #-}
  writeFiles prog = lift . writeFiles prog
  {-# INLINABLE writeFiles #-}

instance (Monoid w, MonadSouffleFileIO m) => MonadSouffleFileIO (WriterT w m) where
  loadFiles prog = lift . loadFiles prog
  {-# INLINABLE loadFiles #-}
  writeFiles prog = lift . writeFiles prog
  {-# INLINABLE writeFiles #-}

instance MonadSouffleFileIO m => MonadSouffleFileIO (StateT s m) where
  loadFiles prog = lift . loadFiles prog
  {-# INLINABLE loadFiles #-}
  writeFiles prog = lift . writeFiles prog
  {-# INLINABLE writeFiles #-}

instance (MonadSouffleFileIO m, Monoid w) => MonadSouffleFileIO (RWST r w s m) where
  loadFiles prog = lift . loadFiles prog
  {-# INLINABLE loadFiles #-}
  writeFiles prog = lift . writeFiles prog
  {-# INLINABLE writeFiles #-}

instance MonadSouffleFileIO m => MonadSouffleFileIO (ExceptT s m) where
  loadFiles prog = lift . loadFiles prog
  {-# INLINABLE loadFiles #-}
  writeFiles prog = lift . writeFiles prog
  {-# INLINABLE writeFiles #-}