module Data.JLD.Flattening.NodeMap (NodeMap, BNMParams (..), buildNodeMap) where

import Data.JLD.Prelude

import Data.JLD.Control.Monad.RES (REST, execREST, withErrorRES')
import Data.JLD.Error (JLDError (..))
import Data.JLD.Model.IRI (isBlankIri)
import Data.JLD.Model.Keyword (Keyword (..), isNotKeyword)
import Data.JLD.Model.NodeMap (NodeMap)
import Data.JLD.Model.NodeMap qualified as N (hasKey2, hasKey3, insert, lookup3, memberArray, modifyArray)
import Data.JLD.Model.NodeObject (isNodeObject)
import Data.JLD.Util (valueIsScalar, valueToArray, valueToNonNullArray)

import Control.Monad.Except (MonadError (..))
import Data.Aeson (Array, Object, Value (..))
import Data.Aeson.Key qualified as K (toText)
import Data.Aeson.KeyMap qualified as KM (filterWithKey, insert, lookup, member, singleton)
import Data.Foldable.WithIndex (iforM_)
import Data.Map.Strict qualified as M (insert, lookup)
import Data.Vector qualified as V (singleton, snoc, uniq)

type BNMT e m = REST BNMEnv (Either (JLDError e) ()) BNMState m

data BNMEnv = BNMEnv
    { bnmEnvActiveGraph :: Text
    , bnmEnvActiveSubject :: Maybe Text
    , bnmEnvActiveProperty :: Maybe Text
    , bnmEnvReferenceNode :: Maybe Object
    }
    deriving (Show)

data BNMState = BNMState
    { bnmStateNodeMap :: NodeMap
    , bnmStateList :: Maybe Array
    , bnmStateIdentifierCounter :: Int
    , bnmStateIdentifierMap :: Map Text Text
    }
    deriving (Show, Eq)

data BNMParams = BNMParams
    { bnmParamsNodeMap :: NodeMap
    , bnmParamsActiveGraph :: Text
    , bnmParamsActiveSubject :: Maybe Text
    , bnmParamsActiveProperty :: Maybe Text
    , bnmParamsList :: Maybe Array
    , bnmParamsReferenceNode :: Maybe Object
    }
    deriving (Show, Eq)

listToObject :: Maybe Array -> Value
listToObject = Object <. KM.singleton (show KeywordList) <. Array <. fromMaybe mempty

bnmModifyNodeMap :: Monad m => (NodeMap -> NodeMap) -> BNMT e m ()
bnmModifyNodeMap fn = modify \s -> s{bnmStateNodeMap = fn (bnmStateNodeMap s)}

bnmModifyList :: Monad m => (Maybe Array -> Maybe Array) -> BNMT e m ()
bnmModifyList fn = modify \s -> s{bnmStateList = fn (bnmStateList s)}

bnmModifyIdentifierCounter :: Monad m => (Int -> Int) -> BNMT e m ()
bnmModifyIdentifierCounter fn = modify \s -> s{bnmStateIdentifierCounter = fn (bnmStateIdentifierCounter s)}

bnmModifyIdentifierMap :: Monad m => (Map Text Text -> Map Text Text) -> BNMT e m ()
bnmModifyIdentifierMap fn = modify \s -> s{bnmStateIdentifierMap = fn (bnmStateIdentifierMap s)}

bnmCreateIdentifier :: Monad m => Maybe Text -> BNMT e m Text
bnmCreateIdentifier Nothing = do
    n <- gets bnmStateIdentifierCounter
    "_:b" <> show n <$ bnmModifyIdentifierCounter (const <| n + 1)
bnmCreateIdentifier (Just identifier) =
    gets (bnmStateIdentifierMap .> M.lookup identifier) >>= \case
        Just nodeId -> pure nodeId
        Nothing -> do
            nodeId <- bnmCreateIdentifier Nothing
            nodeId <$ bnmModifyIdentifierMap (M.insert identifier nodeId)

bnmBuildNodeMap :: Monad m => Value -> (BNMParams -> BNMParams) -> BNMT e m (Maybe Array)
bnmBuildNodeMap value paramsFn = do
    BNMEnv{..} <- ask
    nodeMap <- gets bnmStateNodeMap
    let params p =
            paramsFn
                p
                    { bnmParamsNodeMap = nodeMap
                    , bnmParamsActiveGraph = bnmEnvActiveGraph
                    , bnmParamsActiveSubject = bnmEnvActiveSubject
                    , bnmParamsActiveProperty = bnmEnvActiveProperty
                    }
    (nodeMap', list) <- buildNodeMap value params
    bnmModifyNodeMap <| const nodeMap'
    pure list

buildNodeMap' :: Monad m => Value -> BNMT e m ()
buildNodeMap' element = case element of
    -- 1.
    Array elementArray -> forM_ elementArray (buildNodeMap' .> withErrorRES' (either (Left .> throwError) pure))
    -- 2.
    Object elementObject -> do
        BNMEnv{..} <- ask

        -- 3.
        elementObject' <- case KM.lookup (show KeywordType) elementObject of
            Just type' -> do
                types <-
                    Array <$> forM (valueToArray type') \case
                        String item | isBlankIri item -> String <$> bnmCreateIdentifier (Just item)
                        item -> pure item
                pure <| KM.insert (show KeywordType) types elementObject
            --
            Nothing -> pure elementObject

        if
                -- 4.
                | KM.member (show KeywordValue) elementObject' ->
                    gets bnmStateList >>= \case
                        -- 4.1.
                        Nothing ->
                            gets (bnmStateNodeMap .> N.lookup3 bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty) >>= \case
                                -- 4.1.1.
                                Just (Array activePropertyValue)
                                    | notElem element activePropertyValue ->
                                        bnmModifyNodeMap
                                            <. N.insert bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty
                                            <. Array
                                            <| V.snoc activePropertyValue element
                                    | otherwise -> pure ()
                                -- 4.2.2
                                _ ->
                                    bnmModifyNodeMap
                                        <. N.insert bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty
                                        <. Array
                                        <| V.singleton element
                        -- 4.2.
                        Just list -> bnmModifyList <. const <. Just <| V.snoc list element
                -- 5.
                | Just elemList <- KM.lookup (show KeywordList) elementObject' -> do
                    -- 5.1. 5.2.
                    subList <- listToObject <$> bnmBuildNodeMap elemList id

                    gets bnmStateList >>= \case
                        -- 5.3.
                        Nothing ->
                            bnmModifyNodeMap
                                <. N.modifyArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty
                                <| flip V.snoc subList
                        -- 5.4.
                        Just stateList -> bnmModifyList <. const <. Just <| V.snoc stateList subList
                -- 6.
                | isNodeObject (Object elementObject') -> do
                    id' <- case KM.lookup (show KeywordId) elementObject' of
                        -- 6.1.
                        Just (String idValue) | isBlankIri idValue -> bnmCreateIdentifier <| Just idValue
                        Just _ -> throwError <| Right ()
                        -- 6.2.
                        Nothing -> bnmCreateIdentifier Nothing

                    -- 6.3.
                    nodeMap <- gets bnmStateNodeMap
                    when (N.hasKey2 bnmEnvActiveGraph (Just id') nodeMap) do
                        bnmModifyNodeMap <| N.insert bnmEnvActiveGraph (Just id') (Just <| show KeywordId) (String id')

                    nodeMap' <- gets bnmStateNodeMap
                    if
                            -- 6.5.
                            | Just referenceNode <- bnmEnvReferenceNode ->
                                unless (N.memberArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty (Object referenceNode) nodeMap') do
                                    bnmModifyNodeMap
                                        <. N.modifyArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty
                                        <| flip V.snoc (Object referenceNode)
                            -- 6.6.
                            | isJust bnmEnvActiveProperty -> do
                                -- 6.6.1.
                                let reference = Object <| KM.singleton (show KeywordId) (String id')

                                gets bnmStateList >>= \case
                                    -- 6.6.2.
                                    Nothing ->
                                        unless (N.memberArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty reference nodeMap') do
                                            bnmModifyNodeMap
                                                <. N.modifyArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty
                                                <| flip V.snoc reference
                                    -- 6.6.3.
                                    Just stateList -> bnmModifyList <. const <. Just <| V.snoc stateList reference
                            --
                            | otherwise -> pure ()

                    -- 6.7.
                    case KM.lookup (show KeywordType) elementObject' of
                        Just typeValue -> do
                            nodeType <-
                                Array
                                    <. V.uniq
                                    <. (<> valueToNonNullArray typeValue)
                                    <. fromMaybe mempty
                                    <. fmap valueToNonNullArray
                                    <$> gets (bnmStateNodeMap .> N.lookup3 bnmEnvActiveGraph (Just id') (Just <| show KeywordType))
                            bnmModifyNodeMap <| N.insert bnmEnvActiveGraph (Just id') (Just <| show KeywordType) nodeType
                        --
                        Nothing -> pure ()

                    -- 6.8.
                    nodeMap'' <- gets bnmStateNodeMap
                    case KM.lookup (show KeywordIndex) elementObject' of
                        Just indexValue
                            | N.hasKey3 bnmEnvActiveGraph (Just id') (Just <| show KeywordIndex) nodeMap'' -> throwError <| Left ConflictingIndexes
                            | otherwise -> bnmModifyNodeMap <| N.insert bnmEnvActiveGraph (Just id') (Just <| show KeywordType) indexValue
                        --
                        Nothing -> pure ()

                    -- 6.9.
                    case KM.lookup (show KeywordReverse) elementObject' of
                        Just (Object reverseMap) -> do
                            -- 6.9.1.
                            let referenced = KM.singleton (show KeywordId) (String id')

                            -- 6.9.3.
                            iforM_ reverseMap \key ->
                                valueToArray .> mapM_ \value -> do
                                    void <| bnmBuildNodeMap value \params ->
                                        params
                                            { bnmParamsReferenceNode = Just referenced
                                            , bnmParamsActiveProperty = Just <| K.toText key
                                            }
                                    pure ()
                        --
                        _ -> pure ()

                    -- 6.10.
                    case KM.lookup (show KeywordGraph) elementObject' of
                        Just graphValue ->
                            void <| bnmBuildNodeMap graphValue \params -> params{bnmParamsActiveGraph = id'}
                        --
                        _ -> pure ()

                    -- 6.11.
                    case KM.lookup (show KeywordIncluded) elementObject' of
                        Just includedValue ->
                            void <| bnmBuildNodeMap includedValue id
                        --
                        _ -> pure ()

                    let filteredKeywords = [KeywordId, KeywordType, KeywordIndex, KeywordReverse, KeywordGraph, KeywordIncluded]
                        elementObject'' = KM.filterWithKey (\key _ -> isNotKeyword (K.toText key) filteredKeywords) elementObject'

                    -- 6.12.
                    iforM_ elementObject'' \(K.toText -> property) value ->
                        if value == Null || valueIsScalar value
                            then pure ()
                            else do
                                -- 6.12.1.
                                property' <-
                                    if isBlankIri property
                                        then bnmCreateIdentifier <| Just property
                                        else pure property

                                -- 6.12.2.
                                gets (bnmStateNodeMap .> N.hasKey3 bnmEnvActiveGraph (Just id') (Just property')) >>= \case
                                    True -> pure ()
                                    False -> bnmModifyNodeMap <| N.insert bnmEnvActiveGraph (Just id') (Just property') (Array mempty)

                                void <| bnmBuildNodeMap value \params ->
                                    params
                                        { bnmParamsActiveSubject = Just id'
                                        , bnmParamsActiveProperty = Just property
                                        }

                --
                | otherwise -> pure ()
        pure ()
    --
    _ -> pure ()

buildNodeMap :: Monad m => Value -> (BNMParams -> BNMParams) -> m (NodeMap, Maybe Array)
buildNodeMap document paramsFn = do
    BNMState{..} <- buildNodeMap' document |> execREST env st
    pure (bnmStateNodeMap, bnmStateList)
  where
    BNMParams{..} =
        paramsFn
            BNMParams
                { bnmParamsNodeMap = mempty
                , bnmParamsActiveGraph = show KeywordDefault
                , bnmParamsActiveSubject = Nothing
                , bnmParamsActiveProperty = Nothing
                , bnmParamsList = mempty
                , bnmParamsReferenceNode = Nothing
                }

    env =
        BNMEnv
            { bnmEnvActiveGraph = bnmParamsActiveGraph
            , bnmEnvActiveSubject = bnmParamsActiveSubject
            , bnmEnvActiveProperty = bnmParamsActiveProperty
            , bnmEnvReferenceNode = bnmParamsReferenceNode
            }

    st =
        BNMState
            { bnmStateNodeMap = bnmParamsNodeMap
            , bnmStateList = bnmParamsList
            , bnmStateIdentifierCounter = 1
            , bnmStateIdentifierMap = mempty
            }