{-# HLINT ignore "Redundant flip" #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE BlockArguments #-}

module Datalog.InMemoryDB where

import Control.Exception.Base
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Datalog.DatalogParser (Literal (..), Statement (..), parseDatalog)
import Datalog.Rules
import Datalog.DatalogDB

data InMemoryDB = InMemoryDB
  { relations :: Map RelationId Relation
  , constants :: Set Constant
  } deriving (Show, Eq)

instance DatalogDB InMemoryDB where
  emptyDB :: InMemoryDB
  emptyDB = InMemoryDB
    { relations = Map.empty
    , constants = Set.empty -- the Herbrand universe
    }

  lookupRelation :: InMemoryDB -> Text -> Maybe Relation
  lookupRelation db relationName =
    Map.lookup relationName $ relations db

  insertRelation :: InMemoryDB -> Relation -> InMemoryDB
  insertRelation db relation =
    db {
      relations = Map.insert (_name relation) relation (relations db)
    }

  addConstants :: InMemoryDB -> Set Constant -> InMemoryDB
  addConstants db newConstants =
    db {
      constants = Set.union newConstants (constants db)
    }

withFacts :: DatalogDB db => [Text] -> db
withFacts =
  foldr (addFact . extractFact) emptyDB
 where
  extractFact :: Text -> Literal
  extractFact factText =
    case parseDatalog factText of
      Right (Fact fact) -> fact
      Right otherStatement -> throw $ NonFactException factText otherStatement
      Left ex -> throw $ CannotParseStatementException factText ex

withFactsAndRules :: [Text] -> [Text] -> InMemoryDB
withFactsAndRules facts = foldr (addRule . extractRule) (withFacts facts)

query :: forall db . (DatalogDB db) => db -> Text -> Text
query db qText =
  case parseDatalog qText of
    Right (Query texts literals) -> "#NYI"
    Right otherStatement -> throw $ NonQueryException qText otherStatement
    Left ex -> throw $ CannotParseStatementException qText ex