{-# HLINT ignore "Redundant flip" #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# 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 Text
  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 =
  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 (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 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
  , -- _variableNames :: [Text],
    _headVariables :: [RuleElement]
  , _bodyConstraints :: [BodyConstraint]
  , _db :: NaiveDatabase
  }

maybeConstant :: RuleElement -> Maybe Constant
maybeConstant (RuleElementConstant constant) = Just constant
maybeConstant _ = Nothing

withFactsAndRules :: [Text] -> [Text] -> NaiveDatabase
withFactsAndRules facts rules =
  foldr 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 :: NaiveDatabase -> Relation -> Literal -> RuleContext
  digestHead db relation (Literal neg relationName terms) =
    RuleContext
      { __relation = relation
      , _headVariables = variables'
      , _bodyConstraints = []
      , _db = NaiveDatabase relationMap' constants'
      }
   where
    relationMap :: Map RelationId Relation = relations db
    relationMap' = Map.insert relationName relation relationMap
    extraVariables = toElement <$> terms
    variables' = nub extraVariables
    extraConstants = catMaybes $ maybeConstant <$> variables'
    constants' = Set.union (constants db) $ 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 Set.empty
    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 (Var name) = ConstraintElementIndex (varIndex name)
    toConstraint constant = ConstraintElementConstant constant
  addRule :: (Literal, [Literal]) -> NaiveDatabase  ->NaiveDatabase
  addRule (ruleHead, body) db  =
    NaiveDatabase newRelationMap 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'' = foldl digestBody context' body
    db' = _db context''
    relationMap = relations db'
    variables' = _headVariables context''
    extractVarName :: RuleElement -> Maybe Text
    extractVarName (RuleElementVariable varName) = Just varName
    extractVarName (RuleElementConstant constant) = Nothing
    newRule =
      RelationRule
        { headVariables = catMaybes $ extractVarName <$> variables'
        , bodyElements = fromBodyConstraint <$> (_bodyConstraints context'')
        }
     where
      fromBodyConstraint :: BodyConstraint -> RuleBodyElement
      fromBodyConstraint (BodyConstraint subRelation elements) =
        RuleBodyElement
          { _subRelationId = _name subRelation
          , _ruleElements = toRuleElement <$> elements
          }
      toRuleElement :: ConstraintElement -> RuleElement
      toRuleElement (ConstraintElementConstant constant) = RuleElementConstant constant
      toRuleElement (ConstraintElementIndex index) = variables' !! index
    relation' =
      Relation
        { _name = _name relation
        , _arity = newArity
        , _tuples = _tuples relation
        , _rules = (_rules relation) ++ [newRule]
        }
    newRelationMap = Map.insert relationName relation' relationMap
    constants' = constants db'
  toElement :: Term -> RuleElement
  toElement (Var name) = RuleElementVariable name
  toElement constant = RuleElementConstant constant

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

instance Exception NaiveDatabaseException