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

module Datalog.NaiveDatabase 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
import Datalog.DatalogDB

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


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

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

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

  addRule :: (Literal, [Literal]) -> NaiveDatabase -> NaiveDatabase
  addRule (ruleHead, body) db  =
    NaiveDatabase relationMap' constants' where
      relationName = predName ruleHead
      terms = arguments ruleHead
      newArity = length terms
      relation = lookupRelation0 relationName db newArity Set.empty
      context = digestHead db relation ruleHead
      context' = foldr digestBody context body
      db' = _db context'
      relationMap = relations db'
      relation' = appendRule relation RelationRule { 
          headVariables = _variableNames context'
          , bodyElements = toRuleBodyElement <$> _bodyConstraints context'
        }
      relationMap' = Map.insert relationName relation' relationMap
      constants' = constants db'

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

lookupRelation0 :: DatalogDB db => Text -> db -> Int -> Set [Constant] -> Relation
lookupRelation0 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

lookupRelation00 :: DatalogDB db => 
  Text -> db -> Int -> (Relation -> Relation) -> db
lookupRelation00 relationName db newArity update =
  insertRelation db (update newRelation) 
    where
    newRelation = case lookupRelation db relationName of
      Nothing -> Relation relationName newArity Set.empty []
      Just relation ->
        if _arity relation == newArity then
            relation
          else throw $ BadArityException relationName newArity

lookupRelation000:: DatalogDB db => 
  Text -> db -> Int -> Set [Constant] -> (Relation -> Relation) -> db
lookupRelation000 relationName db newArity tuples update =
  lookupRelation00 relationName db newArity \relation ->
    update relation { 
      _tuples = Set.union tuples $ _tuples relation 
    }

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

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

-- | equivalent(Q,Q) :- . |
-- could be equivalent(Q, 3, 'zzz, Q, R)
-- terms = Var Q, Num 3, Sym zzz, Var Q, Var R
-- want to convert this to:
-- (need constants 3, 'zzz)
-- entries = [RuleElement] = (RuleElement 0), RuleElement Num 3, RuleElement Sym zzz, (RuleElement 0), (RuleElement 1)
-- variableNames = ["Q" "R"]

digestHead :: forall db . (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 :: forall db . (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 = lookupRelation0 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

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

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