content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE FlexibleContexts #-}
  module Main ( main ) where
  
  import qualified Control.Monad.Catch as E
  import Control.Monad.Trans.Class ( lift )
  import Control.Monad.Trans.State.Strict ( evalStateT, StateT, modify, gets )
  import qualified Data.Foldable as F
  import qualified Data.List as L
  import qualified Data.Map as M
  import Data.Maybe ( catMaybes )
  import Data.Sequence ( Seq, (|>) )
  import qualified Data.Sequence as Seq
  import qualified Data.Text as T
  import Data.Typeable ( Typeable )
  import qualified System.Console.Haskeline as HL
  import Text.Printf ( printf )
  
  import Database.Datalog
  
  import qualified Parser as P
  import qualified Commands as C
  
  main :: IO ()
  main = evalStateT (HL.runInputT settings loop) s0
    where
      settings = HL.defaultSettings
      s0 = ReplState { commands = Seq.empty
                     , definedRelations = M.empty
                     }
  
  type ReplM = StateT ReplState IO
  
  data ReplState = ReplState { commands :: !(Seq C.Command)
                             , definedRelations :: !(M.Map String Int)
                             }
  
  -- Each time a (non-query) command is entered, just record it in the
  -- list.  If the command is a query, then "interpret" the whole list
  -- of commands as a DatabaseBuilder action and run it to produce a
  -- Database.  Then execute the query.  Displaying facts is just a
  -- filter over the list of commands.
  
  loop :: HL.InputT ReplM ()
  loop = do
    minput <- HL.getInputLine "% "
    case minput of
      Nothing -> return ()
      Just input -> do
        let cmd = P.parseInput input
        case cmd of
          Left err -> HL.outputStrLn (show err) >> loop
          Right C.DumpFacts -> do
            cs <- lift $ gets commands
            F.forM_ cs $ \c ->
              case c of
                C.AssertFact cl -> HL.outputStrLn (clauseString cl)
                _ -> return ()
            loop
          Right C.DumpRules -> do
            cs <- lift $ gets commands
            F.forM_ cs $ \c -> do
              case c of
                C.AddRule ruleHead ruleBody ->
                  HL.outputStrLn (ruleString ruleHead ruleBody)
                _ -> return ()
            loop
          Right (C.Query qc@(C.Clause name _)) -> do
            erows <- lift $ E.try (evaluateQuery qc)
            case erows of
              Left err -> do
                let errAs :: EvaluationError
                    errAs = err
                HL.outputStrLn (show errAs)
                loop
              Right rows -> do
                F.forM_ rows $ \row -> do
                  let s = L.intercalate ", " row
                  HL.outputStrLn $ printf "%s(%s)" name s
                loop
          Right C.Quit -> return ()
          Right C.Help -> printHelp >> loop
          Right c@(C.AssertFact f) -> do
            ok <- guardArity f
            case ok of
              Nothing ->
                lift $ modify $ \s -> s { commands = commands s |> c }
              Just err -> HL.outputStrLn err
            loop
          Right c@(C.AddRule ruleHead ruleBody) -> do
            hres <- guardArity ruleHead
            bress <- mapM guardArity ruleBody
            case catMaybes (hres : bress) of
              [] -> lift $ modify $ \s -> s { commands = commands s |> c }
              errs -> HL.outputStrLn (unlines errs)
            loop
  
  guardArity :: (Show a) => C.Clause a -> HL.InputT ReplM (Maybe String)
  guardArity f@(C.Clause name args) = do
    rels <- lift $ gets definedRelations
    case M.lookup name rels of
      Just arity | length args == arity -> return Nothing
                 | otherwise ->
                   return $ Just ("Arity mismatch: " ++ show f ++ " should have arity " ++ show arity)
      Nothing -> do
        lift $ modify $ \s -> s { definedRelations = M.insert name (length args) (definedRelations s) }
        return Nothing
  
  printHelp :: HL.InputT ReplM ()
  printHelp =
    HL.outputStrLn $ unlines [ "Datalog REPL"
                             , ""
                             , "Commands"
                             , "  :help - this text"
                             , "  :quit - exit the repl"
                             , "  :facts - print all defined facts"
                             , "  :rules - print all defined rules"
                             , ""
                             , "Syntax"
                             , "  To declare a fact:"
                             , "    relation1(arg1, arg2)."
                             , "  To define a rule:"
                             , "    relation2(X, Y) :- relation1(X, Z), relation1(Z, Y)."
                             , "  To issue a query:"
                             , "    relation2(X, Y)?"
                             , ""
                             , "  Variables are in all caps.  Literals (atoms) begin with a lowercase letter.  Relation names also begin with a lowercase letter."
                             ]
  
  ruleString :: C.Clause C.AnyValue -> [C.Clause C.AnyValue] -> String
  ruleString ruleHead ruleBody =
    concat [ cstring ruleHead
           , " :- "
           , L.intercalate ", " (map cstring ruleBody)
           ]
    where
      cstring (C.Clause name args) =
        let strs = L.intercalate ", " $ map valToString args
        in printf "%s(%s)" name strs
      valToString (C.AVVariable s) = s
      valToString (C.AVLiteral (C.LVString s)) = s
  
  clauseString :: C.Clause C.LiteralValue -> String
  clauseString (C.Clause name lits) = printf "%s(%s)" name strs
    where
      strs = L.intercalate ", " $ map litToString lits
  
  litToString :: C.LiteralValue -> String
  litToString (C.LVString s) = s
  
  pleatM :: (Monad m, F.Foldable f) => a -> f b -> (a -> b -> m a) -> m a
  pleatM seed elts f = F.foldlM f seed elts
  
  evaluateQuery :: C.Clause C.AnyValue -> StateT ReplState IO [[String]]
  evaluateQuery (C.Clause name vals) = do
    cs <- gets commands
    db <- makeDatabase $ do
      _ <- pleatM M.empty cs $ \ !a c -> do
        case c of
          C.AssertFact fact@(C.Clause rel factVals) ->
            case M.lookup rel a of
              Nothing -> do
                let arity = length factVals
                r <- addRelation (T.pack rel) arity
                assertFact r (map litToString factVals)
                lift $ modify $ \s -> s { definedRelations = M.insert rel arity (definedRelations s) }
                return $ M.insert rel (r, arity) a
              Just (r, arity) | arity == length factVals -> do
                assertFact r (map litToString factVals)
                return a
              Just (_, arity) -> E.throwM $ ArityMismatch arity fact
          _ -> return a
      return ()
    queryDatabase db $ do
      _ <- pleatM M.empty cs $ \ !a c -> do
        case c of
          C.AddRule h@(C.Clause headRel headVals) body -> do
            a1 <- checkArityDefs a h
            a2 <- F.foldlM checkArityDefs a1 body
            hr <- inferencePredicate (T.pack headRel)
            let headTerms = map toTerm headVals
                bodies = map toBodyClause body
            assertRule (hr, headTerms) bodies
            return a2
          _ -> return a
      qrel <- inferencePredicate (T.pack name)
      issueQuery qrel (map toTerm vals)
  
  toBodyClause :: C.Clause C.AnyValue -> QueryBuilder ReplM String (Literal Clause String)
  toBodyClause c@(C.Clause rel vals) = do
    checkArity c
    r <- inferencePredicate (T.pack rel)
    lit r (map toTerm vals)
  
  toTerm :: C.AnyValue -> Term String
  toTerm (C.AVVariable v) = LogicVar (T.pack v)
  toTerm (C.AVLiteral (C.LVString l)) = Atom l
  
  checkArityDefs :: M.Map String (Relation, Int)
                    -> C.Clause C.AnyValue
                    -> QueryBuilder ReplM String (M.Map String (Relation, Int))
  checkArityDefs defs c@(C.Clause rel vals) = do
    checkArity c
    case M.lookup rel defs of
      Nothing -> do
        r <- inferencePredicate (T.pack rel)
        return $ M.insert rel (r, length vals) defs
      Just (_, arity) | arity == length vals -> return defs
                      | otherwise -> E.throwM $ ArityMismatch2 arity c
  
  checkArity :: C.Clause C.AnyValue -> QueryBuilder ReplM String ()
  checkArity c@(C.Clause rel vals) = do
    rs <- lift $ gets definedRelations
    case M.lookup rel rs of
      Just arity | carity == arity -> return ()
                 | otherwise -> E.throwM $ ArityMismatch2 arity c
      Nothing -> lift $ modify $ \s -> s { definedRelations = M.insert rel carity (definedRelations s) }
    where
      carity = length vals
  
  data EvaluationError = ArityMismatch Int (C.Clause C.LiteralValue)
                       | ArityMismatch2 Int (C.Clause C.AnyValue)
                       deriving (Eq, Ord, Show, Typeable)
  
  instance E.Exception EvaluationError
  