{-# LANGUAGE OverloadedStrings #-}
module Katip.Scribes.ElasticSearch.Annotations
    ( TypeAnnotated(..)
    -- * Exported for benchmarking
    , deannotateValue
    ) where


-------------------------------------------------------------------------------
import           Control.Applicative as A
import           Data.Aeson
import qualified Data.Foldable       as FT
import qualified Data.HashMap.Strict as HM
import           Data.Maybe
import           Data.Monoid         as M
import           Data.Scientific     (isFloating)
import           Data.Text           (Text)
import qualified Data.Text           as T
-------------------------------------------------------------------------------
import           Katip
-------------------------------------------------------------------------------


-- | Represents a value that can be converted to and from JSON that
-- will type annotate object keys when serializing and strip them out when deserializating
newtype TypeAnnotated a = TypeAnnotated {
      typeAnnotatedValue :: a
    }

instance ToJSON a => ToJSON (TypeAnnotated a) where
  toJSON = annotateValue . toJSON . typeAnnotatedValue


instance ToObject a => ToObject (TypeAnnotated a) where
  toObject = annotateKeys . toObject . typeAnnotatedValue


instance FromJSON a => FromJSON (TypeAnnotated a) where
  parseJSON v = TypeAnnotated A.<$> parseJSON (deannotateValue v)


instance LogItem a => LogItem (TypeAnnotated a) where
  payloadKeys v (TypeAnnotated x) = case payloadKeys v x of
    AllKeys -> AllKeys
    -- Take the key selection, overlap it with the actual keys
    -- produced and annotate them
    SomeKeys ks -> let o = toObject x
                       oInFocus = HM.fromList $ zip ks (repeat Null)
                       final = annotateKeys $ HM.intersection o oInFocus
                   in SomeKeys $ HM.keys final

-------------------------------------------------------------------------------
-- Conversion Functions
-------------------------------------------------------------------------------


annotateValue :: Value -> Value
annotateValue (Object o) = Object $ annotateKeys o
annotateValue (Array a)  = Array (annotateValue <$> a)
annotateValue x          = x


annotateKeys :: Object -> Object
annotateKeys = HM.fromList . map go . HM.toList
  where
    go (k, Object o) = (k, Object $ annotateKeys o)
    go (k, Array a)  = (k, Array (annotateValue <$> a))
    go (k, s@(String _)) = (k M.<> stringAnn, s)
    go (k, n@(Number sci)) = if isFloating sci
                             then (k <> doubleAnn, n)
                             else (k <> longAnn, n)
    go (k, b@(Bool _)) = (k <> booleanAnn, b)
    go (k, Null) = (k <> nullAnn, Null)


deannotateValue :: Value -> Value
deannotateValue (Object o) = Object $ deannotateKeys o
deannotateValue (Array a)  = Array (deannotateValue <$> a)
deannotateValue x          = x


deannotateKeys :: Object -> Object
deannotateKeys = HM.fromList . map go . HM.toList
  where
    go (k, Object o) = (k, Object $ deannotateKeys o)
    go (k, Array a)  = (k, Array (deannotateValue <$> a))
    go (k, v)        = (fromMaybe k k', v)
      where
        k' = FT.asum (stripSuffix <$> suffixes)
        suffixes = [stringAnn, doubleAnn, longAnn, booleanAnn, nullAnn]
        stripSuffix suffix = T.stripSuffix suffix k


-------------------------------------------------------------------------------
-- Annotation Constants
-------------------------------------------------------------------------------


stringAnn :: Text
stringAnn = "::s"

doubleAnn :: Text
doubleAnn = "::d"

longAnn :: Text
longAnn = "::l"

booleanAnn :: Text
booleanAnn = "::b"

nullAnn :: Text
nullAnn = "::n"