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

module Michelson.Untyped.Ext
  ( ExtInstrAbstract (..)
  , StackRef (..)
  , PrintComment (..)
  , TestAssert (..)
  , Var (..)
  , TyVar (..)
  , StackTypePattern (..)
  , StackFn (..)
  , varSet
  , stackTypePatternToList
  ) where

import Data.Aeson.TH (deriveJSON)
import Data.Data (Data(..))
import qualified Data.Set as Set
import qualified Data.Text as T
import Fmt (Buildable(build), Builder, genericF, listF)
import Text.PrettyPrint.Leijen.Text (textStrict)

import Michelson.Printer.Util (RenderDoc(..), renderOpsList)
import Michelson.Untyped.Type
import 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 op =
    STACKTYPE StackTypePattern -- ^ Matches current stack against a type-pattern
  | FN T.Text StackFn [op]     -- ^ A typed stack function (push and pop a @TcExtFrame@)
  | UTEST_ASSERT (TestAssert 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 NFData op => NFData (ExtInstrAbstract op)

instance RenderDoc op => RenderDoc (ExtInstrAbstract op) where
  renderDoc _ =
    \case
      FN _ _ ops -> renderOpsList False ops
      UCOMMENT t  -> textStrict ("/* " <> t <> " */")
      _ -> mempty
  isRenderable =
    \case
      FN {} -> True
      UCOMMENT{} -> True
      _ -> False

instance Buildable op => Buildable (ExtInstrAbstract op) where
  build = genericF

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

newtype Var = Var T.Text deriving stock (Eq, Show, Ord, Data, Generic)

instance NFData Var

instance Buildable Var where
  build = genericF

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

instance NFData TyVar

instance Buildable TyVar where
  build = genericF

-- | 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) -> [Builder]
      pairToList (types, fixed)
        | fixed = map build types
        | otherwise = map build types ++ ["..."]

-- | A stack function that expresses the type signature of a @LetMacro@
data StackFn = StackFn
  { sfnQuantifiedVars :: Maybe (Set Var)
  , sfnInPattern :: StackTypePattern
  , sfnOutPattern :: StackTypePattern
  } deriving stock (Eq, Show, Data, Generic)

instance NFData StackFn

instance Buildable StackFn where
  build = genericF

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

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 op = TestAssert
  { tassName :: T.Text
  , tassComment :: PrintComment
  , tassInstrs :: [op]
  } deriving stock (Eq, Show, Functor, Data, Generic)

instance NFData op => NFData (TestAssert op)

instance Buildable code => Buildable (TestAssert code) where
  build = genericF

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

deriveJSON morleyAesonOptions ''ExtInstrAbstract
deriveJSON morleyAesonOptions ''PrintComment
deriveJSON morleyAesonOptions ''StackTypePattern
deriveJSON morleyAesonOptions ''StackRef
deriveJSON morleyAesonOptions ''StackFn
deriveJSON morleyAesonOptions ''Var
deriveJSON morleyAesonOptions ''TyVar
deriveJSON morleyAesonOptions ''TestAssert