module Data.JLD (
    module Data.JLD.Mime,
    module Data.JLD.Error,
    module Data.JLD.Options,
    JLDExpansionParams (..),
    JLDExpansionState (..),
    JLDFlatteningParams (..),
    expand,
    flatten,
) where

import Data.JLD.Prelude

import Data.JLD.Control.Monad.RES (evalREST, runREST)
import Data.JLD.Error
import Data.JLD.Expansion (JLDEParams (..))
import Data.JLD.Expansion qualified as E (expand)
import Data.JLD.Expansion.Context (buildActiveContext)
import Data.JLD.Expansion.Global (JLDExpansionEnv (..), JLDExpansionState (..))
import Data.JLD.Flattening qualified as F (flatten)
import Data.JLD.Mime
import Data.JLD.Model.ActiveContext (ActiveContext (..), newActiveContext)
import Data.JLD.Model.Keyword (Keyword (..))
import Data.JLD.Model.URI (uriToIri)
import Data.JLD.Options
import Data.JLD.Util (flattenSingletonArray, valueToArray)

import Data.Aeson (Value (..))
import Data.Aeson.KeyMap qualified as KM (lookup, size)
import Text.URI (URI)

data JLDExpansionParams e m = JLDExpansionParams
    { jldExpansionParamsDocumentLoader :: DocumentLoader e m
    , jldExpansionParamsProcessingMode :: JLDVersion
    , jldExpansionParamsMaxRemoteContexts :: Int
    , jldExpansionParamsExpandContext :: Maybe Value
    , jldExpansionParamsFrameExpansion :: Bool
    , jldExpansionParamsState :: JLDExpansionState
    }
    deriving (Show)

expand :: Monad m => Value -> URI -> (JLDExpansionParams () m -> JLDExpansionParams e m) -> m (Either (JLDError e) Value, JLDExpansionState)
expand document baseUrl paramsFn = do
    let JLDExpansionParams{..} =
            paramsFn
                JLDExpansionParams
                    { jldExpansionParamsDocumentLoader = DocumentLoader <. const <. pure <| Left ()
                    , jldExpansionParamsProcessingMode = JLD1_1
                    , jldExpansionParamsMaxRemoteContexts = 20
                    , jldExpansionParamsExpandContext = Nothing
                    , jldExpansionParamsFrameExpansion = False
                    , jldExpansionParamsState =
                        JLDExpansionState
                            { jldExpansionStateContextCache = mempty
                            , jldExpansionStateDocumentCache = mempty
                            }
                    }

        env =
            JLDExpansionEnv
                { jldExpansionEnvDocumentLoader = jldExpansionParamsDocumentLoader
                , jldExpansionEnvProcessingMode = jldExpansionParamsProcessingMode
                , jldExpansionEnvMaxRemoteContexts = jldExpansionParamsMaxRemoteContexts
                }

        activeContext = newActiveContext \ac -> ac{activeContextBaseUrl = Just baseUrl, activeContextBaseIri = uriToIri baseUrl}
        expansionParams p = p{jldeParamsFrameExpansion = jldExpansionParamsFrameExpansion}

    -- 6.
    let maybeExpandContext =
            jldExpansionParamsExpandContext <&> flattenSingletonArray .> \case
                Array expandedContext -> Array expandedContext
                (Object expandedContext) | Just ctx <- KM.lookup (show KeywordContext) expandedContext -> ctx
                expandedContext -> Array <| pure expandedContext

    activeContext' <- case maybeExpandContext of
        Just expandContext ->
            buildActiveContext activeContext expandContext (Just baseUrl) id
                |> evalREST env jldExpansionParamsState
                |> fmap (fromRight activeContext)
        Nothing -> pure activeContext

    -- 8.
    (result, state') <-
        E.expand activeContext' document baseUrl expansionParams
            |> runREST env jldExpansionParamsState

    let result' = case result of
            -- 8.1.
            Right (Object expanded)
                | KM.size expanded == 1
                , Just expanded' <- KM.lookup (show KeywordGraph) expanded ->
                    Right <. Array <| valueToArray expanded'
            -- 8.2.
            Right Null -> Right <| Array mempty
            -- 8.3.
            Right expanded -> Right <. Array <| valueToArray expanded
            --
            Left err -> Left err

    pure (result', state')

data JLDFlatteningParams e m = JLDFlatteningParams
    { jldFlatteningParamsDocumentLoader :: DocumentLoader e m
    , jldFlatteningParamsProcessingMode :: JLDVersion
    , jldFlatteningParamsMaxRemoteContexts :: Int
    , jldFlatteningParamsExpandContext :: Maybe Value
    , jldFlatteningParamsFrameExpansion :: Bool
    , jldFlatteningParamsState :: JLDExpansionState
    , jldFlatteningParamsCompactArrays :: Bool
    }
    deriving (Show)

flatten :: Monad m => Value -> URI -> (JLDFlatteningParams () m -> JLDFlatteningParams e m) -> m (Either (JLDError e) Value, JLDExpansionState)
flatten document baseUrl paramsFn = do
    let JLDFlatteningParams{..} =
            paramsFn
                JLDFlatteningParams
                    { jldFlatteningParamsDocumentLoader = DocumentLoader <. const <. pure <| Left ()
                    , jldFlatteningParamsProcessingMode = JLD1_1
                    , jldFlatteningParamsMaxRemoteContexts = 20
                    , jldFlatteningParamsExpandContext = Nothing
                    , jldFlatteningParamsFrameExpansion = False
                    , jldFlatteningParamsState =
                        JLDExpansionState
                            { jldExpansionStateContextCache = mempty
                            , jldExpansionStateDocumentCache = mempty
                            }
                    , jldFlatteningParamsCompactArrays = True
                    }
        expansionParams =
            JLDExpansionParams
                { jldExpansionParamsDocumentLoader = jldFlatteningParamsDocumentLoader
                , jldExpansionParamsProcessingMode = jldFlatteningParamsProcessingMode
                , jldExpansionParamsMaxRemoteContexts = jldFlatteningParamsMaxRemoteContexts
                , jldExpansionParamsExpandContext = jldFlatteningParamsExpandContext
                , jldExpansionParamsFrameExpansion = jldFlatteningParamsFrameExpansion
                , jldExpansionParamsState = jldFlatteningParamsState
                }
    (result, state') <- expand document baseUrl (const expansionParams)
    case result of
        Left err -> pure (Left err, state')
        Right expanded -> fmap (,state') <. runExceptT <| F.flatten expanded