{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications  #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Redundant flip" #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Datalog.NaiveDatabase where

import Data.Map (Map)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Datalog.DatalogParser(parseDatalog, Literal(..), Statement(..), Term (..), Head (HeadSingle))
import qualified Data.Map as Map
import Text.Megaparsec (ParseErrorBundle)
import Data.Void
import Control.Exception.Base
import Data.List
import Data.Maybe

data Value =
  ValueInt Int |
  ValueSymbol String

data NaiveDatabase = NaiveDatabase {
  relations :: Map RelationId Relation,
  constants :: Set Constant
}

data RuleElement = -- entry occurring in a head or body relation - constant or variable
  RuleElementConstant Constant |
  RuleElementVariable Text
  deriving (Show, Eq)

data RuleBodyElement = RuleBodyElement {
  _subRelationId :: RelationId,
  _ruleElements :: [RuleElement]
} deriving (Show, Eq)

data RelationRule = RelationRule {
  headVariables :: [Text],
  bodyElements :: [RuleBodyElement]
} deriving (Show, Eq)

data Relation = Relation {
  _name :: RelationId,
  _arity :: Int,
  _tuples :: Set [Constant],
  _rules :: [RelationRule]
} deriving (Show, Eq)

-- Our constants will be the terms of the Datalog grammar - ints/variables/symbols
type Constant = Term
type RelationId = Text

emptyDB :: NaiveDatabase
emptyDB = NaiveDatabase {
  relations = Map.empty,
  constants = Set.empty   -- the Herbrand universe
}

withFacts :: [Text] -> NaiveDatabase
withFacts facts =
  foldl addFact emptyDB (extractFact <$> facts) 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
    addFact :: NaiveDatabase -> Literal -> NaiveDatabase
    addFact (NaiveDatabase relationMap constantSet) (Literal neg relationName terms) =
      NaiveDatabase newRelationMap newConstantSet where
        newArity = length terms
        newRelation = lookupRelation relationName relationMap newArity (Set.singleton terms)
        newRelationMap = Map.insert relationName newRelation relationMap
        newConstantSet = Set.union constantSet $ Set.fromList terms

lookupRelation :: RelationId -> Map RelationId Relation -> Int -> Set [Term] -> Relation
lookupRelation relationName relationMap newArity tuples =
    case Map.lookup relationName relationMap of
      Nothing -> Relation relationName newArity tuples []
      Just relation ->
        if (_arity relation == newArity)
          then
            let newTuples = Set.union tuples $ _tuples relation
            in Relation relationName newArity newTuples [] 
          else throw $ BadArityException relationName newArity

data ConstraintElement = -- entry occurring in a rule body constraint - constant or variable index
  ConstraintElementConstant Constant |
  ConstraintElementIndex Int
  deriving (Show, Eq)

data BodyConstraint = BodyConstraint {
  _relation :: Relation,
  _elements :: [ConstraintElement]
}

data RuleContext = RuleContext {
  __relation :: Relation,
  _headVariables :: [RuleElement],
  _bodyConstraints :: [BodyConstraint],
  _db :: NaiveDatabase
}

toElement :: Term -> RuleElement
-- toElement constant@(Var _) = RuleElementConstant constant
-- toElement constant@(Num _) = RuleElementConstant constant
toElement (Var name) = RuleElementVariable name
toElement constant = RuleElementConstant constant

maybeConstant :: RuleElement -> Maybe Constant
maybeConstant (RuleElementConstant constant) = Just constant
maybeConstant _ = Nothing

withFactsAndRules :: [Text] -> [Text] -> NaiveDatabase
withFactsAndRules facts rules =
  foldl addRule (withFacts facts) (extractRule <$> rules) where
    extractRule:: Text -> (Literal, [Literal])
    extractRule ruleText =
      case (parseDatalog ruleText) of
        Right (Rule (HeadSingle lhs) rhs) -> (lhs, rhs)
        Right otherStatement -> throw $ NonRuleException ruleText otherStatement
        Left ex -> throw $ CannotParseStatementException ruleText ex
    digestHead :: RuleContext -> Literal -> RuleContext
    digestHead  context (Literal neg relationName terms) =
      RuleContext { 
        __relation = relation, 
        _headVariables = variables', 
        _bodyConstraints = _bodyConstraints context, 
        _db = NaiveDatabase relationMap' constants'
      } where
        newArity = length terms
        relation = __relation context
        relationMap :: Map RelationId Relation = (relations (_db context))
        relationMap' = Map.insert relationName relation relationMap
        extraVariables = toElement <$> terms
        extraConstants = catMaybes $ maybeConstant <$> extraVariables
        variables' = nub $ _headVariables context ++ extraVariables
        constants' = Set.union (constants (_db context)) $ Set.fromList extraConstants
    digestBody :: RuleContext -> Literal -> RuleContext
    digestBody context (Literal neg subRelationName terms) =
      RuleContext { 
        __relation = relation, 
        _headVariables = variables', 
        _bodyConstraints = constraints', 
        _db = NaiveDatabase relationMap' constants'
      }  where
        relation = __relation context
        newArity = length terms
        subRelation = lookupRelation subRelationName relationMap newArity Set.empty
        relationMap :: Map RelationId Relation = relations (_db context)
        relationMap' = Map.insert subRelationName subRelation relationMap
        extraVariables = toElement <$> terms
        extraConstants = catMaybes $ maybeConstant <$> extraVariables
        variables' = nub $ _headVariables context ++ extraVariables
        constants' = Set.union (constants (_db context)) $ Set.fromList extraConstants
        constraints = _bodyConstraints context
        newConstraint = BodyConstraint subRelation (toConstraint <$> terms)
        constraints' = constraints ++ [newConstraint]
        varIndex :: Text -> Int
        varIndex name =
          case elemIndex (RuleElementVariable name) variables' of
            Just index -> index
            Nothing -> throw $ VariableLookupException name variables'
        toConstraint :: Term -> ConstraintElement
        -- toConstraint thing = ConstraintElementIndex 0
        toConstraint (Var name) = ConstraintElementIndex (varIndex name)
        toConstraint constant = ConstraintElementConstant constant
    addRule :: NaiveDatabase -> (Literal, [Literal]) -> NaiveDatabase
    addRule  db (ruleHead, body) =
      NaiveDatabase newRelationMap constants' where
        relationName = predName ruleHead
        terms = arguments ruleHead
        newArity = length terms
        relation = lookupRelation relationName (relations (_db context)) newArity Set.empty
        context = RuleContext { 
          __relation = relation, 
          _headVariables = [], 
          _bodyConstraints = [], 
          _db = db
        }
        context' = digestHead context ruleHead
        context'' = foldl digestBody context' body
        db' = _db context''
        relationMap = relations db'
        variables' = _headVariables context''
        varIndex :: Text -> Int -- TODO unify with the above
        varIndex name =
          case elemIndex (RuleElementVariable name) variables' of
            Just index -> index
            Nothing -> throw $ VariableLookupException name variables'
        toConstraintElement :: RuleElement -> ConstraintElement
        toConstraintElement (RuleElementConstant constant) = ConstraintElementConstant constant
        toConstraintElement (RuleElementVariable varName) = ConstraintElementIndex (varIndex varName)
        extractVarName :: RuleElement -> Maybe Text
        extractVarName (RuleElementVariable varName) = Just varName
        extractVarName (RuleElementConstant constant) = Nothing
        newRule = RelationRule {
          headVariables = catMaybes $ extractVarName <$> variables',
          bodyElements = fromBodyConstraint <$> (_bodyConstraints context'')
        } where
          fromBodyConstraint :: BodyConstraint -> RuleBodyElement
          fromBodyConstraint (BodyConstraint subRelation elements) =
            RuleBodyElement {
              _subRelationId = _name subRelation, 
              _ruleElements = toRuleElement <$> elements
            } 
          toRuleElement :: ConstraintElement -> RuleElement
          toRuleElement (ConstraintElementConstant constant) = RuleElementConstant constant
          toRuleElement (ConstraintElementIndex index) = variables' !! index
        relation' = Relation {
          _name = _name relation,
          _arity = newArity,
          _tuples = _tuples relation,
          _rules = (_rules relation) ++ [newRule]
        }
        newRelationMap = Map.insert relationName relation' relationMap
        constants' = constants db'

query :: NaiveDatabase -> 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

data NaiveDatabaseException
  = CannotParseStatementException Text (ParseErrorBundle Text Void) |
    NonFactException Text Statement |
    NonRuleException Text Statement |
    NonQueryException Text Statement |
    BadArityException Text Int |
    VariableLookupException Text [RuleElement] |
    UnexpectedConstantException Constant
  deriving (Show)

instance Exception NaiveDatabaseException