{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ViewPatterns #-}

-- |
-- Context is a collection of attributes that can be referenced in flag
-- evaluations and analytics events.
--
-- To create a Context of a single kind, such as a user, you may use
-- 'makeContext'.
--
-- To create an LDContext with multiple kinds, use 'makeMultiContext'.
--
-- Additional properties can be set on a single-kind context using the set
-- methods found in this module.
--
-- Each method will always return a Context. However, that Context may be
-- invalid. You can check the validity of the resulting context, and the
-- associated errors by calling 'isValid' and 'getError'.
module LaunchDarkly.Server.Context.Internal
    ( Context (..)
    , SingleContext (..)
    , MultiContext (..)
    , makeContext
    , makeMultiContext
    , withName
    , withAnonymous
    , withAttribute
    , withPrivateAttributes
    , getKey
    , getKeys
    , getCanonicalKey
    , getKinds
    , redactContext
    )
where

import Data.Aeson (FromJSON, Result (Success), ToJSON, Value (..), fromJSON, parseJSON, toJSON, withObject, (.:), (.:?))
import Data.Aeson.Types (Parser, prependFailure, typeMismatch)
import Data.Function ((&))
import Data.Generics.Product (getField, setField)
import qualified Data.HashSet as HS
import Data.List (sortBy)
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Set (Set)
import qualified Data.Set as S
import Data.Text (Text, intercalate, replace, unpack)
import qualified Data.Vector as V
import GHC.Generics (Generic)
import LaunchDarkly.AesonCompat (KeyMap, deleteKey, emptyObject, foldrWithKey, fromList, insertKey, keyMapUnion, lookupKey, mapValues, objectKeys, singleton, toList)
import LaunchDarkly.Server.Config (Config)
import LaunchDarkly.Server.Reference (Reference)
import qualified LaunchDarkly.Server.Reference as R

-- | data record for the Context type
data Context
    = Single SingleContext
    | Multi MultiContext
    | Invalid {error :: !Text}
    deriving (Generic, Show, Eq)

instance ToJSON Context where
    toJSON (Single c) = toJSON c
    toJSON (Multi c) = toJSON c
    toJSON (Invalid c) = toJSON c

instance FromJSON Context where
    parseJSON a@(Object o) =
        case lookupKey "kind" o of
            Nothing -> parseLegacyUser a
            Just (String "multi") -> parseMultiContext a
            Just _ -> parseSingleContext a
    parseJSON invalid = prependFailure "parsing Context failed, " (typeMismatch "Object" invalid)

data SingleContext = SingleContext
    { key :: !Text
    , fullKey :: !Text
    , kind :: !Text
    , name :: !(Maybe Text)
    , anonymous :: !Bool
    , attributes :: !(Maybe (KeyMap Value))
    , privateAttributes :: !(Maybe (Set Reference))
    }
    deriving (Generic, Show, Eq)

instance ToJSON SingleContext where
    toJSON = (toJsonObject True)

data MultiContext = MultiContext
    { fullKey :: !Text
    , contexts :: !(KeyMap SingleContext)
    }
    deriving (Generic, Show, Eq)

instance ToJSON MultiContext where
    toJSON (MultiContext {contexts}) =
        mapValues (\c -> toJsonObject False c) contexts
            & insertKey "kind" "multi"
            & Object

-- |
-- Create a single kind context from the provided hash.
--
-- The provided hash must match the format as outlined in the [SDK
-- documentation](https://docs.launchdarkly.com/sdk/features/user-config).
makeContext :: Text -> Text -> Context
makeContext "" _ = Invalid {error = "context key must not be empty"}
makeContext key kind = makeSingleContext key kind

-- This function is used internally to create a context with legacy key
-- validation rules; namely, a legacy context is allowed to have an empty key.
-- No other type of context is. Users of this SDK can only use the makeContext
-- to create a single-kind context, which includes the non-empty key
-- restriction.
makeSingleContext :: Text -> Text -> Context
makeSingleContext _ "" = Invalid {error = "context kind must not be empty"}
makeSingleContext _ "kind" = Invalid {error = "context kind cannot be 'kind'"}
makeSingleContext _ "multi" = Invalid {error = "context kind cannot be 'multi'"}
makeSingleContext key kind
    | (all (`elem` ['a' .. 'z'] ++ ['A' .. 'Z'] ++ ['0' .. '9'] ++ ['.', '-', '_']) (unpack kind)) == False = Invalid {error = "context kind contains disallowed characters"}
    | otherwise =
        Single
            SingleContext
                { key = key
                , fullKey = canonicalizeKey key kind
                , kind = kind
                , name = Nothing
                , anonymous = False
                , attributes = Nothing
                , privateAttributes = Nothing
                }

-- |
-- Create a multi-kind context from the list of Contexts provided.
--
-- A multi-kind context is comprised of two or more single kind contexts. You
-- cannot include a multi-kind context instead another multi-kind context.
--
-- Additionally, the kind of each single-kind context must be unique. For
-- instance, you cannot create a multi-kind context that includes two user kind
-- contexts.
--
-- If you attempt to create a multi-kind context from one single-kind context,
-- this method will return the single-kind context instead of a new multi-kind
-- context wrapping that one single-kind.
makeMultiContext :: [Context] -> Context
makeMultiContext [] = Invalid {error = "multi-kind contexts require at least one single-kind context"}
makeMultiContext [c@(Single _)] = c
makeMultiContext contexts =
    let singleContexts = mapMaybe unwrapSingleContext contexts
        sorted = sortBy (\lhs rhs -> compare (kind lhs) (kind rhs)) singleContexts
        kinds = HS.fromList $ map kind singleContexts
     in case (length contexts, length singleContexts, length kinds) of
            (a, b, _) | a /= b -> Invalid {error = "multi-kind contexts can only contain single-kind contexts"}
            (a, _, c) | a /= c -> Invalid {error = "multi-kind contexts cannot contain two single-kind contexts with the same kind"}
            _ ->
                Multi
                    MultiContext
                        { fullKey = intercalate ":" $ map (\c -> canonicalizeKey (key c) (kind c)) sorted
                        , contexts = fromList $ map (\c -> ((kind c), c)) singleContexts
                        }

-- |
-- Sets the name attribute for a single-kind context.
--
-- Calling this method on an invalid or multi-kind context is a no-op.
withName :: Text -> Context -> Context
withName name (Single c) = Single $ setField @"name" (Just name) c
withName _ c = c

-- |
-- Sets the anonymous attribute for a single-kind context.
--
-- Calling this method on an invalid or multi-kind context is a no-op.
withAnonymous :: Bool -> Context -> Context
withAnonymous anonymous (Single c) = Single $ setField @"anonymous" anonymous c
withAnonymous _ c = c

-- |
-- Sets the value of any attribute for the context.
--
-- This includes only attributes that are addressable in evaluations -- not
-- metadata such as private attributes. For example, if the attribute name is
-- "privateAttributes", you will be setting an attribute with that name which
-- you can use in evaluations or to record data for your own purposes, but it
-- will be unrelated to 'withPrivateAttributes'.
--
-- If attribute name is "privateAttributeNames", it is ignored and no attribute
-- is set.
--
-- This method uses the Value type to represent a value of any JSON type: null,
-- boolean, number, string, array, or object. For all attribute names that do
-- not have special meaning to LaunchDarkly, you may use any of those types.
-- Values of different JSON types are always treated as different values: for
-- instance, null, false, and the empty string "" are not the same, and the
-- number 1 is not the same as the string "1".
--
-- The following attribute names have special restrictions on their value
-- types, and any value of an unsupported type will be ignored (leaving the
-- attribute unchanged):
--
-- - "name": Must be a string.
-- - "anonymous": Must be a boolean.
--
-- The attribute name "_meta" is not allowed, because it has special meaning in
-- the JSON schema for contexts; any attempt to set an attribute with this name
-- has no effect.
--
-- The attribute names "kind" and "key" are not allowed. They must be provided
-- during the initial context creation. See 'makeContext'.
--
-- Values that are JSON arrays or objects have special behavior when referenced
-- in flag/segment rules.
--
-- For attributes that aren't subject to the special restrictions mentioned
-- above, a value of Null is equivalent to removing any current non-default
-- value of the attribute. Null is not a valid attribute value in the
-- LaunchDarkly model; any expressions in feature flags that reference an
-- attribute with a null value will behave as if the attribute did not exist.
--
-- Calling this method on an invalid or multi-kind context is a no-op.
withAttribute :: Text -> Value -> Context -> Context
withAttribute "key" _ c = c
withAttribute "kind" _ c = c
withAttribute "name" (String value) c = withName value c
withAttribute "name" Null (Single c) = Single $ c {name = Nothing}
withAttribute "name" _ c = c
withAttribute "anonymous" (Bool value) c = withAnonymous value c
withAttribute "anonymous" _ c = c
withAttribute "_meta" _ c = c
withAttribute "privateAttributeNames" _ c = c
withAttribute _ Null c@(Single SingleContext {attributes = Nothing}) = c
withAttribute attr value (Single c@(SingleContext {attributes = Nothing})) =
    Single $ c {attributes = Just $ singleton attr value}
withAttribute attr Null (Single c@(SingleContext {attributes = Just attrs})) =
    Single $ c {attributes = Just $ deleteKey attr attrs}
withAttribute attr value (Single c@(SingleContext {attributes = Just attrs})) =
    Single $ c {attributes = Just $ insertKey attr value attrs}
withAttribute _ _ c = c

-- |
-- Sets the private attributes for a single-kind context.
--
-- Calling this method on an invalid or multi-kind context is a no-op.
withPrivateAttributes :: Set Reference -> Context -> Context
withPrivateAttributes attrs (Single c)
    | S.null attrs = Single $ c {privateAttributes = Nothing}
    | otherwise = Single $ c {privateAttributes = Just attrs}
withPrivateAttributes _ c = c

-- Given a key and kind, generate a canonical key.
--
-- In a multi-kind context, each individual context should theoretically
-- contain the same key. To address this restriction, we generate a canonical
-- key that includes the context's kind. However, if the kind is "user", we
-- omit the kind inclusion to maintain backwards compatibility.
canonicalizeKey :: Text -> Text -> Text
canonicalizeKey key "user" = key
canonicalizeKey key kind = kind <> ":" <> (replace "%" "%25" key & replace ":" "%3A")

unwrapSingleContext :: Context -> Maybe SingleContext
unwrapSingleContext (Single c) = Just c
unwrapSingleContext _ = Nothing

-- Internally used convenience function to retrieve a context's key.
--
-- This method is functionally equivalent to @fromMaybe "" $ getValue "key"@,
-- it's just nicer to use.
getKey :: Context -> Text
getKey (Single c) = key c
getKey _ = ""

-- Internally used convenience function for retrieving all context keys,
-- indexed by their kind.
--
-- A single kind context will return a single element map containing its kind
-- and key. Multi-kind contexts will return a map of kind / key pairs for each
-- of its sub-contexts. An invalid context will return the empty map.
getKeys :: Context -> KeyMap Text
getKeys (Single c) = singleton (kind c) (key c)
getKeys (Multi (MultiContext {contexts})) = mapValues key contexts
getKeys _ = emptyObject

-- Internally used convenience function to retrieve a context's fully qualified
-- key.
getCanonicalKey :: Context -> Text
getCanonicalKey (Single c) = getField @"fullKey" c
getCanonicalKey (Multi c) = getField @"fullKey" c
getCanonicalKey _ = ""

-- Internally used convenience function for retrieving a list of context kinds
-- in the provided context.
--
-- A single kind context will return a single element list containing only that
-- one kind. Multi-kind contexts will return a list of kinds for each of its
-- sub-contexts. An invalid context will return the empty list.
getKinds :: Context -> [Text]
getKinds (Single c) = [kind c]
getKinds (Multi (MultiContext {contexts})) = objectKeys contexts
getKinds _ = []

-- Internally used function for encoding a SingleContext into a JSON object.
--
-- This functionality has been extracted into this separate function because we
-- need to control whether or not the kind property will be included in the
-- final output. If we didn't have this restriction, we could simply inline
-- this function on the SingleContext.
toJsonObject :: Bool -> SingleContext -> Value
toJsonObject includeKind context =
    Object $ fromList $ (getMapOfRedactableProperties context ++ getMapOfRequiredProperties includeKind context)

-- Contexts can be broken into two different types of attributes -- those which
-- can be redacted, and those which can't.
--
-- This method will return a list of name / value pairs which represent the
-- attributes which are eligible for redaction. The other half of the context
-- can be retrieved through the getMapOfRequiredProperties function.
getMapOfRedactableProperties :: SingleContext -> [(Text, Value)]
getMapOfRedactableProperties (SingleContext {name = Nothing, attributes = Nothing}) = []
getMapOfRedactableProperties (SingleContext {name = Nothing, attributes = Just attrs}) = toList attrs
getMapOfRedactableProperties (SingleContext {name = Just n, attributes = Just attrs}) = ("name", String n) : (toList attrs)
getMapOfRedactableProperties (SingleContext {name = Just n, attributes = Nothing}) = [("name", String n)]

-- Contexts can be broken into two different types of attributes -- those which
-- can be redacted, and those which can't.
--
-- This method will return a list of name / value pairs which represent the
-- attributes which cannot be redacted. The other half of the context can be
-- retrieved through the getMapOfRedactableProperties function.
getMapOfRequiredProperties :: Bool -> SingleContext -> [(Text, Value)]
getMapOfRequiredProperties includeKind SingleContext {key, kind, anonymous, privateAttributes} =
    filter
        ((/=) Null . snd)
        [ ("key", toJSON $ key)
        , ("kind", toJSON $ if includeKind then String kind else Null)
        , ("anonymous", toJSON $ if anonymous then Bool True else Null)
        , ("_meta", maybe Null toJSON privateAttributes)
        ,
            ( "_meta"
            , case privateAttributes of
                Nothing -> Null
                Just attrs -> toJSON $ singleton "privateAttributes" (Array $ V.fromList $ map toJSON $ S.elems attrs)
            )
        ]

-- Internally used function to decode a JSON object using the legacy user
-- scheme into a modern single-kind "user" context.
parseLegacyUser :: Value -> Parser Context
parseLegacyUser = withObject "LegacyUser" $ \o -> do
    (key :: Text) <- o .: "key"
    (secondary :: Maybe Text) <- o .:? "secondary"
    (ip :: Maybe Text) <- o .:? "ip"
    (country :: Maybe Text) <- o .:? "country"
    (email :: Maybe Text) <- o .:? "email"
    (firstName :: Maybe Text) <- o .:? "firstName"
    (lastName :: Maybe Text) <- o .:? "lastName"
    (avatar :: Maybe Text) <- o .:? "avatar"
    (name :: Maybe Text) <- o .:? "name"
    (anonymous :: Maybe Bool) <- o .:? "anonymous"
    (custom :: Maybe (KeyMap Value)) <- o .:? "custom"
    (privateAttributeNames :: Maybe [Text]) <- o .:? "privateAttributeNames"
    let context =
            makeSingleContext key "user"
                & withAttribute "secondary" (fromMaybe Null (String <$> secondary))
                & withAttribute "ip" (fromMaybe Null (String <$> ip))
                & withAttribute "country" (fromMaybe Null (String <$> country))
                & withAttribute "email" (fromMaybe Null (String <$> email))
                & withAttribute "firstName" (fromMaybe Null (String <$> firstName))
                & withAttribute "lastName" (fromMaybe Null (String <$> lastName))
                & withAttribute "avatar" (fromMaybe Null (String <$> avatar))
                & withAttribute "name" (fromMaybe Null (String <$> name))
                & withAttribute "anonymous" (fromMaybe Null (Bool <$> anonymous))
                & withPrivateAttributes (S.fromList $ map R.makeLiteral $ fromMaybe [] privateAttributeNames)
     in return $ foldrWithKey (\k v c -> withAttribute k v c) context (fromMaybe emptyObject custom)

-- Internally used function to decode a JSON object using the new context
-- scheme into a modern single-kind context.
parseSingleContext :: Value -> Parser Context
parseSingleContext = withObject "SingleContext" $ \o -> do
    (key :: Text) <- o .: "key"
    (kind :: Text) <- o .: "kind"
    (meta :: Maybe (KeyMap Value)) <- o .:? "_meta"
    (privateAttributes :: Maybe [Text]) <- (fromMaybe emptyObject meta) .:? "privateAttributes"
    let context =
            makeContext key kind
                & withPrivateAttributes (S.fromList $ map R.makeReference $ fromMaybe [] privateAttributes)
     in return $ foldrWithKey (\k v c -> withAttribute k v c) context o

-- Internally used function to decode a JSON object using the new context
-- scheme into a modern multi-kind context.
parseMultiContext :: Value -> Parser Context
parseMultiContext = withObject "MultiContext" $ \o -> do
    let contextLists = toList $ deleteKey "kind" o
        contextObjectLists = mapMaybe (\(k, v) -> case (k, v) of (_, Object obj) -> Just (k, obj); _ -> Nothing) contextLists
        results = map (\(kind, obj) -> fromJSON $ Object $ insertKey "kind" (String kind) obj) contextObjectLists
        single = mapMaybe (\result -> case result of Success r -> Just r; _ -> Nothing) results
     in case (length contextLists, length single) of
            (a, b) | a /= b -> return $ Invalid {error = "multi-kind context JSON contains non-single-kind contexts"}
            (_, _) -> return $ makeMultiContext single

-- Internally used function which performs context attribute redaction.
redactContext :: Config -> Context -> Value
redactContext _ (Invalid _) = Null
redactContext config (Multi MultiContext {contexts}) =
    mapValues (\context -> redactSingleContext False context (getAllPrivateAttributes config context)) contexts
        & insertKey "kind" "multi"
        & Object
        & toJSON
redactContext config (Single context) =
    toJSON $ redactSingleContext True context (getAllPrivateAttributes config context)

-- Apply redaction requirements to a SingleContext type.
redactSingleContext :: Bool -> SingleContext -> Set Reference -> Value
redactSingleContext includeKind context privateAttributes =
    let State {context = redactedContext, redacted} = foldr applyRedaction State {context = fromList $ getMapOfRedactableProperties context, redacted = []} privateAttributes
        redactedValues = Array $ V.fromList $ map String redacted
        required = fromList $ getMapOfRequiredProperties includeKind context
     in case redacted of
            [] -> Object $ keyMapUnion redactedContext required
            _ -> Object $ keyMapUnion redactedContext (insertKey "_meta" (Object $ singleton "redactedAttributes" redactedValues) required)

-- Internally used convenience function for creating a Set of References which
-- can redact all top level values in a provided context.
--
-- Given the context:
-- {
--      "kind": "user",
--      "key": "user-key",
--      "name": "Sandy",
--      "address": {
--          "city": "Chicago"
--      }
-- }
--
-- getAllTopLevelRedactableNames context would yield the set ["name",
-- "address"].
getAllTopLevelRedactableNames :: SingleContext -> Set Reference
getAllTopLevelRedactableNames SingleContext {name = Nothing, attributes = Nothing} = S.empty
getAllTopLevelRedactableNames SingleContext {name = Just _, attributes = Nothing} = S.singleton $ R.makeLiteral "name"
getAllTopLevelRedactableNames SingleContext {name = Nothing, attributes = Just attrs} = S.fromList $ map R.makeLiteral $ objectKeys attrs
getAllTopLevelRedactableNames SingleContext {name = Just _, attributes = Just attrs} = S.fromList $ (R.makeLiteral "name") : (map R.makeLiteral $ objectKeys attrs)

-- Internally used convenience function to return a set of references which
-- would apply all redaction rules.
--
-- If allAttributesPrivate is True in the config, this will return a set which
-- covers the entire context.
getAllPrivateAttributes :: Config -> SingleContext -> Set Reference
getAllPrivateAttributes (getField @"allAttributesPrivate" -> True) context = getAllTopLevelRedactableNames context
getAllPrivateAttributes config SingleContext {privateAttributes = Nothing} = getField @"privateAttributeNames" config
getAllPrivateAttributes config SingleContext {privateAttributes = Just attrs} = S.union (getField @"privateAttributeNames" config) attrs

-- Internally used storage type for returning both the resulting redacted
-- context and the list of any attributes which were redacted.
data State = State
    { context :: KeyMap Value
    , redacted :: ![Text]
    }

-- Internally used store type for managing some state while the redaction
-- process is recursing.
data RedactState = RedactState
    { context :: KeyMap Value
    , reference :: Reference
    , redacted :: ![Text]
    }

-- Kick off the redaction process by priming the recursive redaction state.
applyRedaction :: Reference -> State -> State
applyRedaction reference State {context, redacted} =
    let (RedactState {context = c, redacted = r}) = redactComponents (R.getComponents reference) 0 RedactState {context, redacted, reference}
     in State {context = c, redacted = r}

-- Recursively apply redaction rules
redactComponents :: [Text] -> Int -> RedactState -> RedactState
-- If there are no components left to explore, then we can just return the
-- current state of things. This branch should never actually execute.
-- References aren't valid if there isn't at least one component, and we don't
-- recurse in the single component case. We just include it here for
-- completeness.
redactComponents [] _ state = state
-- kind, key, and anonymous are top level attributes that cannot be redacted.
redactComponents ["kind"] 0 state = state
redactComponents ["key"] 0 state = state
redactComponents ["anonymous"] 0 state = state
-- If we have a single component, then we are either trying to redact a simple
-- top level item, or we have recursed through all reference component parts
-- until the last one. We determine which of those situations we are in through
-- use of the 'level' parameter. 'level' = 0 means we are at the top level of
-- the call stack.
--
-- If we have a single component and we have found it in the current context
-- map, then we know we can redact it.
--
-- If we do not find it in the context, but we are at the top level (and thus
-- making a simple redaction), we consider that a successful redaction.
--
-- Otherwise, if there is no match and we aren't at the top level, the
-- redaction has failed and so we can just return the current state unmodified.
redactComponents [x] level state@(RedactState {context, reference, redacted}) = case (level, lookupKey x context) of
    (_, Just _) -> state {context = deleteKey x context, redacted = (R.getRawPath reference) : redacted}
    (0, _) -> state {redacted = (R.getRawPath reference) : redacted}
    _ -> state
redactComponents (x : xs) level state@(RedactState {context}) = case lookupKey x context of
    Just (Object o) ->
        let substate@(RedactState {context = subcontext}) = redactComponents xs (level + 1) (state {context = o})
         in substate {context = insertKey x (Object $ subcontext) context}
    _ -> state