-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Michelson annotations in untyped model.

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

  -- * Annotation Set
  , AnnotationSet
  , emptyAnnSet
  , fullAnnSet
  , isNoAnnSet
  , minAnnSetSize
  , singleAnnSet
  , singleGroupAnnSet

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

  -- * Creation and conversions
  , noAnn
  , ann
  , mkAnnotation
  , specialVarAnns
  , specialFieldAnn
  , isValidAnnStart
  , isValidAnnBodyChar
  , unifyAnn
  , ifAnnUnified
  , disjoinVn
  , convAnn
  ) where

import Data.Aeson.TH (deriveJSON)
import Data.Char (isAlpha, isAscii, isDigit, isNumber)
import Data.Data (Data(..))
import Data.Default (Default(..))
import qualified Data.Kind as Kind
import qualified Data.Text as T
import Data.Typeable ((:~:)(..), eqT)
import Fmt (Buildable(build))
import Instances.TH.Lift ()
import Language.Haskell.TH.Lift (deriveLift)
import Text.PrettyPrint.Leijen.Text (Doc, hsep, textStrict, (<+>))
import qualified Text.Show

import Michelson.Printer.Util (RenderDoc(..), buildRenderDoc, doesntNeedParens, printDocS)
import 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 = AnnotationUnsafe { unAnnotation :: Text }
  deriving stock (Eq, Data, Functor, Generic)
  deriving newtype (IsString)

instance NFData (Annotation tag)

pattern Annotation :: Text -> Annotation tag
pattern Annotation ann <- AnnotationUnsafe ann

{-# COMPLETE Annotation :: Annotation #-}

instance Default (Annotation tag) where
  def = noAnn

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

-- | An 'AnnotationSet' contains all the type/field/variable '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 '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 'ann', they are omitted
--   - if one or more 'noAnn' precede a non-empty 'ann', they are kept
--   - every non-empty 'ann' is obviously kept
-- This is why order for each group is important as well as separation of
-- different groups of 'Annotation's.
data AnnotationSet = AnnotationSet
  { asTypes  :: [TypeAnn]
  , asFields :: [FieldAnn]
  , asVars   :: [VarAnn]
  } deriving stock Eq

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 'Annotation'.
emptyAnnSet :: AnnotationSet
emptyAnnSet = AnnotationSet [] [] []

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

-- | An 'AnnotationSet' with several '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 'Annotation'.
fullAnnSet :: [TypeAnn] -> [FieldAnn] -> [VarAnn] -> AnnotationSet
fullAnnSet asTypes asFields asVars = AnnotationSet {..}

-- | Returns 'True' if all '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 '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 '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 '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) []

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

class Typeable (tag :: Kind.Type) => KnownAnnTag tag where
  annPrefix :: Text

instance KnownAnnTag tag => Show (Annotation tag) where
  show = printDocS True . renderDoc doesntNeedParens

data TypeTag
data FieldTag
data VarTag
data SomeTag

type TypeAnn = Annotation TypeTag
type FieldAnn = Annotation FieldTag
type VarAnn = Annotation VarTag
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 Show AnnotationSet where
  show = printDocS True . renderDoc doesntNeedParens

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

instance Buildable AnnotationSet where
  build = buildRenderDoc

-- | Renders a single 'Annotation', this is used in every rendering instance of it.
-- Note that this also renders empty ones/'noAnn's because a single '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 '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
--------------------------------------------------------------------------------

noAnn :: Annotation a
noAnn = AnnotationUnsafe ""

-- | Makes an `Annotation` from its textual value, prefix (%/@/:) excluded
-- Throws an error if the given `Text` contains invalid characters
ann :: HasCallStack => Text -> Annotation a
ann = either error id . mkAnnotation

-- | Makes an `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 $ AnnotationUnsafe text
  | text == specialFieldAnn = Right $ AnnotationUnsafe text
  | otherwise = do
    suffix <- case T.uncons text of
      Just (h, tl) | isValidAnnStart h -> Right tl
      Just (h, _) -> Left $ T.snoc "Invalid first character: " h
      _ -> Right ""
    maybe (Right $ AnnotationUnsafe text) (Left . T.snoc "Invalid character: ") $
      T.find (not . isValidAnnBodyChar) suffix

-- | 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 == "" = ann $ a <> b
    | otherwise          = ann $ a <> "." <> b

instance Monoid VarAnn where
    mempty = noAnn

unifyAnn :: Annotation tag -> Annotation tag -> Maybe (Annotation tag)
unifyAnn (Annotation ann1) (Annotation ann2)
  | ann1 == "" || ann2 == "" = Just $ ann $ ann1 <> ann2
  | ann1 == ann2 = Just $ ann ann1
  | otherwise  = Nothing

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

disjoinVn :: VarAnn -> (VarAnn, VarAnn)
disjoinVn (Annotation a) = case T.findIndex (== '.') $ T.reverse a of
  Just ((n - 1 -) -> pos) -> (ann $ T.take pos a, ann $ T.drop (pos + 1) a)
  Nothing                 -> (noAnn, ann a)
  where
    n = T.length a

convAnn :: Annotation tag1 -> Annotation tag2
convAnn (Annotation a) = ann a

pattern WithAnn :: Annotation tag -> Annotation tag
pattern WithAnn ann <- ann@(Annotation (toString -> _:_))

deriveJSON morleyAesonOptions ''Annotation
deriveLift ''Annotation