{-# 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

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)
    }

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 
    }
    

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

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 BodyConstraint = BodyConstraint
  { _subRelation :: Relation
  , _elements :: [RuleElement]
  }

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

appendRule :: Relation -> RelationRule -> Relation
appendRule relation rule =
  relation {
    _rules = rule : (_rules relation)
  }
  -- Relation { _name = _name relation
  -- , _arity = _arity relation
  -- , _tuples = _tuples relation
  -- , _rules = rule : (_rules relation)
  -- } 

toRuleBodyElement :: BodyConstraint -> RuleBodyElement
toRuleBodyElement (BodyConstraint subRelation elements) =
  RuleBodyElement { 
    _subRelationId = _name subRelation
    , _ruleElements = elements
    }

withFactsAndRules :: [Text] -> [Text] -> NaiveDatabase
withFactsAndRules facts = foldr (addRule . extractRule) (withFacts facts)
 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 :: NaiveDatabase -> Relation -> Literal -> RuleContext
digestHead db relation (Literal neg relationName terms) =
  RuleContext
    { __relation = relation
    , _variableNames = variableNames
    , _headEntries = entries'
    , _bodyConstraints = []
    , _db = NaiveDatabase relationMap' constants'
    }
 where
  variableNames = nub $ extractVariableNames terms
  headTermToElement :: Term -> RuleElement
  headTermToElement (Var name) = 
    RuleElementVariable $ lookupVariable name variableNames
  headTermToElement constant = RuleElementConstant constant
  relationMap :: Map RelationId Relation = relations db
  relationMap' = Map.insert relationName relation relationMap
  extraVariables = headTermToElement <$> terms
  entries' = nub extraVariables
  extraConstants = mapMaybe extractConstant entries' where
    extractConstant :: RuleElement -> Maybe Constant
    extractConstant (RuleElementConstant constant) = Just constant
    extractConstant _ = Nothing
  constants' = Set.union (constants db) $ Set.fromList extraConstants
digestBody :: Literal -> RuleContext -> RuleContext
digestBody (Literal neg subRelationName subTerms) context =
  context { 
    _variableNames = variableNames
    , _bodyConstraints = newConstraint : constraints
    , _db = NaiveDatabase relationMap' constants'
    }
 where
  variableNames = nub $ _variableNames context ++ extractVariableNames subTerms
  newArity = length subTerms
  subRelation = lookupRelation0 subRelationName (_db context) newArity Set.empty
  relationMap :: Map RelationId Relation = relations (_db context)
  relationMap' = Map.insert subRelationName subRelation relationMap
  extraConstants = mapMaybe constantFromTerm subTerms where
    constantFromTerm :: Term -> Maybe Constant
    constantFromTerm (Var _) = Nothing
    constantFromTerm constant = Just constant
  constants' = Set.union (constants (_db context)) $ 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

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'

extractVariableNames :: [Term] -> [Text]
extractVariableNames = mapMaybe extractVariableName where
  extractVariableName :: Term -> Maybe Text
  extractVariableName (Var name) = Just name
  extractVariableName _ = Nothing

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 [Text]
  | UnexpectedConstantException Constant
  deriving (Show)

instance Exception NaiveDatabaseException