-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Michelson annotations in untyped model.

{-# LANGUAGE DeriveLift #-}

module Morley.Michelson.Untyped.Annotation
  ( Annotation (..)
  , VarAnns (..)
  , pattern Annotation

  -- * Annotation Set
  , AnnotationSet(..)
  , annsCount
  , emptyAnnSet
  , firstAnn
  , fullAnnSet
  , isNoAnnSet
  , minAnnSetSize
  , secondAnn
  , singleAnnSet
  , singleGroupAnnSet
  , minimizeAnnSet

  -- * Rendering
  , KnownAnnTag(..)
  , TypeAnn
  , FieldAnn
  , VarAnn
  , RootAnn
  , TypeTag
  , FieldTag
  , VarTag
  , AnyAnn(..)

  -- * Creation and conversions
  , noAnn
  , annQ
  , varAnnQ
  , fieldAnnQ
  , typeAnnQ
  , mkAnnotation
  , specialVarAnns
  , specialFieldAnn
  , isValidAnnStart
  , isValidAnnBodyChar
  , renderAnyAnns
  ) where

import Prelude hiding (lift)

import Data.Aeson.TH (deriveJSON)
import Data.Char (isAlpha, isAscii, isDigit, isNumber)
import Data.Data (Data(..))
import Data.Default (Default(..))
import Data.Text qualified as T
import Data.Typeable (eqT, (:~:)(..))
import Fmt (Doc, build, (<+>))
import Instances.TH.Lift ()
import Language.Haskell.TH qualified as TH
import Language.Haskell.TH.Quote qualified as TH
import Language.Haskell.TH.Syntax (Lift(..))
import Prettyprinter (hsep)
import Text.Show qualified as T
import Type.Reflection (tyConName, typeRep, typeRepTyCon)

import Morley.Michelson.Printer.Util (RenderDoc(..))
import Morley.Util.Aeson

-- | Generic Type/Field/Variable Annotation
--
-- As per Michelson documentation, this type has an invariant:
-- (except for the first character, here parametrized in the type @tag@) the
-- allowed character set is the one matching the following regexp:
-- @%|@%%|%@|[@:%][_0-9a-zA-Z][_0-9a-zA-Z\.%@]*
newtype Annotation tag = UnsafeAnnotation { unAnnotation :: Text }
  deriving stock (Eq, Data, Functor, Generic, Lift)

instance NFData (Annotation tag)

-- | Unidirectional pattern synonym used to pattern-match on t'Annotation'
-- without invoking 'UnsafeAnnotation'
pattern Annotation :: Text -> Annotation tag
pattern Annotation ann <- UnsafeAnnotation ann

{-# COMPLETE Annotation :: Annotation #-}

instance Default (Annotation tag) where
  def = noAnn

-- | Either one or two variable annotations
data VarAnns
  = OneVarAnn VarAnn
  | TwoVarAnns VarAnn VarAnn
  deriving stock (Generic, Show)
  deriving anyclass (NFData)

--------------------------------------------------------------------------------
-- Annotation Set
--------------------------------------------------------------------------------

-- | An 'AnnotationSet' contains all the type/field/variable t'Annotation's
-- , with each group in order, associated with an entity.
-- Note that in its rendering/show instances the unnecessary annotations will be
-- omitted, as well as in some of the functions operating with it.
-- Necessary t'Annotation's are the ones strictly required for a consistent
-- representation.
-- In particular, for each group (t/f/v):
--   - if all annotations are 'noAnn' they are all omitted
--   - if one or more 'noAnn' follow a non-empty t'Annotation', they are omitted
--   - if one or more 'noAnn' precede a non-empty t'Annotation', they are kept
--   - every non-empty t'Annotation' is obviously kept
-- This is why order for each group is important as well as separation of
-- different groups of t'Annotation's.
data AnnotationSet = AnnotationSet
  { asTypes  :: [TypeAnn]
  , asFields :: [FieldAnn]
  , asVars   :: [VarAnn]
  } deriving stock (Eq, Show)

instance Semigroup AnnotationSet where
  (AnnotationSet ts1 fs1 vs1) <> (AnnotationSet ts2 fs2 vs2) = AnnotationSet {..}
    where
      asTypes  = ts1 <> ts2
      asFields = fs1 <> fs2
      asVars   = vs1 <> vs2

instance Monoid AnnotationSet where
  mempty = emptyAnnSet

-- | An 'AnnotationSet' without any t'Annotation'.
emptyAnnSet :: AnnotationSet
emptyAnnSet = AnnotationSet [] [] []

-- | An 'AnnotationSet' with only a single t'Annotation' (of any kind).
singleAnnSet :: forall tag. KnownAnnTag tag => Annotation tag -> AnnotationSet
singleAnnSet an = singleGroupAnnSet [an]

-- | An 'AnnotationSet' with several t'Annotation's of the same kind.
singleGroupAnnSet :: forall tag. KnownAnnTag tag => [Annotation tag] -> AnnotationSet
singleGroupAnnSet ans = AnnotationSet {..}
  where
    asTypes = case eqT @tag @TypeTag of Just Refl -> ans; Nothing -> []
    asFields = case eqT @tag @FieldTag of Just Refl -> ans; Nothing -> []
    asVars = case eqT @tag @VarTag of Just Refl -> ans; Nothing -> []

-- | An 'AnnotationSet' built from all 3 kinds of t'Annotation'.
fullAnnSet :: [TypeAnn] -> [FieldAnn] -> [VarAnn] -> AnnotationSet
fullAnnSet asTypes asFields asVars = AnnotationSet {..}

-- | Returns 'True' if all t'Annotation's in the Set are unnecessary/empty/'noAnn'.
-- False otherwise.
isNoAnnSet :: AnnotationSet -> Bool
isNoAnnSet annSet = null asTypes && null asFields && null asVars
  where AnnotationSet {..} = minimizeAnnSet annSet

-- | Returns the amount of t'Annotation's that are necessary for a consistent
-- representation. See 'AnnotationSet'.
minAnnSetSize :: AnnotationSet -> Int
minAnnSetSize annSet = length asTypes + length asFields + length asVars
  where AnnotationSet {..} = minimizeAnnSet annSet

-- | Removes all unnecessary t'Annotation's. See 'AnnotationSet'.
minimizeAnnSet :: AnnotationSet -> AnnotationSet
minimizeAnnSet (AnnotationSet ts fs vs) = AnnotationSet {..}
  where
    asTypes = trimEndNoAnn ts
    asFields = trimEndNoAnn fs
    asVars = trimEndNoAnn vs

-- | Removes all unnecessary t'Annotation's from a list of the same type
trimEndNoAnn :: [Annotation tag] -> [Annotation tag]
trimEndNoAnn = foldr (\a lst -> if null lst && a == noAnn then [] else a : lst) []

-- | Returns the number of annotations in 'AnnotationSet' for each type.
annsCount :: AnnotationSet -> (Int, Int, Int)
annsCount (AnnotationSet tas fas vas) = (length tas, length fas, length vas)

-- | Returns the first annotation in a list of annotations of a specific type
-- in 'AnnotationSet', or 'noAnn' if this list is empty.
firstAnn :: (KnownAnnTag tag) => AnnotationSet -> Annotation tag
firstAnn = getAnn (\case [] -> noAnn; a : _ -> a)

-- | Returns the second annotation in a list of annotations of a specific type
-- in 'AnnotationSet', or 'noAnn' if this list contains less than 2 elements.
secondAnn :: (KnownAnnTag tag) => AnnotationSet -> Annotation tag
secondAnn = getAnn (\case [] -> noAnn; [_] -> noAnn; _ : a : _ -> a)

-- | Retrieves an annotation of a specific type from 'AnnotationSet' using
-- the passed function.
getAnn :: forall tag. (KnownAnnTag tag)
       => ([Annotation tag] -> Annotation tag)
       -> AnnotationSet
       -> Annotation tag
getAnn getter annSet = case eqT @tag @TypeTag of
  Just Refl -> getter $ asTypes annSet
  Nothing -> case eqT @tag @FieldTag of
    Just Refl -> getter $ asFields annSet
    Nothing -> case eqT @tag @VarTag of
      Just Refl -> getter $ asVars annSet
      Nothing -> error "Impossible"

--------------------------------------------------------------------------------
-- Rendering
--------------------------------------------------------------------------------

-- | A constraint representing that type-level annotation tag is known at
-- compile-time
class Typeable (tag :: Type) => KnownAnnTag tag where
  annPrefix :: Text
  -- ^ What prefix is used for the given annotation type (identified by @tag@) in Michelson code,
  -- i.e. @%@ for field annotations, @\@@ for variable annotations, @:@ for type annotations

instance Typeable tag => Show (Annotation tag) where
  showsPrec d (Annotation text) =
    T.showParen (d > app_prec) $ T.showString $ toString $
      "UnsafeAnnotation @" <> tag <> " \"" <> text <> "\""
    where
      app_prec = 10
      tag = toText . tyConName . typeRepTyCon $ typeRep @tag

-- | Type-level tag for type annotations
data TypeTag
-- | Type-level tag for field annotations
data FieldTag
-- | Type-level tag for variable annotations
data VarTag

-- | A convenience synonym for type t'Annotation'
type TypeAnn = Annotation TypeTag
-- | A convenience synonym for field t'Annotation'
type FieldAnn = Annotation FieldTag
-- | A convenience synonym for variable t'Annotation'
type VarAnn = Annotation VarTag

-- | Field annotation for the entire parameter.
type RootAnn = Annotation FieldTag

instance KnownAnnTag FieldTag where
  annPrefix = "%"
instance KnownAnnTag VarTag where
  annPrefix = "@"
instance KnownAnnTag TypeTag where
  annPrefix = ":"

-- | Datatype representing arbitrary annotation.
data AnyAnn = AnyAnnType TypeAnn | AnyAnnField FieldAnn | AnyAnnVar VarAnn
  deriving stock (Eq, Show, Data, Generic)
  deriving anyclass NFData

instance KnownAnnTag tag => RenderDoc (Annotation tag) where
  renderDoc _ = renderAnn

instance RenderDoc AnyAnn where
  renderDoc _ = renderAnyAnn

instance RenderDoc AnnotationSet where
  renderDoc _ (AnnotationSet {..}) =
    renderAnnGroup asTypes <+> renderAnnGroup asFields <+> renderAnnGroup asVars

-- | Renders a single t'Annotation', this is used in every rendering instance of it.
-- Note that this also renders empty ones/'noAnn's because a single t'Annotation'
-- does not have enough context to know if it can be omitted, use 'singleAnnSet'
-- if you want to hide it instead.
renderAnn :: forall tag. KnownAnnTag tag => Annotation tag -> Doc
renderAnn (Annotation text) = build $ annPrefix @tag <> text

renderAnyAnn :: AnyAnn -> Doc
renderAnyAnn = \case
  AnyAnnType x -> renderAnn x
  AnyAnnField x -> renderAnn x
  AnyAnnVar x -> renderAnn x

renderAnyAnns :: [AnyAnn] -> Doc
renderAnyAnns = hsep . map renderAnyAnn

-- | Renders a list of t'Annotation's, omitting unnecessary empty ones/'noAnn'.
-- This is used (3 times) to render an 'AnnotationSet'.
renderAnnGroup :: KnownAnnTag tag => [Annotation tag] -> Doc
renderAnnGroup = hsep . map renderAnn . trimEndNoAnn

--------------------------------------------------------------------------------
-- Creation and conversions
--------------------------------------------------------------------------------

-- | Empty t'Annotation', i.e. no annotation
noAnn :: Annotation a
noAnn = UnsafeAnnotation ""

-- | Makes an t'Annotation' from its textual value, prefix (%/@/:) excluded
-- Returns a 'Text' error message if the given 'Text' contains invalid characters
mkAnnotation :: Text -> Either Text (Annotation a)
mkAnnotation text
  -- TODO [#48]: these are special annotations and should not be always allowed
  | text `elem` specialVarAnns = Right $ UnsafeAnnotation text
  | text == specialFieldAnn = Right $ UnsafeAnnotation text
  | otherwise = do
    suffix <- case T.uncons text of
      Just (h, tl) | isValidAnnStart h -> Right tl
      Just (h, _) -> Left $ "Invalid first character: '" <> one h <> "'"
      _ -> Right ""
    maybe (Right $ UnsafeAnnotation text) (\c -> Left $ "Invalid character: '" <> one c <> "'") $
      T.find (not . isValidAnnBodyChar) suffix

-- |
-- >>> :t [annQ||]
-- ... :: forall {k} {tag :: k}. Annotation tag
--
-- >>> :t [annQ|abc|]
-- ... :: forall {k} {tag :: k}. Annotation tag
annQ :: TH.QuasiQuoter
annQ = annQImpl Nothing

-- |
-- >>> :t [typeAnnQ||]
-- ... :: TypeAnn
--
-- >>> :t [typeAnnQ|abc|]
-- ... :: TypeAnn
typeAnnQ :: TH.QuasiQuoter
typeAnnQ = annQImpl (Just [t|TypeAnn|])

-- |
-- >>> :t [fieldAnnQ||]
-- ... :: FieldAnn
--
-- >>> :t [fieldAnnQ|abc|]
-- ... :: FieldAnn
fieldAnnQ :: TH.QuasiQuoter
fieldAnnQ = annQImpl (Just [t|FieldAnn|])

-- |
-- >>> :t [varAnnQ||]
-- ... :: VarAnn
--
-- >>> :t [varAnnQ|abc|]
-- ... :: VarAnn
varAnnQ :: TH.QuasiQuoter
varAnnQ = annQImpl (Just [t|VarAnn|])

annQImpl :: Maybe TH.TypeQ -> TH.QuasiQuoter
annQImpl annTypeMb = TH.QuasiQuoter
  {
    TH.quoteExp = \s ->
      case (mkAnnotation $ toText @String s) of
        Left err -> fail $ toString err
        Right _ -> case annTypeMb of
          Nothing -> [e| (UnsafeAnnotation $ fromString s) |]
          Just annType -> [e| (UnsafeAnnotation $ fromString s :: $(annType)) |]
  , TH.quotePat = \s ->
       case (mkAnnotation $ toText @String s) of
         Left err -> fail $ toString err
         Right _ -> case annTypeMb of
           Nothing -> [p| UnsafeAnnotation $(TH.litP $ TH.StringL s) |]
           Just annType -> [p| (UnsafeAnnotation $(TH.litP $ TH.StringL s) :: $(annType)) |]
  , TH.quoteType = \_ -> fail "Cannot use this QuasiQuoter at type position"
  , TH.quoteDec = \_  -> fail "Cannot use this QuasiQuoter at declaration position"
  }

-- | List of all the special Variable Annotations, only allowed in @CAR@ and @CDR@
-- instructions, prefix (@) excluded.
-- These do not respect the rules of `isValidAnnStart` and `isValidAnnBodyChar`.
specialVarAnns :: [Text]
specialVarAnns = ["%%","%"]

-- | The only special Field Annotation, only allowed in @PAIR@, @LEFT@ and
-- @RIGHT@ instructions, prefix (%) excluded.
-- This does not respect the rules of `isValidAnnStart` and `isValidAnnBodyChar`.
specialFieldAnn :: Text
specialFieldAnn = "@"


-- | Checks if a `Char` is valid to be the first of an annotation, prefix
-- (%/@/:) excluded, the ones following should be checked with
-- `isValidAnnBodyChar` instead.
-- Note that this does not check Special Annotations, see `specialVarAnns`
-- and `specialFieldAnn`
isValidAnnStart :: Char -> Bool
isValidAnnStart x = (isAscii x && (isAlpha x || isDigit x)) || x == '_'

-- | Checks if a `Char` is valid to be part of an annotation, following a valid
-- first character (see `isValidAnnStart`) and the prefix (%/@/:).
-- Note that this does not check Special Annotations, see `specialVarAnns`
-- and `specialFieldAnn`
isValidAnnBodyChar :: Char -> Bool
isValidAnnBodyChar x =
  isValidAnnStart x || (isAscii x && isNumber x) ||  x `elem` (".%@" :: String)

instance Semigroup VarAnn where
  Annotation a <> Annotation b
    | a == "" || b == "" = UnsafeAnnotation $ a <> b
    | otherwise          = UnsafeAnnotation $ a <> "." <> b

instance Monoid VarAnn where
    mempty = noAnn

deriveJSON morleyAesonOptions ''Annotation
deriveJSON morleyAesonOptions ''AnyAnn