-- | Michelson instructions in untyped model.

module Michelson.Untyped.Instr
  ( InstrAbstract (..)
  , ExpandedOp (..)
  , ExpandedInstr
  , flattenExpandedOp

  -- * Contract's address
  , OriginationOperation (..)
  , mkContractAddress
  ) where

import Prelude hiding (EQ, GT, LT)

import qualified Data.Aeson as Aeson
import qualified Data.ByteString.Lazy as BSL
import Data.Data (Data(..))
import Fmt (Buildable(build), (+|), (|+))
import Generics.SYB (everywhere, mkT)
import Text.PrettyPrint.Leijen.Text
  (Doc, align, braces, enclose, indent, line, nest, space, text, (<$$>), (<+>))

import Michelson.ErrorPos (InstrCallStack)
import Michelson.Printer.Util
  (RenderDoc(..), buildRenderDoc, doesntNeedParens, needsParens, renderOpsList, spaces)
import Michelson.Untyped.Annotation (Annotation, FieldAnn, TypeAnn, VarAnn, noAnn, renderWEAnn)
import Michelson.Untyped.Contract (Contract'(..))
import Michelson.Untyped.Ext (ExtInstrAbstract)
import Michelson.Untyped.Type (Comparable, Type)
import Michelson.Untyped.Value (Value'(..))
import Tezos.Address (Address, mkContractAddressRaw)
import Tezos.Core (Mutez)
import Tezos.Crypto (KeyHash)

-------------------------------------
-- Types after macroexpander
-------------------------------------

type ExpandedInstr = InstrAbstract ExpandedOp

data ExpandedOp
  = PrimEx ExpandedInstr
  | SeqEx [ExpandedOp]
  | WithSrcEx InstrCallStack ExpandedOp
  deriving stock (Show, Eq, Data, Generic)

instance RenderDoc ExpandedOp where
  renderDoc pn (WithSrcEx _ op) = renderDoc pn op
  renderDoc pn (PrimEx i) = renderDoc pn i
  renderDoc _  (SeqEx i) = renderOpsList False i
  isRenderable =
    \case PrimEx i -> isRenderable i
          WithSrcEx _ op -> isRenderable op
          _ -> True

instance Buildable ExpandedOp where
  build (WithSrcEx _ op) = build op
  build (PrimEx expandedInstr) = "<PrimEx: "+|expandedInstr|+">"
  build (SeqEx expandedOps)    = "<SeqEx: "+|expandedOps|+">"

-- | Flatten all 'SeqEx' in 'ExpandedOp'. This function is mostly for
-- testing. It returns instructions with the same logic, but they are
-- not strictly equivalent, because they are serialized differently
-- (grouping instructions into sequences affects the way they are
-- PACK'ed).
flattenExpandedOp :: ExpandedOp -> [ExpandedInstr]
flattenExpandedOp =
  \case
    PrimEx i -> [flattenInstr i]
    SeqEx ops -> concatMap flattenExpandedOp ops
    WithSrcEx _ op -> flattenExpandedOp op
  where
    flattenInstr :: ExpandedInstr -> ExpandedInstr
    flattenInstr = everywhere (mkT flattenOps)

    flattenOps :: [ExpandedOp] -> [ExpandedOp]
    flattenOps [] = []
    flattenOps (SeqEx s : xs) = s ++ flattenOps xs
    flattenOps (x@(PrimEx _) : xs) = x : flattenOps xs
    flattenOps (WithSrcEx _ op : xs) = op : flattenOps xs

-------------------------------------
-- Abstract instruction
-------------------------------------

-- | Michelson instruction with abstract parameter `op`.  This
-- parameter is necessary, because at different stages of our pipeline
-- it will be different. Initially it can contain macros and
-- non-flattened instructions, but then it contains only vanilla
-- Michelson instructions.
data InstrAbstract op
  = EXT               (ExtInstrAbstract op)
  | DROPN              Word
  -- ^ "DROP n" instruction.
  -- Note: reference implementation permits int16 here.
  -- Negative numbers are parsed successfully there, but rejected later.
  -- Morley is more permissive, so we use 'Word' here,
  -- i. e. permit more positive numbers. We do not permit negative numbers
  -- at type level.
  -- In practice, probably nobody will ever have numbers greater than ≈1000
  -- here, at least due to gas limits.
  -- Same reasoning applies to other instructions which have a numeric
  -- parameter representing number of elements on stack.
  | DROP
  -- ^ 'DROP' is essentially as special case for 'DROPN', but we need
  -- both because they are packed differently.
  | DUP               VarAnn
  | SWAP
  | DIG               Word
  | DUG               Word
  | PUSH              VarAnn Type (Value' op)
  | SOME              TypeAnn VarAnn
  | NONE              TypeAnn VarAnn Type
  | UNIT              TypeAnn VarAnn
  | IF_NONE           [op] [op]
  | PAIR              TypeAnn VarAnn FieldAnn FieldAnn
  | CAR               VarAnn FieldAnn
  | CDR               VarAnn FieldAnn
  | LEFT              TypeAnn VarAnn FieldAnn FieldAnn Type
  | RIGHT             TypeAnn VarAnn FieldAnn FieldAnn Type
  | IF_LEFT           [op] [op]
  | NIL               TypeAnn VarAnn Type
  | CONS              VarAnn -- TODO add TypeNote param
  | IF_CONS           [op] [op]
  | SIZE              VarAnn
  | EMPTY_SET         TypeAnn VarAnn Comparable
  | EMPTY_MAP         TypeAnn VarAnn Comparable Type
  | EMPTY_BIG_MAP     TypeAnn VarAnn Comparable Type
  | MAP               VarAnn [op]
  | ITER              [op]
  | MEM               VarAnn
  | GET               VarAnn
  | UPDATE            VarAnn
  | IF                [op] [op]
  | LOOP              [op]
  | LOOP_LEFT         [op]
  | LAMBDA            VarAnn Type Type [op]
  -- TODO check on alphanet whether we can pass TypeNote
  | EXEC              VarAnn
  | APPLY             VarAnn
  | DIP               [op]
  | DIPN              Word [op]
  | FAILWITH
  | CAST              VarAnn Type
  | RENAME            VarAnn
  | PACK              VarAnn
  | UNPACK            TypeAnn VarAnn Type
  | CONCAT            VarAnn
  | SLICE             VarAnn
  | ISNAT             VarAnn
  | ADD               VarAnn
  | SUB               VarAnn
  | MUL               VarAnn
  | EDIV              VarAnn
  | ABS               VarAnn
  | NEG               VarAnn
  | LSL               VarAnn
  | LSR               VarAnn
  | OR                VarAnn
  | AND               VarAnn
  | XOR               VarAnn
  | NOT               VarAnn
  | COMPARE           VarAnn
  | EQ                VarAnn
  | NEQ               VarAnn
  | LT                VarAnn
  | GT                VarAnn
  | LE                VarAnn
  | GE                VarAnn
  | INT               VarAnn
  | SELF              VarAnn FieldAnn
  | CONTRACT          VarAnn FieldAnn Type
  | TRANSFER_TOKENS   VarAnn
  | SET_DELEGATE      VarAnn
  | CREATE_CONTRACT   VarAnn VarAnn (Contract' op)
  | IMPLICIT_ACCOUNT  VarAnn
  | NOW               VarAnn
  | AMOUNT            VarAnn
  | BALANCE           VarAnn
  | CHECK_SIGNATURE   VarAnn
  | SHA256            VarAnn
  | SHA512            VarAnn
  | BLAKE2B           VarAnn
  | HASH_KEY          VarAnn
  | STEPS_TO_QUOTA    VarAnn
  | SOURCE            VarAnn
  | SENDER            VarAnn
  | ADDRESS           VarAnn
  | CHAIN_ID          VarAnn
  deriving stock (Eq, Show, Functor, Data, Generic)

instance (RenderDoc op) => RenderDoc (InstrAbstract op) where
  renderDoc pn = \case
    EXT extInstr            -> renderDoc pn extInstr
    DROP                    -> "DROP"
    DROPN n                 -> "DROP" <+> text (show n)
    DUP va                  -> "DUP" <+> renderAnnot va
    SWAP                    -> "SWAP"
    DIG n                   -> "DIG" <+> text (show n)
    DUG n                   -> "DUG" <+> text (show n)
    PUSH va t v             -> "PUSH" <+> renderAnnot va <+> renderTy t <+> renderDoc needsParens v
    SOME ta va              -> "SOME" <+> renderAnnot ta <+> renderAnnot va
    NONE ta va t            -> "NONE" <+> renderAnnot ta <+> renderAnnot va <+> renderTy t
    UNIT ta va              -> "UNIT" <+> renderAnnot ta <+> renderAnnot va
    IF_NONE x y             -> "IF_NONE" <+> nest 9 (renderOps x) <$$> spaces 8 <> nest 9 (renderOps y)
    PAIR ta va fa1 fa2      -> "PAIR" <+> renderAnnot ta <+> renderAnnot va <+> renderTwoAnnotations fa1 fa2
    CAR va fa               -> "CAR" <+> renderAnnot va <+> renderAnnot fa
    CDR va fa               -> "CDR" <+> renderAnnot va <+> renderAnnot fa
    LEFT ta va fa1 fa2 t    -> "LEFT" <+> renderAnnot ta <+> renderAnnot va <+> renderTwoAnnotations fa1 fa2 <+> renderTy t
    RIGHT ta va fa1 fa2 t   -> "RIGHT" <+> renderAnnot ta <+> renderAnnot va <+> renderTwoAnnotations fa1 fa2 <+> renderTy t
    IF_LEFT x y             -> "IF_LEFT" <+> nest 9 (renderOps x) <$$> spaces 8 <> nest 9 (renderOps y)
    NIL ta va t             -> "NIL" <+> renderAnnot ta <+> renderAnnot va <+> renderTy t
    CONS va                 -> "CONS" <+> renderAnnot va
    IF_CONS x y             -> "IF_CONS" <+> nest 9 (renderOps x) <$$> spaces 8 <> nest 9 (renderOps y)
    SIZE va                 -> "SIZE" <+> renderAnnot va
    EMPTY_SET ta va t       -> "EMPTY_SET" <+> renderAnnot ta <+> renderAnnot va <+> renderComp t
    EMPTY_MAP ta va c t     -> "EMPTY_MAP" <+> renderAnnot ta <+> renderAnnot va <+> renderComp c <+> renderTy t
    EMPTY_BIG_MAP ta va c t -> "EMPTY_BIG_MAP" <+> renderAnnot ta <+> renderAnnot va <+> renderComp c <+> renderTy t
    MAP va s                -> "MAP" <+> renderAnnot va <$$> spaces 4 <> nest 5 (renderOps s)
    ITER s                  -> "ITER" <+> nest 6 (renderOps s)
    MEM va                  -> "MEM" <+> renderAnnot va
    GET va                  -> "GET" <+> renderAnnot va
    UPDATE va               -> "UPDATE" <+> renderAnnot va
    IF x y                  -> "IF" <+> nest 4 (renderOps x) <$$> spaces 3 <> nest 4 (renderOps y)
    LOOP s                  -> "LOOP" <+> nest 6 (renderOps s)
    LOOP_LEFT s             -> "LOOP_LEFT" <+> nest 11 (renderOps s)
    LAMBDA va t r s         -> "LAMBDA" <+> renderAnnot va <+> renderTy t <+> renderTy r <$$> spaces 7 <> nest 8 (renderOps s)
    EXEC va                 -> "EXEC" <+> renderAnnot va
    APPLY va                -> "APPLY" <+> renderAnnot va
    DIP s                   -> "DIP" <+> nest 5 (renderOps s)
    DIPN n s                -> "DIP" <+> text (show n) <> line <> indent (3 + (length $ show @Text n)) (renderOps s)
    FAILWITH                -> "FAILWITH"
    CAST va t               -> "CAST" <+> renderAnnot va <+> renderTy t
    RENAME va               -> "RENAME" <+> renderAnnot va
    PACK va                 -> "PACK" <+> renderAnnot va
    UNPACK ta va t          -> "UNPACK" <+> renderAnnot ta <+> renderAnnot va <+> renderTy t
    CONCAT va               -> "CONCAT" <+> renderAnnot va
    SLICE va                -> "SLICE" <+> renderAnnot va
    ISNAT va                -> "ISNAT" <+> renderAnnot va
    ADD va                  -> "ADD" <+> renderAnnot va
    SUB va                  -> "SUB" <+> renderAnnot va
    MUL va                  -> "MUL" <+> renderAnnot va
    EDIV va                 -> "EDIV" <+> renderAnnot va
    ABS va                  -> "ABS" <+> renderAnnot va
    NEG va                  -> "NEG" <+> renderAnnot va
    LSL va                  -> "LSL" <+> renderAnnot va
    LSR va                  -> "LSR" <+> renderAnnot va
    OR  va                  -> "OR" <+> renderAnnot va
    AND va                  -> "AND" <+> renderAnnot va
    XOR va                  -> "XOR" <+> renderAnnot va
    NOT va                  -> "NOT" <+> renderAnnot va
    COMPARE va              -> "COMPARE" <+> renderAnnot va
    EQ va                   -> "EQ" <+> renderAnnot va
    NEQ va                  -> "NEQ" <+> renderAnnot va
    LT va                   -> "LT" <+> renderAnnot va
    GT va                   -> "GT" <+> renderAnnot va
    LE va                   -> "LE" <+> renderAnnot va
    GE va                   -> "GE" <+> renderAnnot va
    INT va                  -> "INT" <+> renderAnnot va
    SELF va fa              -> "SELF" <+> renderAnnot va <+> renderAnnot fa
    CONTRACT va fa t        -> "CONTRACT" <+> renderAnnot va <+> renderAnnot fa <+> renderTy t
    TRANSFER_TOKENS va      -> "TRANSFER_TOKENS" <+> renderAnnot va
    SET_DELEGATE va         -> "SET_DELEGATE" <+> renderAnnot va
    CREATE_CONTRACT va1 va2 contract -> let
      body = enclose space space $ align $ (renderDoc doesntNeedParens contract)
      in "CREATE_CONTRACT" <+> renderTwoAnnotations va1 va2 <$$> (indent 2 $ braces $ body)
    IMPLICIT_ACCOUNT va     -> "IMPLICIT_ACCOUNT" <+> renderAnnot va
    NOW va                  -> "NOW" <+> renderAnnot va
    AMOUNT va               -> "AMOUNT" <+> renderAnnot va
    BALANCE va              -> "BALANCE" <+> renderAnnot va
    CHECK_SIGNATURE va      -> "CHECK_SIGNATURE" <+> renderAnnot va
    SHA256 va               -> "SHA256" <+> renderAnnot va
    SHA512 va               -> "SHA512" <+> renderAnnot va
    BLAKE2B va              -> "BLAKE2B" <+> renderAnnot va
    HASH_KEY va             -> "HASH_KEY" <+> renderAnnot va
    STEPS_TO_QUOTA va       -> "STEPS_TO_QUOTA" <+> renderAnnot va
    SOURCE va               -> "SOURCE" <+> renderAnnot va
    SENDER va               -> "SENDER" <+> renderAnnot va
    ADDRESS va              -> "ADDRESS" <+> renderAnnot va
    CHAIN_ID va             -> "CHAIN_ID" <+> renderAnnot va
    where
      renderTy = renderDoc @Type needsParens
      renderComp = renderDoc @Comparable needsParens
      renderOps = renderOpsList False
      renderTwoAnnotations ann1 ann2 =
        -- If both annotations are noAnn we don't want to explicitly print them
        if (ann1 /= noAnn || ann2 /= noAnn)
        then renderWEAnn ann1 <+> renderWEAnn ann2
        else renderDoc doesntNeedParens ann1 <+> renderDoc doesntNeedParens ann2

      renderAnnot :: RenderDoc (Annotation tag) => Annotation tag -> Doc
      renderAnnot = renderDoc @(Annotation _) needsParens

  isRenderable = \case
    EXT extInstr -> isRenderable extInstr
    _ -> True

instance (RenderDoc op, Buildable op) => Buildable (InstrAbstract op) where
  build = \case
    EXT ext -> build ext
    mi -> buildRenderDoc mi

----------------------------------------------------------------------------
-- Contract's address computation
--
-- Note: it might be a bit weird place for this functionality, but it's the
-- lowest layer where all necessary Michelson types are defined. We may
-- reconsider it later.
----------------------------------------------------------------------------

-- | Data necessary to originate a contract.
data OriginationOperation = OriginationOperation
  { ooOriginator :: Address
  -- ^ Originator of the contract.
  , ooDelegate :: Maybe KeyHash
  -- ^ Optional delegate.
  , ooBalance :: Mutez
  -- ^ Initial balance of the contract.
  , ooStorage :: Value' ExpandedOp
  -- ^ Initial storage value of the contract.
  , ooContract :: Contract' ExpandedOp
  -- ^ The contract itself.
  } deriving stock (Show, Generic)

-- | Compute address of a contract from its origination operation.
--
-- TODO [TM-62] It's certainly imprecise, real Tezos implementation doesn't
-- use JSON, but we don't need precise format yet, so we just use some
-- serialization format (JSON because we have necessary instances already).
mkContractAddress :: OriginationOperation -> Address
mkContractAddress = mkContractAddressRaw . BSL.toStrict . Aeson.encode

----------------------------------------------------------------------------
-- JSON serialization
----------------------------------------------------------------------------

instance Aeson.ToJSON ExpandedOp
instance Aeson.FromJSON ExpandedOp
instance Aeson.ToJSON op => Aeson.ToJSON (InstrAbstract op)
instance Aeson.FromJSON op => Aeson.FromJSON (InstrAbstract op)
instance Aeson.FromJSON OriginationOperation
instance Aeson.ToJSON OriginationOperation