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

module Datalog.DatalogDB where

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

class DatalogDB db where
  emptyDB :: db
  lookupRelation :: db -> Text -> Maybe Relation
  insertRelation :: db -> Relation -> db
  addConstants :: db -> Set Constant -> db

lookupRelationArity :: DatalogDB db => Text -> db -> Int -> Set [Constant] -> Relation
lookupRelationArity relationName db newArity tuples =
  case lookupRelation db relationName of
    Nothing -> Relation relationName newArity tuples []
    Just relation ->
      if _arity relation == newArity then
          let newTuples = Set.union tuples $ _tuples relation
           in relation { _tuples = newTuples }
        else throw $ BadArityException relationName newArity

addFact :: (DatalogDB db) => Literal -> db -> db
addFact (Literal neg relationName terms) db =
  insertRelation (addConstants db extraConstants) newRelation
  where
  newArity = length terms
  newRelation = lookupRelationArity relationName db newArity (Set.singleton terms)
  extraConstants = Set.fromList terms

addRule :: (DatalogDB db) => (Literal, [Literal]) -> db -> db
addRule (ruleHead, body) db  =
  insertRelation db' newRelation
  where
    relationName = predName ruleHead
    terms = arguments ruleHead
    newArity = length terms
    relation = lookupRelationArity relationName db newArity Set.empty
    context = digestHead db relation ruleHead
    context' = foldr digestBody context body
    db' = _db context'
    newRelation = appendRule relation RelationRule { 
        headVariables = _variableNames context'
        , bodyElements = toRuleBodyElement <$> _bodyConstraints context'
      }

data (DatalogDB db) => RuleContext db = RuleContext
  { __relation :: Relation
  , _variableNames :: [Text]
  , _headEntries :: [RuleElement]
  , _bodyConstraints :: [BodyConstraint]
  , _db :: db
  }

digestHead :: (DatalogDB db) => db -> Relation -> Literal -> RuleContext db
digestHead db relation (Literal neg relationName terms) =
  RuleContext
    { __relation = relation
    , _variableNames = variableNames
    , _headEntries = entries'
    , _bodyConstraints = []
    , _db = insertRelation (addConstants db extraConstants) relation
    }
 where
  variableNames = nub $ extractVariableNames terms
  entries' = nub $ (headTermToElement variableNames) <$> terms
  extraConstants = Set.fromList $ mapMaybe extractConstant entries' where
    extractConstant :: RuleElement -> Maybe Constant
    extractConstant (RuleElementConstant constant) = Just constant
    extractConstant _ = Nothing
digestBody :: (DatalogDB db) => Literal -> RuleContext db -> RuleContext db
digestBody (Literal neg subRelationName subTerms) context =
  context { 
    _variableNames = variableNames
    , _bodyConstraints = newConstraint : constraints
    , _db = insertRelation (addConstants db constants') subRelation
    }
 where
  db = _db context
  variableNames = nub $ _variableNames context ++ extractVariableNames subTerms
  newArity = length subTerms
  subRelation = lookupRelationArity subRelationName (_db context) newArity Set.empty
  extraConstants = mapMaybe constantFromTerm subTerms where
    constantFromTerm :: Term -> Maybe Constant
    constantFromTerm (Var _) = Nothing
    constantFromTerm constant = Just constant
  constants' = Set.fromList extraConstants
  constraints = _bodyConstraints context
  newConstraint = BodyConstraint subRelation subRuleElements where
    subRuleElements = toRuleElement <$> subTerms
    toRuleElement :: Term -> RuleElement
    toRuleElement (Var varName) = RuleElementVariable $ lookupVariable varName variableNames
    toRuleElement constant = RuleElementConstant constant

lookupVariable :: Text -> [Text] -> Int
lookupVariable varName variableNames = 
  case elemIndex varName variableNames of
      Just index -> index
      Nothing -> throw $ VariableLookupException varName variableNames

headTermToElement :: [Text] -> Term -> RuleElement
headTermToElement variableNames (Var name) = 
  RuleElementVariable $ lookupVariable name variableNames
headTermToElement _ constant = RuleElementConstant constant