State.hs

module Effectful.OpenTelemetry.Tracing.Trace.State
    ( Key
    , key
    , Value
    , value
    , State
    , fromList
    , fromText
    , toText
    , fromByteString
    , toByteString
    , toList
    , lookup
    , insert
    , delete
    )
where

import Data.Aeson (ToJSON (..))
import Data.ByteString (ByteString)
import Data.Char qualified as Char
import Data.Foldable qualified as Foldable
import Data.List qualified as List
import Data.List.Extra qualified as List
import Data.Maybe (fromJust, mapMaybe)
import Data.Sequence (Seq)
import Data.Sequence qualified as Seq
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import GHC.IsList (IsList)
import GHC.IsList qualified as GHC
import Proto3.Wire.Encode.Class qualified as Proto
import Prelude hiding (lookup)

-- | W3C TraceState Key
--
-- See <https://www.w3.org/TR/trace-context/#key the W3C specification>.
newtype Key = Key Text
    deriving newtype (Show, Eq, Ord)

-- | W3C TraceState Value
--
-- See <https://www.w3.org/TR/trace-context/#value the W3C specification>.
newtype Value = Value Text
    deriving newtype (Show, Eq)

-- | W3C TraceState
--
-- May hold at most 32 unique 'Key's.
-- The 'Semigroup' operation for 'State' prefers values from the left operand.
newtype State = State (Seq (Key, Value))
    deriving newtype (Show, Eq, Monoid)

instance Semigroup State where
    State left <> State right =
        State . Seq.take 32 $ left <> Seq.filter (not . isInLeft . fst) right
      where
        isInLeft = (`elem` (fst <$> left))

instance IsList State where
    type Item State = (Key, Value)
    fromList = fromList
    toList = toList

instance ToJSON State where
    toJSON = toJSON . toText

instance {-# OVERLAPPING #-} Proto.EncodeField State where
    encodeField n = Proto.encodeField n . toText

toList :: State -> [(Key, Value)]
toList (State s) = Foldable.toList s

fromList :: [(Key, Value)] -> State
fromList = State . Seq.fromList . List.take 32 . List.nubOrdOn fst

-- | Parse a 'State' from a comma-separated list of 'key'@=@'value' pairs.
-- Drops any key-value pairs that cannot be parsed.
-- Evaluates to an empty state if the 'Text' contains no parseable key-value pairs.
fromText :: Text -> State
fromText = fromList . mapMaybe (kv . Text.strip) . Text.splitOn ","
  where
    kv :: Text -> Maybe (Key, Value)
    kv t = case Text.breakOn "=" t of
        ( key -> Just k
            , Text.stripPrefix "=" -> fromJust -> value -> Just v
            ) -> Just (k, v)
        _ -> Nothing

toText :: State -> Text
toText = Text.intercalate "," . fmap (\(Key k, Value v) -> k <> "=" <> v) . toList

toByteString :: State -> ByteString
toByteString = Text.encodeUtf8 . toText

-- | Parse a 'State' from a comma-separated list of 'key'@=@'value' pairs.
-- Drops any key-value pairs that cannot be parsed.
-- Evaluates to an empty state if the 'ByteString' contains no parseable key-value pairs.
fromByteString :: ByteString -> State
fromByteString = either (const mempty) fromText . Text.decodeUtf8'

-- | Construct a 'Key', or 'Nothing' on invalid input.
key :: Text -> Maybe Key
key (Text.strip -> t)
    | isValidKey = Just (Key t)
    | otherwise = Nothing
  where
    isValidKey
        | Text.length t > 256 = False
        | otherwise = case Text.splitOn "@" t of
            [simpleKey] -> isValidSimpleKey simpleKey
            [tenant, system] -> isValidTenantKey tenant system
            _ -> False

    isValidTenantKey tenant system =
        not (Text.null tenant)
            && Text.length tenant <= 241
            && Text.length system <= 14
            && isValidSimpleKey tenant
            && isValidSimpleKey system

    isValidSimpleKey k = case Text.uncons k of
        Just (firstChar, rest) -> Char.isAsciiLower firstChar && Text.all isSimpleChar rest
        Nothing -> False

    isSimpleChar c = Char.isAsciiLower c || Char.isDigit c || c `elem` ['_', '-', '*', '/']

-- | Construct a 'Value', or 'Nothing' on invalid input.
value :: Text -> Maybe Value
value (Text.strip -> t)
    | isValidValue = Just (Value t)
    | otherwise = Nothing
  where
    isValidValue = Text.length t < 256 && Text.all isValidChar t
    isValidChar c =
        let code = fromEnum c
         in code >= 0x20 && code <= 0x7E && c /= ',' && c /= '='

-- | Lookup a 'Value' by 'Key'.
lookup :: Key -> State -> Maybe Value
lookup k = List.lookup k . toList

-- | Insert or update a 'Key'-'Value' pair.
-- If the key already exists, it is moved to the left-most position in 'State'.
-- If the 'State' contains 32 keys before insertion, the right-most key will be removed.
insert :: Key -> Value -> State -> State
insert k v s = State (pure (k, v)) <> s

-- | Remove a 'Value' by 'Key' if it exists.
delete :: Key -> State -> State
delete k (State s) = State $ Seq.filter ((k /=) . fst) s