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

-- | Michelson annotations in untyped model.

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

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

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

  -- * Creation and conversions
  , noAnn
  , annQ
  , varAnnQ
  , fieldAnnQ
  , typeAnnQ
  , mkAnnotation
  , specialVarAnns
  , specialFieldAnn
  , isValidAnnStart
  , isValidAnnBodyChar
  , orAnn
  , unifyAnn
  , unifyPairFieldAnn
  , convergeVarAnns
  , ifAnnUnified
  , convAnn
  ) where

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 (Buildable(build))
import Instances.TH.Lift ()
import Language.Haskell.TH qualified as TH
import Language.Haskell.TH.Lift (deriveLift)
import Language.Haskell.TH.Quote qualified as TH
import Text.PrettyPrint.Leijen.Text (Doc, hsep, textStrict, (<+>))
import Text.Show qualified as T
import Type.Reflection (tyConName, typeRep, typeRepTyCon)

import Morley.Michelson.Printer.Util (RenderDoc(..), buildRenderDoc)
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)

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
-- | Type-level tag for "some" annotations, i.e. those that are unknown at compile-time
data SomeTag

-- | 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
-- | A convenience synonym for "some" t'Annotation', i.e. its type is unknown at compile-time
type SomeAnn = Annotation SomeTag

-- | 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 = ":"

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

instance KnownAnnTag tag => Buildable (Annotation tag) where
  build = buildRenderDoc

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

instance Buildable AnnotationSet where
  build = buildRenderDoc

-- | 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) = textStrict $ annPrefix @tag <> text

-- | 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

-- | Returns the first annotation if it's not empty, or the second one otherwise.
--
-- > "a" `orAnn` "b" == "a"
-- > "a" `orAnn` ""  == "a"
-- > ""  `orAnn` "b" == "b"
-- > ""  `orAnn` ""  == ""
orAnn :: Annotation t -> Annotation t -> Annotation t
orAnn a b = bool a b (a == def)

-- | Given two type or field annotations, attempt to converge them by joining
-- these annotations with the following rule:
-- 1. If either annotation is empty, an empty annotation is returned;
-- 2. If both annotations are equal, return this annotation;
-- 3. Otherwise, returns 'Nothing'.
--
-- This function is used primarily for type-checking and attempts to imitate the
-- reference implementation's observed behavior with annotations.
unifyAnn :: Annotation tag -> Annotation tag -> Maybe (Annotation tag)
unifyAnn a@(Annotation ann1) (Annotation ann2)
  | ann1 == "" || ann2 == "" = Just noAnn
  | ann1 == ann2 = Just $ a
  | otherwise  = Nothing

-- | Given two field annotations where one of them is used in CAR or CDR,
-- attempt to converge them by joining these annotations with the following rule:
-- 1. If either annotation is empty, return the non-empty one (or empty if both are empty);
-- 2. If both annotations are equal, return this annotation;
-- 3. Otherwise, returns 'Nothing'.
--
-- This function is used primarily for type-checking and attempts to imitate the
-- reference implementation's observed behavior with field annotations when CAR
-- and CDR are used with pairs.
unifyPairFieldAnn :: FieldAnn -> FieldAnn -> Maybe FieldAnn
unifyPairFieldAnn a1@(Annotation ann1) a2@(Annotation ann2)
  | ann1 == "" || ann2 == "" = Just $ a1 `orAnn` a2
  | ann1 == ann2 = Just a1
  | otherwise = Nothing

-- | Keeps an annotation if and only if the two of them are equal and returns an
-- empty annotation otherwise.
convergeVarAnns :: VarAnn -> VarAnn -> VarAnn
convergeVarAnns ann1 ann2
  | ann1 == ann2 = ann1
  | otherwise    = noAnn

ifAnnUnified :: Annotation tag -> Annotation tag -> Bool
ifAnnUnified a1 a2 = isJust $ a1 `unifyAnn` a2

-- | Convert annotation from one type to another
convAnn :: Annotation tag1 -> Annotation tag2
convAnn (Annotation a) = UnsafeAnnotation a

-- | Unidirectional pattern synonym matching only non-empty annotations
pattern WithAnn :: Annotation tag -> Annotation tag
pattern WithAnn ann <- ann@(Annotation (toString -> _:_))

deriveJSON morleyAesonOptions ''Annotation
deriveLift ''Annotation