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

import Data.JLD.Prelude

import Data.JLD.Control.Monad.RES (REST, runREST, withErrorRES')
import Data.JLD.Error (JLDError (..))
import Data.JLD.Flattening.Global (JLDFlatteningT)
import Data.JLD.Model.IRI (isBlankIri)
import Data.JLD.Model.Keyword (Keyword (..), isKeywordLike, isNotKeyword)
import Data.JLD.Model.NodeMap (NodeMap, PropertyMap)
import Data.JLD.Model.NodeMap qualified as N (hasKey2, hasKey3, insert, lookup2, 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 (FoldableWithIndex (..), iforM_)
import Data.Map.Strict qualified as M (insert, lookup)
import Data.Vector qualified as V (singleton, snoc, uniq)
import Debug.Pretty.Simple (pTraceShowM)

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
    , bnmParamsIdentifierCounter :: Int
    , bnmParamsIdentifierMap :: Map Text Text
    }
    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
    identifierCounter <- gets bnmStateIdentifierCounter
    identifierMap <- gets bnmStateIdentifierMap
    let params p =
            paramsFn
                p
                    { bnmParamsNodeMap = nodeMap
                    , bnmParamsActiveGraph = bnmEnvActiveGraph
                    , bnmParamsIdentifierCounter = identifierCounter
                    , bnmParamsIdentifierMap = identifierMap
                    }
    (nodeMap', list, identifierCounter', identifierMap') <-
        buildNodeMap value params |> runExceptT >=> \case
            Left err -> throwError <| Left err
            Right a -> pure a
    modify \st ->
        st
            { bnmStateNodeMap = nodeMap'
            , bnmStateIdentifierCounter = identifierCounter'
            , bnmStateIdentifierMap = identifierMap'
            }
    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 \params ->
                            params
                                { bnmParamsActiveSubject = bnmEnvActiveSubject
                                , bnmParamsActiveProperty = bnmEnvActiveProperty
                                , bnmParamsList = Just mempty
                                , bnmParamsReferenceNode = bnmEnvReferenceNode
                                }

                    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)
                            | otherwise -> pure idValue
                        Just _ -> throwError <| Right ()
                        -- 6.2.
                        Nothing -> bnmCreateIdentifier Nothing

                    -- 6.3.
                    bnmModifyNodeMap \nodeMap ->
                        if N.hasKey2 bnmEnvActiveGraph (Just id') nodeMap
                            then nodeMap
                            else N.insert bnmEnvActiveGraph (Just id') (Just <| show KeywordId) (String id') nodeMap

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

                                gets bnmStateList >>= \case
                                    -- 6.6.2.
                                    Nothing -> bnmModifyNodeMap \nodeMap ->
                                        if N.memberArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty reference nodeMap
                                            then nodeMap
                                            else N.modifyArray bnmEnvActiveGraph bnmEnvActiveSubject bnmEnvActiveProperty (`V.snoc` reference) nodeMap
                                    -- 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.
                                bnmModifyNodeMap \nodeMap ->
                                    if N.hasKey3 bnmEnvActiveGraph (Just id') (Just property') nodeMap
                                        then nodeMap
                                        else N.insert bnmEnvActiveGraph (Just id') (Just property') (Array mempty) nodeMap

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

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

buildNodeMap :: Monad m => Value -> (BNMParams -> BNMParams) -> JLDFlatteningT e m (NodeMap, Maybe Array, Int, Map Text Text)
buildNodeMap document paramsFn = do
    (result, BNMState{..}) <- buildNodeMap' document |> runREST env st
    case result of
        Left (Left err) -> throwError err
        _ -> pure (bnmStateNodeMap, bnmStateList, bnmStateIdentifierCounter, bnmStateIdentifierMap)
  where
    BNMParams{..} =
        paramsFn
            BNMParams
                { bnmParamsNodeMap = mempty
                , bnmParamsActiveGraph = show KeywordDefault
                , bnmParamsActiveSubject = Nothing
                , bnmParamsActiveProperty = Nothing
                , bnmParamsList = mempty
                , bnmParamsReferenceNode = Nothing
                , bnmParamsIdentifierCounter = 0
                , bnmParamsIdentifierMap = mempty
                }

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

    st =
        BNMState
            { bnmStateNodeMap = bnmParamsNodeMap
            , bnmStateList = bnmParamsList
            , bnmStateIdentifierCounter = bnmParamsIdentifierCounter
            , bnmStateIdentifierMap = bnmParamsIdentifierMap
            }

mergeNodeMaps :: NodeMap -> NodeMap
mergeNodeMaps = foldl' (ifoldl' go) mempty
  where
    go :: Maybe Text -> NodeMap -> PropertyMap -> NodeMap
    go subjectKey result = ifoldl' (go' subjectKey) result'
      where
        result' = case N.lookup2 (show KeywordMerged) subjectKey result of
            Just _ -> result
            Nothing -> N.insert (show KeywordMerged) subjectKey (Just <| show KeywordId) (maybe Null String subjectKey) result

    go' :: Maybe Text -> Maybe Text -> NodeMap -> Value -> NodeMap
    go' subjectKey propertyKey result property
        | propertyKey /= Just (show KeywordType) && maybe False isKeywordLike propertyKey =
            N.insert (show KeywordMerged) subjectKey propertyKey property result
        | otherwise =
            N.insert (show KeywordMerged) subjectKey propertyKey array result
      where
        array =
            Array
                <. (<> valueToArray property)
                <. maybe mempty valueToArray
                <| N.lookup3 (show KeywordMerged) subjectKey propertyKey result