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

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)

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 Int
  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 =
  foldr 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 :: Literal -> NaiveDatabase -> NaiveDatabase
  addFact (Literal neg relationName terms) (NaiveDatabase relationMap constantSet) =
    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 BodyConstraint = BodyConstraint
  { _subRelation :: Relation
  , _elements :: [RuleElement]
  }

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

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 = 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 =
    RuleContext
      { __relation = __relation context
      , _variableNames = variableNames
      , _headEntries = _headEntries context
      , _bodyConstraints = newConstraint : constraints
      , _db = NaiveDatabase relationMap' constants'
      }
   where
    variableNames = nub $ _variableNames context ++ extractVariableNames subTerms
    newArity = length subTerms
    subRelation = lookupRelation subRelationName relationMap 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 = lookupRelation relationName (relations db) newArity Set.empty
    context = digestHead db relation ruleHead
    context' = foldr digestBody context body
    db' = _db context'
    relationMap = relations db'
    newRule =
      RelationRule
        { headVariables = _variableNames context'
        , bodyElements = toRuleBodyElement <$> _bodyConstraints context'
        } where
        toRuleBodyElement :: BodyConstraint -> RuleBodyElement
        toRuleBodyElement (BodyConstraint subRelation elements) =
          RuleBodyElement { 
            _subRelationId = _name subRelation
            , _ruleElements = elements
            }
    relation' =
      Relation
        { _name = _name relation
        , _arity = newArity
        , _tuples = _tuples relation
        , _rules = newRule : _rules relation
        }
    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