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

{-# LANGUAGE UndecidableSuperClasses #-}

-- | Type and field annotations for Lorentz types.
module Lorentz.Annotation
  ( AnnOptions (..)
  , defaultAnnOptions
  , dropPrefixThen
  , appendTo
  , toCamel
  , toPascal
  , toSnake

  , ctorNameToAnnWithOptions
  , FollowEntrypointFlag (..)
  , GenerateFieldAnnFlag (..)

  , HasAnnotation (..)
  , GHasAnnotation (..)
  , gGetAnnotationNoField
  , insertTypeAnn
  ) where

import Data.Char (isUpper)
import Data.Text qualified as T
import Data.Text.Manipulate (toCamel, toPascal, toSnake)
import GHC.Generics qualified as G

import Morley.Michelson.Text
import Morley.Michelson.Typed
  (BigMap, BigMapId, ContractRef(..), EpAddress, KnownIsoT, Notes(..), Operation, Ticket, ToT,
  insertTypeAnn, starNotes)
import Morley.Michelson.Typed.Haskell.Value (GValueType)
import Morley.Michelson.Untyped (FieldAnn, TypeAnn, VarAnn, mkAnnotation, noAnn)
import Morley.Tezos.Address
import Morley.Tezos.Core
import Morley.Tezos.Crypto
import Morley.Util.Named
import Morley.Util.Text
import Morley.Util.TypeLits

----------------------------------------------------------------------------
-- Annotation Customization
----------------------------------------------------------------------------

-- | Allow customization of field annotation generated for a type
-- when declaring its 'HasAnnotation' instance.
data AnnOptions = AnnOptions
  { fieldAnnModifier :: Text -> Text
  }

defaultAnnOptions :: AnnOptions
defaultAnnOptions = AnnOptions id

-- | Drops the field name prefix from a field.
-- We assume a convention of the prefix always being lower case,
-- and the first letter of the actual field name being uppercase.
-- It also accepts another function which will be applied directly
-- after dropping the prefix.
dropPrefixThen :: (Text -> Text) -> Text -> Text
dropPrefixThen f = f . T.dropWhile (Prelude.not . isUpper)

-- | @appendTo suffix fields field@ appends the given suffix to @field@
-- if the field exists in the @fields@ list.
appendTo :: Text -> [Text] -> Text -> Text
appendTo suffix fields field
  | field `elem` fields = field <> suffix
  | otherwise = field

----------------------------------------------------------------------------
-- Typeclasses related to Annotaiton Generation
----------------------------------------------------------------------------

ctorNameToAnnWithOptions :: forall ctor. (KnownSymbol ctor, HasCallStack) => AnnOptions -> FieldAnn
ctorNameToAnnWithOptions o = unsafe . mkAnnotation . fieldAnnModifier o $ headToLower $ (symbolValT' @ctor)

-- | Used in `GHasAnnotation` and `HasAnnotation` as a flag to track
-- whether or not it directly follows an entrypoint to avoid introducing
-- extra entrypoints.
data FollowEntrypointFlag = FollowEntrypoint | NotFollowEntrypoint

-- | Used in `GHasAnnotation` as a flag to track whether or not field/constructor
-- annotations should be generated.
data GenerateFieldAnnFlag = GenerateFieldAnn | NotGenerateFieldAnn

-- | Use this in the instance of @HasAnnotation@ when field annotations
-- should not be generated.
gGetAnnotationNoField
    :: forall a. (GHasAnnotation (G.Rep a), GValueType (G.Rep a) ~ ToT a)
    => FollowEntrypointFlag -> Notes (ToT a)
gGetAnnotationNoField = \_ -> gGetAnnotation @(G.Rep a) defaultAnnOptions NotFollowEntrypoint NotGenerateFieldAnn ^. _1

-- | This class defines the type and field annotations for a given type. Right now
-- the type annotations come from names in a named field, and field annotations are
-- generated from the record fields.
class HasAnnotation a where
  getAnnotation :: FollowEntrypointFlag -> Notes (ToT a)
  default getAnnotation
    :: (GHasAnnotation (G.Rep a), GValueType (G.Rep a) ~ ToT a)
    => FollowEntrypointFlag
    -> Notes (ToT a)
  getAnnotation b = gGetAnnotation @(G.Rep a) (annOptions @a) b GenerateFieldAnn ^. _1

  annOptions :: AnnOptions
  default annOptions :: AnnOptions
  annOptions = defaultAnnOptions

instance (HasAnnotation a, KnownSymbol name)
  => HasAnnotation (NamedF Identity a name) where
  getAnnotation b = insertTypeAnn (symbolAnn @name) $
    getAnnotation @a b
    where
      symbolAnn :: forall s. KnownSymbol s => TypeAnn
      symbolAnn = unsafe . mkAnnotation $ symbolValT' @s

instance (HasAnnotation (Maybe a), KnownSymbol name)
  => HasAnnotation (NamedF Maybe a name) where
  getAnnotation b = getAnnotation @(NamedF Identity (Maybe a) name) b

-- Primitive instances
instance (HasAnnotation a) => HasAnnotation (Maybe a) where
  getAnnotation _ = NTOption noAnn (getAnnotation @a NotFollowEntrypoint)

instance HasAnnotation ()

instance HasAnnotation Integer where
  getAnnotation _ = starNotes

instance HasAnnotation Natural where
  getAnnotation _ = starNotes

instance HasAnnotation MText where
  getAnnotation _ = starNotes

instance HasAnnotation Bool where
  getAnnotation _ = starNotes

instance HasAnnotation ByteString where
  getAnnotation _ = starNotes

instance HasAnnotation Mutez where
  getAnnotation _ = starNotes

instance HasAnnotation Address where
  getAnnotation _ = starNotes

instance HasAnnotation TxRollupL2Address where
  getAnnotation _ = starNotes

instance HasAnnotation EpAddress where
  getAnnotation _ = starNotes

instance HasAnnotation KeyHash where
  getAnnotation _ = starNotes

instance HasAnnotation Timestamp where
  getAnnotation _ = starNotes

instance HasAnnotation PublicKey where
  getAnnotation _ = starNotes

instance HasAnnotation Signature where
  getAnnotation _ = starNotes

instance HasAnnotation ChainId where
  getAnnotation _ = starNotes

instance (HasAnnotation a) => HasAnnotation (ContractRef a) where
  getAnnotation _ = NTContract noAnn (getAnnotation @a NotFollowEntrypoint)

instance (HasAnnotation d) => HasAnnotation (Ticket d) where
  getAnnotation _ = NTTicket noAnn (getAnnotation @d NotFollowEntrypoint)

instance (HasAnnotation k, HasAnnotation v) => HasAnnotation (Map k v) where
  getAnnotation _ = NTMap noAnn (getAnnotation @k NotFollowEntrypoint) (getAnnotation @v NotFollowEntrypoint)

instance (HasAnnotation k, HasAnnotation v) => HasAnnotation (BigMap k v) where
  getAnnotation _ = NTBigMap noAnn (getAnnotation @k NotFollowEntrypoint) (getAnnotation @v NotFollowEntrypoint)

instance (HasAnnotation k, HasAnnotation v) => HasAnnotation (BigMapId k v) where
  getAnnotation _ = starNotes

instance (KnownIsoT v) => HasAnnotation (Set v) where
  getAnnotation _ = starNotes

instance (HasAnnotation a) => HasAnnotation [a] where
  getAnnotation _ = NTList noAnn (getAnnotation @a NotFollowEntrypoint)

instance HasAnnotation Operation where
  getAnnotation _ = starNotes

instance HasAnnotation Chest where
  getAnnotation _ = starNotes

instance HasAnnotation ChestKey where
  getAnnotation _ = starNotes

instance (HasAnnotation a, HasAnnotation b) => HasAnnotation (Either a b)

instance (HasAnnotation a, HasAnnotation b) => HasAnnotation (a, b)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c) => HasAnnotation (a, b, c)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d) => HasAnnotation (a, b, c, d)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e)
  => HasAnnotation (a, b, c, d, e)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e, HasAnnotation f)
  => HasAnnotation (a, b, c, d, e, f)
instance
  ( HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e
  , HasAnnotation f, HasAnnotation g)
  => HasAnnotation (a, b, c, d, e, f, g)

-- | A Generic @HasAnnotation@ implementation
class GHasAnnotation a where
  gGetAnnotation :: AnnOptions -> FollowEntrypointFlag -> GenerateFieldAnnFlag -> (Notes (GValueType a), FieldAnn, VarAnn)

instance GHasAnnotation G.U1 where
  gGetAnnotation _ _ _ = (starNotes, noAnn, noAnn)

instance (GHasAnnotation x)
  => GHasAnnotation (G.M1 G.S ('G.MetaSel 'Nothing b c d) x)
  where
  gGetAnnotation o b b2 = gGetAnnotation @x o b b2

instance (GHasAnnotation x, KnownSymbol a)
  => GHasAnnotation (G.M1 G.S ('G.MetaSel ('Just a) b c d) x)
  where
  gGetAnnotation o b b2 = case b2 of
    GenerateFieldAnn -> (gGetAnnotation @x o b b2 ^. _1, ctorNameToAnnWithOptions @a o, noAnn)
    NotGenerateFieldAnn -> (gGetAnnotation @x o b b2 ^. _1, noAnn, noAnn)

instance (GHasAnnotation x, KnownSymbol a) => GHasAnnotation (G.M1 G.C ('G.MetaCons a _p _f) x) where
  gGetAnnotation o b b2 =
    ( gGetAnnotation @x o b b2 ^. _1
    , ctorNameToAnnWithOptions @a o
    , noAnn
    )

instance (GHasAnnotation x) => GHasAnnotation (G.M1 G.D i1 x) where
  gGetAnnotation o b b2 = gGetAnnotation @x o b b2

instance
    ( GHasAnnotation x
    , GHasAnnotation y
    )
  =>
    GHasAnnotation (x G.:+: y)
  where

  gGetAnnotation o followEntrypointFlag generateAnnFlag =
    let (xTypeAnn, xFieldAnn, xVarAnn) = gGetAnnotation @x o followEntrypointFlag generateAnnFlag
        (yTypeAnn, yFieldAnn, yVarAnn) = gGetAnnotation @y o followEntrypointFlag generateAnnFlag
    in case (followEntrypointFlag, generateAnnFlag) of
          (NotFollowEntrypoint, GenerateFieldAnn) ->
            ( NTOr noAnn
                xFieldAnn yFieldAnn
                xTypeAnn yTypeAnn
            , noAnn
            , xVarAnn
            )
          _ ->
            ( NTOr noAnn
                noAnn noAnn
                xTypeAnn yTypeAnn
            , noAnn
            , yVarAnn
            )

instance
    ( GHasAnnotation x
    , GHasAnnotation y
    )
  =>
    GHasAnnotation (x G.:*: y)
  where
  gGetAnnotation o _ b2 =
    let  (xTypeAnn, xFieldAnn, xVarAnn) = gGetAnnotation @x o NotFollowEntrypoint b2
         (yTypeAnn, yFieldAnn, yVarAnn) = gGetAnnotation @y o NotFollowEntrypoint b2
    in
      ( NTPair noAnn
          xFieldAnn yFieldAnn
          xVarAnn yVarAnn
          xTypeAnn yTypeAnn
      , noAnn
      , noAnn
      )

instance (HasAnnotation x) => GHasAnnotation (G.Rec0 x) where
  gGetAnnotation _ b _ = (getAnnotation @x b, noAnn, noAnn)