-- | Functions for working with Meta, the default annotation type in Hydra

module Hydra.Meta where

import Hydra.Core
import Hydra.Compute
import Hydra.CoreDecoding
import Hydra.CoreEncoding
import Hydra.Common
import Hydra.Monads
import Hydra.Mantle
import qualified Hydra.Impl.Haskell.Dsl.Terms as Terms

import qualified Data.Map as M
import qualified Data.Maybe as Y


aggregateAnnotations :: (a -> Maybe (Annotated a Meta)) -> a -> Meta
aggregateAnnotations getAnn t = Meta $ M.fromList $ addMeta [] t
  where
    addMeta m t = case getAnn t of
      Nothing -> m
      Just (Annotated t' (Meta other)) -> addMeta (m ++ M.toList other) t'

getAnnotation :: String -> Meta -> Maybe (Term Meta)
getAnnotation key (Meta m) = M.lookup key m

getAttr :: String -> Flow s (Maybe (Term Meta))
getAttr key = Flow q
  where
    q s0 t0 = FlowState (Just $ M.lookup key $ traceOther t0) s0 t0

getAttrWithDefault :: String -> Term Meta -> Flow s (Term Meta)
getAttrWithDefault key def = Y.fromMaybe def <$> getAttr key

getDescription :: Meta -> GraphFlow Meta (Y.Maybe String)
getDescription meta = case getAnnotation metaDescription meta of
  Nothing -> pure Nothing
  Just term -> case term of
    TermLiteral (LiteralString s) -> pure $ Just s
    _ -> fail $ "unexpected value for " ++ show metaDescription ++ ": " ++ show term

getTermAnnotation :: Context Meta -> String -> Term Meta -> Y.Maybe (Term Meta)
getTermAnnotation cx key = getAnnotation key . termMetaInternal

getTermDescription :: Term Meta -> GraphFlow Meta (Y.Maybe String)
getTermDescription = getDescription . termMetaInternal

getType :: Meta -> GraphFlow Meta (Y.Maybe (Type Meta))
getType meta = case getAnnotation metaType meta of
  Nothing -> pure Nothing
  Just dat -> Just <$> decodeType dat

getTypeDescription :: Type Meta -> GraphFlow Meta (Y.Maybe String)
getTypeDescription = getDescription . typeMetaInternal

metaAnnotationClass :: AnnotationClass Meta
metaAnnotationClass = AnnotationClass {
    annotationClassDefault = Meta M.empty,
    annotationClassEqual = (==),
    annotationClassCompare = \m1 m2 -> toComparison $ m1 `compare` m2,
    annotationClassShow = show,
    annotationClassRead = read,

    -- TODO: simplify
    annotationClassTermMeta = termMetaInternal,
    annotationClassTypeMeta = typeMetaInternal,
    annotationClassTermDescription = getTermDescription,
    annotationClassTypeDescription = getTypeDescription,
    annotationClassTermType = getType . termMetaInternal,
    annotationClassSetTermDescription = setTermDescription,
    annotationClassSetTermType = setTermType,
    annotationClassTypeOf = getType,
    annotationClassSetTypeOf = setType}
  where
    toComparison c = case c of
      LT -> ComparisonLessThan
      EQ -> ComparisonEqualTo
      GT -> ComparisonGreaterThan

metaDescription :: String
metaDescription = "description"

metaType :: String
metaType = "type"

nextCount :: String -> Flow s Int
nextCount attrName = do
  count <- getAttrWithDefault attrName (Terms.int32 0) >>= Terms.expectInt32
  putAttr attrName (Terms.int32 $ count + 1)
  return count

putAttr :: String -> Term Meta -> Flow s ()
putAttr key val = Flow q
  where
    q s0 t0 = FlowState (Just ()) s0 (t0 {traceOther = M.insert key val $ traceOther t0})

setAnnotation :: String -> Y.Maybe (Term Meta) -> Meta -> Meta
setAnnotation key val (Meta m) = Meta $ M.alter (const val) key m

setDescription :: Y.Maybe String -> Meta -> Meta
setDescription d = setAnnotation metaDescription (Terms.string <$> d)

setTermAnnotation :: Context Meta -> String -> Y.Maybe (Term Meta) -> Term Meta -> Term Meta
setTermAnnotation cx key val term = if meta == annotationClassDefault (contextAnnotations cx)
    then term'
    else TermAnnotated $ Annotated term' meta
  where
    term' = stripTerm term
    meta = setAnnotation key val $ termMetaInternal term

setTermDescription :: Context Meta -> Y.Maybe String -> Term Meta -> Term Meta
setTermDescription cx d = setTermAnnotation cx metaDescription (Terms.string <$> d)

setTermType :: Context Meta -> Y.Maybe (Type Meta) -> Term Meta -> Term Meta
setTermType cx d = setTermAnnotation cx metaType (encodeType <$> d)

setType :: Y.Maybe (Type Meta) -> Meta -> Meta
setType mt = setAnnotation metaType (encodeType <$> mt)

setTypeAnnotation :: Context Meta -> String -> Y.Maybe (Term Meta) -> Type Meta -> Type Meta
setTypeAnnotation cx key val typ = if meta == annotationClassDefault (contextAnnotations cx)
    then typ'
    else TypeAnnotated $ Annotated typ' meta
  where
    typ' = stripType typ
    meta = setAnnotation key val $ typeMetaInternal typ

setTypeDescription :: Context Meta -> Y.Maybe String -> Type Meta -> Type Meta
setTypeDescription cx d = setTypeAnnotation cx metaDescription (Terms.string <$> d)

termMetaInternal :: Term Meta -> Meta
termMetaInternal = aggregateAnnotations $ \t -> case t of
  TermAnnotated a -> Just a
  _ -> Nothing

typeMetaInternal :: Type Meta -> Meta
typeMetaInternal = aggregateAnnotations $ \t -> case t of
  TypeAnnotated a -> Just a
  _ -> Nothing