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

-- | Morley extensions to the Michelson language
module Morley.Michelson.Untyped.Ext
  ( ExtInstrAbstract (..)
  , StackRef (..)
  , PrintComment (..)
  , TestAssert (..)
  , Var (..)
  , TyVar (..)
  , StackTypePattern (..)
  , varSet
  , stackTypePatternToList
  ) where

import Data.Aeson.TH (deriveJSON)
import Data.Aeson.Types
  (FromJSON(..), ToJSON(..), genericParseJSON, genericToEncoding, genericToJSON)
import Data.Data (Data(..))
import Data.Set qualified as Set
import Data.Text qualified as T
import Fmt (Buildable(build), Doc, blockListF, blockMapF, listF, nameF)
import Prettyprinter (align, enclose)

import Morley.Michelson.Printer.Util (RenderDoc(..))
import Morley.Michelson.Untyped.HoistInstr
import Morley.Michelson.Untyped.Type
import Morley.Util.Aeson

-- | Implementation-specific instructions embedded in a @NOP@ primitive, which
-- mark a specific point during a contract's typechecking or execution.
--
-- These instructions are not allowed to modify the contract's stack, but may
-- impose additional constraints that can cause a contract to report errors in
-- type-checking or testing.
--
-- Additionaly, some implementation-specific language features such as
-- type-checking of @LetMacro@s are implemented using this mechanism
-- (specifically @FN@ and @FN_END@).
data ExtInstrAbstract f op =
    STACKTYPE StackTypePattern -- ^ Matches current stack against a type-pattern
  | UTEST_ASSERT (TestAssert f op)   -- ^ Copy the current stack and run an inline assertion on it
  | UPRINT PrintComment         -- ^ Print a comment with optional embedded @StackRef@s
  | UCOMMENT Text        -- ^ A comment in Michelson code
  deriving stock (Eq, Show, Data, Generic, Functor)

instance HoistInstr ExtInstrAbstract where
  hoistInstr f = \case
    STACKTYPE x -> STACKTYPE x
    UTEST_ASSERT x -> UTEST_ASSERT $ hoistInstr f x
    UPRINT x -> UPRINT x
    UCOMMENT x -> UCOMMENT x

instance (Foldable f, Buildable op) => Buildable (ExtInstrAbstract f op)

instance NFData (f op) => NFData (ExtInstrAbstract f op)

instance RenderDoc (ExtInstrAbstract f op) where
  renderDoc _ =
    \case
      UCOMMENT t -> enclose "/* " " */" $ align $ build t
      _ -> mempty
  isRenderable =
    \case
      UCOMMENT{} -> True
      _ -> False

-- | A reference into the stack.
newtype StackRef = StackRef Natural
  deriving stock (Eq, Show, Data, Generic)

instance NFData StackRef

instance Buildable StackRef where
  build (StackRef i) = "%[" <> show i <> "]"

-- | A (named) type variable
newtype Var = Var T.Text
  deriving stock (Eq, Show, Ord, Data, Generic)
  deriving anyclass Buildable

instance NFData Var

-- | A type-variable or a type-constant
data TyVar =
    VarID Var
  | TyCon Ty
  deriving stock (Eq, Show, Data, Generic)
  deriving anyclass Buildable

instance NFData TyVar

-- | A stack pattern-match
data StackTypePattern
 = StkEmpty
 | StkRest
 | StkCons TyVar StackTypePattern
  deriving stock (Eq, Show, Data, Generic)

instance NFData StackTypePattern

-- | Convert 'StackTypePattern' to a list of types. Also returns
-- 'Bool' which is 'True' if the pattern is a fixed list of types and
-- 'False' if it's a pattern match on the head of the stack.
stackTypePatternToList :: StackTypePattern -> ([TyVar], Bool)
stackTypePatternToList StkEmpty = ([], True)
stackTypePatternToList StkRest = ([], False)
stackTypePatternToList (StkCons t pat) =
  first (t :) $ stackTypePatternToList pat

instance Buildable StackTypePattern where
  build = listF . pairToList . stackTypePatternToList
    where
      pairToList :: ([TyVar], Bool) -> [Doc]
      pairToList (types, fixed)
        | fixed = map build types
        | otherwise = map build types ++ ["..."]

-- | Get the set of variables in a stack pattern
varSet :: StackTypePattern -> Set Var
varSet = \case
  StkEmpty -> Set.empty
  StkRest -> Set.empty
  (StkCons (VarID v) stk) -> v `Set.insert` (varSet stk)
  (StkCons _ stk) -> varSet stk

-- | A comment with optional embedded 'StackRef's. Used with @PRINT@ extended instruction.
newtype PrintComment = PrintComment
  { unUPrintComment :: [Either T.Text StackRef]
  } deriving stock (Eq, Show, Data, Generic)

instance NFData PrintComment

instance Buildable PrintComment where
  build = foldMap (either build build) . unUPrintComment

-- | An inline test assertion
data TestAssert f (op :: Type) = TestAssert
  { tassName :: T.Text
  , tassComment :: PrintComment
  , tassInstrs :: f op
  } deriving stock (Eq, Show, Functor, Data, Generic)

instance HoistInstr TestAssert where
  hoistInstr f TestAssert{..} = TestAssert{tassInstrs = f tassInstrs, ..}

instance NFData (f op) => NFData (TestAssert f op)

instance (Foldable f, Buildable op) => Buildable (TestAssert f op) where
  build (TestAssert name comment instrs) =
    nameF "TestAssert" $ blockMapF
      [ ("Name" :: Text, build name)
      , ("Comment", build comment)
      , ("Instrs", blockListF instrs)
      ]


-------------------------------------
-- Aeson instances
-------------------------------------

deriveJSON morleyAesonOptions ''StackRef
deriveJSON morleyAesonOptions ''PrintComment

instance FromJSON (f op) => FromJSON (TestAssert f op) where
    parseJSON = genericParseJSON morleyAesonOptions
instance ToJSON (f op) => ToJSON (TestAssert f op) where
    toJSON     = genericToJSON morleyAesonOptions
    toEncoding = genericToEncoding morleyAesonOptions

deriveJSON morleyAesonOptions ''Var
deriveJSON morleyAesonOptions ''TyVar
deriveJSON morleyAesonOptions ''StackTypePattern

instance FromJSON (f op) => FromJSON (ExtInstrAbstract f op) where
    parseJSON = genericParseJSON morleyAesonOptions
instance ToJSON (f op) => ToJSON (ExtInstrAbstract f op) where
    toJSON     = genericToJSON morleyAesonOptions
    toEncoding = genericToEncoding morleyAesonOptions