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

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

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) = 
      case (elemIndex name) variableNames of
          Just index -> RuleElementVariable index
          Nothing -> throw $ VariableLookupException 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 maybeConstant entries'
    constants' = Set.union (constants db) $ Set.fromList extraConstants
  digestBody :: Literal -> RuleContext -> RuleContext
  digestBody (Literal neg subRelationName terms) context =
    RuleContext
      { __relation = relation
      , _variableNames = variableNames
      , _headEntries = variables'
      , _bodyConstraints = constraints'
      , _db = NaiveDatabase relationMap' constants'
      }
   where
    relation = __relation context
    variableNames = _variableNames context
    bodyTermToElement :: Term -> RuleElement
    bodyTermToElement (Var name) = 
      case (elemIndex name) variableNames of
          Just index -> RuleElementVariable index
          Nothing -> throw $ VariableLookupException name variableNames
    bodyTermToElement constant = RuleElementConstant constant
    newArity = length terms
    subRelation = lookupRelation subRelationName relationMap newArity Set.empty
    relationMap :: Map RelationId Relation = relations (_db context)
    relationMap' = Map.insert subRelationName subRelation relationMap
    extraVariables = bodyTermToElement <$> terms
    extraConstants = mapMaybe maybeConstant extraVariables
    variables' = nub $ _headEntries context ++ extraVariables
    constants' = Set.union (constants (_db context)) $ Set.fromList extraConstants
    constraints = _bodyConstraints context
    newConstraint = BodyConstraint subRelation variables'
    constraints' = newConstraint : constraints
  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