{-# 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 qualified Data.Text as T
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 (Exception)
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 RelationRule = RelationRule {
  headVariables :: [RuleElement],
  body :: [(Relation, [RuleElement])]
} deriving (Show, Eq)

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

-- newtype RelationId = RelationId Text
--   deriving (Eq, Ord, Show)

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

-- newtype Constant = Constant Text
--   deriving (Eq, Ord, Show)

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 terms
        newRelationMap = Map.insert relationName newRelation relationMap
        newConstantSet = Set.union constantSet $ Set.fromList terms

lookupRelation :: RelationId -> Map RelationId Relation -> Int -> [Term] -> Relation
lookupRelation relationName relationMap newArity terms =
    case Map.lookup relationName relationMap of
      Nothing -> Relation newArity (Set.singleton terms) []
      Just relation ->
        if (_arity relation == newArity)
          then Relation (length terms) (Set.singleton terms) []
          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 terms
        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 terms
        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)

        -- toBodyConstraint :: (Relation, [RuleElement]) -> BodyConstraint
        -- toBodyConstraint (subRelation, ruleElements) =
        --   BodyConstraint {
        --     _relation = subRelation,
        --     _elements = toConstraintElement <$> ruleElements
        --   }
        newRule = RelationRule {
          headVariables = variables',
          body = fromBodyConstraint <$> (_bodyConstraints context'')
        } where
          fromBodyConstraint :: BodyConstraint -> (Relation, [RuleElement])
          fromBodyConstraint (BodyConstraint subRelation elements) =
            (subRelation, toRuleElement <$> elements) 
          toRuleElement :: ConstraintElement -> RuleElement
          toRuleElement (ConstraintElementConstant constant) = RuleElementConstant constant
          toRuleElement (ConstraintElementIndex index) = variables' !! index
        -- input is [(Relation, [RuleElement])]
        -- bodyConstraint has
        --   _relation :: Relation,
        --   _elements :: [ConstraintElement]
        relation' = 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]
  deriving (Show)

instance Exception NaiveDatabaseException