-- | Module, carrying logic of @PACK@ instruction.
--
-- This is nearly symmetric to adjacent Unpack.hs module.
module Michelson.Interpret.Pack
  ( packCode'
  , packT'
  , packValue
  , packValue'
    -- * Serializers used in morley-client
  , encodeValue'
  , packNotedT'
    -- * Internals
  , encodeIntPayload
  ) where

import Prelude hiding (EQ, GT, LT)

import Control.Exception (assert)
import qualified Data.Binary.Put as Bi
import qualified Data.Bits as Bits
import qualified Data.ByteArray as ByteArray
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Map as Map
import Data.Singletons (SingI(..), demote)

import Michelson.Text
import Michelson.Typed
import Michelson.Untyped.Annotation
  (Annotation(..), FieldAnn, TypeAnn, VarAnn, noAnn, renderWEAnn)
import Tezos.Address (Address(..), ContractHash(..))
import Tezos.Core (ChainId(..), Mutez(..), timestampToSeconds)
import Tezos.Crypto (KeyHash(..), KeyHashTag(..), PublicKey(..), signatureToBytes)
import qualified Tezos.Crypto.Ed25519 as Ed25519
import qualified Tezos.Crypto.P256 as P256
import qualified Tezos.Crypto.Secp256k1 as Secp256k1
import Util.Peano (peanoValSing)

-- | Serialize a value given to @PACK@ instruction.
packValue :: PackedValScope t => Value t -> LByteString
packValue x = "\x05" <> encodeValue x

-- | Same as 'packValue', for strict bytestring.
packValue' :: PackedValScope t => Value t -> ByteString
packValue' = LBS.toStrict . packValue

encodeValue' :: (SingI t, HasNoOp t) => Value t -> ByteString
encodeValue' = LBS.toStrict . encodeValue

packT' :: forall (t :: T). SingI t => ByteString
packT' = LBS.toStrict $ encodeT' @t

packCode' :: Instr inp out -> ByteString
packCode' = LBS.toStrict . encodeInstrs

surround :: LByteString -> LByteString -> ByteString -> LByteString
surround prefix suffix main = prefix <> LBS.fromStrict main <> suffix

-- | Generic serializer.
--
-- We don't require @HasNoBigMap@ constraint here since the big_map serialization
-- is only prohibited in @PACK@ instructions, however, we still want to be able to
-- serialize big_map e.g. in order to transform typed value to low-level Micheline
-- representation.
-- TODO: Serialize chain operations properly as well since they actually also have
-- byte representation.
encodeValue :: forall t. (SingI t, HasNoOp t) => Value t -> LByteString
encodeValue val = case (val, sing @t) of
  (VC cval, _) -> encodeCValue cval
  (VKey s, _) -> encodeBytes . LBS.fromStrict $ case s of
    PublicKeyEd25519 pk -> "\x00" <> Ed25519.publicKeyToBytes pk
    PublicKeySecp256k1 pk -> "\x01" <> Secp256k1.publicKeyToBytes pk
    PublicKeyP256 pk -> "\x02" <> P256.publicKeyToBytes pk
  (VUnit, _) -> "\x03\x0b"
  (VSignature x, _) -> encodeBytes . LBS.fromStrict $ signatureToBytes x
  (VChainId x, _) ->
    encodeBytes . LBS.fromStrict $
      ByteArray.convert (unChainId x)
  (VOption (Just x), STOption _) -> "\x05\x09" <> encodeValue x
  (VOption Nothing, _) -> "\x03\x06"
  (VList xs, STList _) -> encodeList encodeValue xs
  (VSet xs, _) -> encodeList encodeCValue (toList xs)
  (VContract addr sepc, _) -> encodeEpAddress $ EpAddress addr (sepcName sepc)
  (VPair (v1, v2), STPair l _) ->
    case checkOpPresence l of
      OpAbsent -> "\x07\x07" <> encodeValue v1 <> encodeValue v2
  (VOr (Left v), STOr l _) ->
    case checkOpPresence l of
      OpAbsent -> "\x05\x05" <> encodeValue v
  (VOr (Right v), STOr l _) ->
    case checkOpPresence l of
      OpAbsent-> "\x05\x08" <> encodeValue v
  (VLam lam, _) -> encodeInstrs $ rfAnyInstr lam
  (VMap m, STMap _ _) -> encodeMap m
  (VBigMap m, STBigMap _ _) -> encodeMap m

encodeCValue :: CValue t -> LByteString
encodeCValue = \case
  CvInt x        -> encodeNumeric x
  CvNat x        -> encodeNumeric x
  CvString text  -> encodeString text
  CvBytes bytes  -> encodeBytes (LBS.fromStrict bytes)
  CvMutez x      -> encodeNumeric (unMutez x)
  CvBool True    -> "\x03\x0a"
  CvBool False   -> "\x03\x03"
  CvKeyHash kh   -> encodeBytes $ encodeKeyHashRaw kh
  CvTimestamp x  -> encodeNumeric (timestampToSeconds @Integer x)
  CvAddress addr -> encodeEpAddress addr

encodeLength :: Int -> LByteString
encodeLength = Bi.runPut . Bi.putWord32be . fromIntegral

-- | Lift encoded list content to an entire encoded list.
encodeAsList :: LByteString -> LByteString
encodeAsList bs = encodeLength (length bs) <> bs

-- | Encode a list-like structure.
encodeList :: (a -> LByteString) -> [a] -> LByteString
encodeList encodeElem l = "\x02" <> encodeAsList (LBS.concat $ map encodeElem l)

-- | Encode a text.
encodeString :: MText -> LByteString
encodeString text = "\x01" <> encodeAsList (encodeUtf8 $ unMText text)

-- | Encode some raw data.
encodeBytes :: LByteString -> LByteString
encodeBytes bs =
  "\x0a" <> encodeAsList bs

encodeEpName :: EpName -> LByteString
encodeEpName = encodeUtf8 . unAnnotation . epNameToRefAnn

-- | Encode some map.
encodeMap :: (SingI v, HasNoOp v) => Map (CValue k) (Value v) -> LByteString
encodeMap m =
  encodeList (\(k, v) -> "\x07\x04" <> encodeCValue k <> encodeValue v) (Map.toList m)

encodeKeyHashRaw :: KeyHash -> LByteString
encodeKeyHashRaw kh = (<> LBS.fromStrict (khBytes kh)) $
  case khTag kh of
    KeyHashEd25519 -> "\x00"
    KeyHashSecp256k1 -> "\x01"
    KeyHashP256 -> "\x02"

encodeAddress :: Address -> LByteString
encodeAddress = \case
  KeyAddress keyHash -> "\x00" <> (encodeKeyHashRaw keyHash)
  ContractAddress (ContractHash address) -> surround "\x01" "\x00" address

encodeEpAddress :: EpAddress -> LByteString
encodeEpAddress (EpAddress addr epName) =
  encodeBytes $ encodeAddress addr <> encodeEpName epName

-- | Encode contents of a given number.
encodeIntPayload :: Integer -> LByteString
encodeIntPayload = LBS.pack . toList . doEncode True
  where
    {- Numbers are represented as follows:

    byte 0:         1              _         ______   ||  lowest digits
            has continuation  is negative   payload   ||
                                                      ||
    byte 1:         1                       _______   ||
    ...             1                       _______   ||
    byte n:         0                       _______   ||
            has continuation                payload   \/  highest digits
    -}
    doEncode :: Bool -> Integer -> NonEmpty Word8
    doEncode isFirst a
      | a >= byteWeight =
          let (hi, lo) = a `divMod` byteWeight
              byte = Bits.setBit (fromIntegral @_ @Word8 lo) 7
          in byte :| toList (doEncode False hi)
      | a >= 0 =
          one (fromIntegral @_ @Word8 a)
      | otherwise = assert isFirst $
          let h :| t = doEncode True (-a)
          in Bits.setBit h 6 :| t
      where
        byteWeight = if isFirst then 64 else 128

-- | Encode an int-like value.
encodeNumeric :: Integral i => i -> LByteString
encodeNumeric i = "\x00" <> encodeIntPayload (fromIntegral i)

-- | Encode a code block.
encodeInstrs :: Instr inp out -> LByteString
encodeInstrs = encodeList id . one . encodeInstr

-- | Encode an instruction.
encodeInstr :: forall inp out. Instr inp out -> LByteString
encodeInstr = \case
  InstrWithNotes n a ->
    encodeNotedInstr a n []
  InstrWithVarNotes varNotes a ->
    encodeVarNotedInstr a varNotes
  FrameInstr _ i ->
    encodeInstr i
  Seq a b ->
    encodeInstr a <> encodeInstr b
  Nop ->
    mempty
  Nested i ->
    encodeInstrs i
  DocGroup _ i ->
    encodeInstrs i
  Ext _ ->
    ""
  DROP ->
    "\x03\x20"
  DROPN s ->
    "\x05\x20" <> encodeNumeric (peanoValSing s)
  DUP ->
    "\x03\x21"
  SWAP ->
    "\x03\x4c"
  DIG s ->
    "\x05\x70" <> encodeNumeric (peanoValSing s)
  DUG s ->
    "\x05\x71" <> encodeNumeric (peanoValSing s)
  PUSH (a :: Value t) ->
    "\x07\x43" <> encodeT' @t <> encodeValue a
  SOME ->
    "\x03\x46"
  NONE | _ :: Proxy ('TOption t ': s) <- Proxy @out ->
    "\x05\x3e" <> encodeT' @t
  UNIT ->
    "\x03\x4f"
  IF_NONE a b ->
    "\x07\x2f" <> encodeInstrs a <> encodeInstrs b
  PAIR ->
    "\x03\x42"
  (AnnCAR fn) ->
    encodeWithAnns [] [fn] [] "\x03\x16"
  (AnnCDR fn) ->
    encodeWithAnns [] [fn] [] "\x03\x17"
  LEFT | _ :: Proxy ('TOr l r ': s) <- Proxy @out ->
    "\x05\x33" <> encodeT' @r
  RIGHT | _ :: Proxy ('TOr l r ': s) <- Proxy @out ->
    "\x05\x44" <> encodeT' @l
  IF_LEFT a b ->
    "\x07\x2e" <> encodeInstrs a <> encodeInstrs b
  NIL | _ :: Proxy ('TList t ': s) <- Proxy @out ->
    "\x05\x3d" <> encodeT' @t
  CONS ->
    "\x03\x1b"
  IF_CONS a b ->
    "\x07\x2d" <> encodeInstrs a <> encodeInstrs b
  SIZE ->
    "\x03\x45"
  EMPTY_SET | _ :: Proxy ('TSet t ': s) <- Proxy @out ->
    "\x05\x24" <> encodeT' @('Tc t)
  EMPTY_MAP | _ :: Proxy ('TMap k v ': s) <- Proxy @out ->
    "\x07\x23" <> encodeT' @('Tc k) <> encodeT' @v
  EMPTY_BIG_MAP | _ :: Proxy ('TBigMap k v ': s) <- Proxy @out ->
    "\x07\x72" <> encodeT' @('Tc k) <> encodeT' @v
  MAP a ->
    "\x05\x38" <> encodeInstrs a
  ITER a ->
    "\x05\x52" <> encodeInstrs a
  MEM ->
    "\x03\x39"
  GET ->
    "\x03\x29"
  UPDATE ->
    "\x03\x50"
  IF a b ->
    "\x07\x2c" <> encodeInstrs a <> encodeInstrs b
  LOOP a ->
    "\x05\x34" <> encodeInstrs a
  LOOP_LEFT a ->
    "\x05\x53" <> encodeInstrs a
  LAMBDA (v :: Value ('TLambda i o)) ->
    "\x09\x31" <>
    encodeAsList (encodeT' @i <> encodeT' @o <> encodeValue v) <>
    encodeLength 0  -- encoding of a Variable Annotation (that we don't support)
  EXEC ->
    "\x03\x26"
  APPLY ->
    "\x03\x73"
  DIP a ->
    "\x05\x1f" <> encodeInstrs a
  DIPN s a ->
    "\x07\x1f" <> encodeNumeric (peanoValSing s) <> encodeInstrs a
  FAILWITH ->
    "\x03\x27"
  CAST | _ :: Proxy (t ': s) <- Proxy @out ->
    "\x05\x57" <> encodeT' @t
  RENAME ->
    "\x03\x58"
  PACK ->
    "\x03\x0c"
  UNPACK | _ :: Proxy ('TOption t ': s) <- Proxy @out ->
    "\x05\x0d" <> encodeT' @t
  CONCAT ->
    "\x03\x1a"
  CONCAT' ->
    "\x03\x1a"
  SLICE ->
    "\x03\x6f"
  ISNAT ->
    "\x03\x56"
  ADD ->
    "\x03\x12"
  SUB ->
    "\x03\x4b"
  MUL ->
    "\x03\x3a"
  EDIV ->
    "\x03\x22"
  ABS ->
    "\x03\x11"
  NEG ->
    "\x03\x3b"
  LSL ->
    "\x03\x35"
  LSR ->
    "\x03\x36"
  OR ->
    "\x03\x41"
  AND ->
    "\x03\x14"
  XOR ->
    "\x03\x51"
  NOT ->
    "\x03\x3f"
  COMPARE ->
    "\x03\x19"
  EQ ->
    "\x03\x25"
  NEQ ->
    "\x03\x3c"
  LT ->
    "\x03\x37"
  GT ->
    "\x03\x2a"
  LE ->
    "\x03\x32"
  GE ->
    "\x03\x28"
  INT ->
    "\x03\x30"
  SELF sepc ->
    case sepcName sepc of
      DefEpName ->  "\x03\x49"
      epName -> "\x04\x49" <> encodeEpName epName
  CONTRACT ns ep | _ :: Proxy ('TOption ('TContract t) ': s) <- Proxy @out ->
    encodeWithAnns [] [epNameToRefAnn ep] [] $ "\x05\x55" <> encodeNotedT' @t ns
  TRANSFER_TOKENS ->
    "\x03\x4d"
  SET_DELEGATE ->
    "\x03\x4e"
  CREATE_CONTRACT (FullContract instr np ng)
    | _ :: Instr '[ 'TPair p g ] '[ 'TPair ('TList 'TOperation) g ] <- instr ->
    let contents =
          [ "\x05\x00" <> encodeNotedT' @p (unParamNotes np)
          , "\x05\x01" <> encodeNotedT' @g ng
          , "\x05\x02" <> encodeInstrs instr
          ]
    -- TODO [TM-96] These ^ should be encoded in the same order in which
    -- they appear in the original code
    in "\x05\x1d" <> encodeList id contents
  IMPLICIT_ACCOUNT ->
   "\x03\x1e"
  NOW ->
   "\x03\x40"
  AMOUNT ->
   "\x03\x13"
  BALANCE ->
   "\x03\x15"
  CHECK_SIGNATURE ->
   "\x03\x18"
  SHA256 ->
   "\x03\x0f"
  SHA512 ->
   "\x03\x10"
  BLAKE2B ->
   "\x03\x0e"
  HASH_KEY ->
   "\x03\x2b"
  STEPS_TO_QUOTA ->
   "\x03\x4a"
  SOURCE ->
   "\x03\x47"
  SENDER ->
   "\x03\x48"
  ADDRESS ->
   "\x03\x54"
  CHAIN_ID ->
   "\x03\x75"

-- | Iff there are non-empty annotations it increments the value's tag and
-- appends the encoded annotations.
encodeWithAnns :: [TypeAnn] -> [FieldAnn] -> [VarAnn] -> LByteString -> LByteString
encodeWithAnns tns fns vns encodedInput
  | null encodedInput = encodedInput
  | null annsList = encodedInput
  | otherwise = inputIncrem <> encodedAnns
  where
    trimEndNoAnn a lst = if null lst && a == noAnn then [] else a : lst
    tnsText = map (show . renderWEAnn) $ foldr trimEndNoAnn [] tns
    fnsText = map (show . renderWEAnn) $ foldr trimEndNoAnn [] fns
    vnsText = map (show . renderWEAnn) $ foldr trimEndNoAnn [] vns
    annsList = tnsText <> fnsText <> vnsText
    encodedAnns = encodeAsList . encodeUtf8 $ unwords annsList
    inputIncrem = (1 + LBS.head encodedInput) `LBS.cons` LBS.tail encodedInput

-- | Encode an instruction with variable annotations
encodeVarNotedInstr :: Instr inp out -> NonEmpty VarAnn -> LByteString
encodeVarNotedInstr i vns = case i of
  InstrWithNotes n a -> encodeNotedInstr a n (toList vns)
  _ -> encodeWithAnns [] [] (toList vns) $ encodeInstr i

-- | Encode an instruction with Annotations
encodeNotedInstr
  :: forall inp out. Instr inp out -> PackedNotes out -> [VarAnn] -> LByteString
encodeNotedInstr a (PackedNotes n _) vns = case (a, Proxy @out, n) of
  (SOME, _, NTOption tn _ns) ->
    encodeWithAnns [tn] [] vns $ encodeInstr a
  (NONE, _ :: Proxy ('TOption t ': s), NTOption tn ns) ->
    encodeWithAnns [tn] [] vns $ "\x05\x3e" <> encodeNotedT' @t ns
  (UNIT, _, NTUnit tn) ->
    encodeWithAnns [tn] [] vns $ encodeInstr a
  (PAIR, _, NTPair tn fn1 fn2 _ns1 _ns2) ->
    encodeWithAnns [tn] [fn1, fn2] vns $ encodeInstr a
  (LEFT, _ :: Proxy ('TOr l r ': s), NTOr tn fn1 fn2 _ns1 ns2) ->
    encodeWithAnns [tn] [fn1, fn2] vns $ "\x05\x33" <> encodeNotedT' @r ns2
  (RIGHT, _ :: Proxy ('TOr l r ': s), NTOr tn fn1 fn2 ns1 _ns2) ->
    encodeWithAnns [tn] [fn1, fn2] vns $ "\x05\x44" <> encodeNotedT' @l ns1
  (NIL, _ :: Proxy ('TList t ': s), NTList tn ns) ->
    encodeWithAnns [tn] [] vns $ "\x05\x3d" <> encodeNotedT' @t ns
  (EMPTY_SET, _ :: Proxy ('TSet t ': s), NTSet tn1 tn2) ->
    encodeWithAnns [tn1] [] vns $ "\x05\x24" <> encodeNotedT' @('Tc t) (NTc tn2)
  (EMPTY_MAP, _ :: Proxy ('TMap k v ': s), NTMap tn1 tn2 ns) ->
    encodeWithAnns [tn1] [] vns $
      "\x07\x23" <> encodeNotedT' @('Tc k) (NTc tn2) <> encodeNotedT' @v ns
  (EMPTY_BIG_MAP, _ :: Proxy ('TBigMap k v ': s), NTBigMap tn1 tn2 ns) ->
    encodeWithAnns [tn1] [] vns $
      "\x07\x72" <> encodeNotedT' @('Tc k) (NTc tn2) <> encodeNotedT' @v ns
  (PUSH (v :: Value t), _, NTc tn) ->
    "\x07\x43" <> encodeNotedT' @t (NTc tn) <> encodeValue v
  (LAMBDA (v :: Value ('TLambda i o)), _, NTLambda _tn ns1 ns2) ->
    "\x09\x31" <>
    encodeAsList (encodeNotedT' @i ns1 <> encodeNotedT' @o ns2 <> encodeValue v) <>
    encodeLength 0  -- encoding of a Variable Annotation (that we don't support)
  (CAST, _ :: Proxy (t ': s), NTc tn) ->
    "\x05\x57" <> encodeNotedT' @t (NTc tn)
  (UNPACK, _ :: Proxy ('TOption t ': s), NTOption tn ns) ->
    encodeWithAnns [tn] [] vns $ "\x05\x0d" <> encodeNotedT' @t ns
  -- NOTE: `CONTRACT` may be part of an `InstrWithNotes` with `NTOption`, but is
  -- taken care of in `encodeInstr` anyway (because it contains the note itself)
  _ -> encodeInstr a

packNotedT' :: forall (t :: T). SingI t => Notes t -> ByteString
packNotedT' = LBS.toStrict . encodeNotedT'

encodeNotedT' :: forall (t :: T). SingI t => Notes t -> LByteString
encodeNotedT' = encodeNotedST (sing @t) noAnn

-- Note: to encode field annotations we have to accept them as an additional
-- parameter because they are stored in the parent's `Notes t`, e.g. see STPair.
encodeNotedST :: Sing t -> FieldAnn -> Notes t -> LByteString
encodeNotedST st fn n = case (st, n) of
  (STc ct, NTc tn) ->
    encodeCTWithAnns (fromSingCT ct) tn fn
  (STKey, NTKey tn) ->
    encodeWithAnns [tn] [fn] [] $ "\x03\x5c"
  (STUnit, NTUnit tn) ->
    encodeWithAnns [tn] [fn] [] $ "\x03\x6c"
  (STSignature, NTSignature tn) ->
    encodeWithAnns [tn] [fn] [] $ "\x03\x67"
  (STChainId, NTChainId tn) ->
    encodeWithAnns [tn] [fn] [] $ "\x03\x74"
  (STOption a, NTOption tn ns) ->
    encodeWithAnns [tn] [fn] [] $ "\x05\x63" <> encodeNotedST a noAnn ns
  (STList a, NTList tn ns) ->
    encodeWithAnns [tn] [fn] [] $ "\x05\x5f" <> encodeNotedST a noAnn ns
  (STSet a, NTSet tn1 tn2) ->
    encodeWithAnns [tn1] [fn] [] $ "\x05\x66" <> encodeCTWithAnns (fromSingCT a) tn2 noAnn
  (STOperation, NTOperation tn) ->
    encodeWithAnns [tn] [fn] [] $ "\x03\x6d"
  (STContract a, NTContract tn ns) ->
    encodeWithAnns [tn] [fn] [] $ "\x05\x5a" <> encodeNotedST a noAnn ns
  (STPair a b, NTPair tn fn1 fn2 ns1 ns2) ->
    encodeWithAnns [tn] [fn] [] $
      "\x07\x65" <> encodeNotedST a fn1 ns1 <> encodeNotedST b fn2 ns2
  (STOr a b, NTOr tn fn1 fn2 ns1 ns2) ->
    encodeWithAnns [tn] [fn] [] $
      "\x07\x64" <> encodeNotedST a fn1 ns1 <> encodeNotedST b fn2 ns2
  (STLambda a r, NTLambda tn ns1 ns2) ->
    encodeWithAnns [tn] [fn] [] $
      "\x07\x5e" <> encodeNotedST a noAnn ns1 <> encodeNotedST r noAnn ns2
  (STMap k v, NTMap tn1 tn2 ns) ->
    encodeWithAnns [tn1] [fn] [] $
      "\x07\x60" <> encodeCTWithAnns (fromSingCT k) tn2 noAnn <> encodeNotedST v noAnn ns
  (STBigMap k v, NTBigMap tn1 tn2 ns) ->
    encodeWithAnns [tn1] [fn] [] $
      "\x07\x61" <> encodeCTWithAnns (fromSingCT k) tn2 noAnn <> encodeNotedST v noAnn ns

encodeCTWithAnns :: CT -> TypeAnn -> FieldAnn -> LByteString
encodeCTWithAnns ct tn fn = encodeWithAnns [tn] [fn] [] $ encodeCT ct

encodeT :: T -> LByteString
encodeT = \case
  Tc ct -> encodeCT ct
  TKey  -> "\x03\x5c"
  TUnit -> "\x03\x6c"
  TSignature -> "\x03\x67"
  TChainId -> "\x03\x74"
  TOption t -> "\x05\x63" <> encodeT t
  TList t -> "\x05\x5f" <> encodeT t
  TSet t -> "\x05\x66" <> encodeCT t
  TOperation -> "\x03\x6d"
  TContract t -> "\x05\x5a" <> encodeT t
  TPair a b -> "\x07\x65" <> encodeT a <> encodeT b
  TOr a b -> "\x07\x64" <> encodeT a <> encodeT b
  TLambda a r -> "\x07\x5e" <> encodeT a <> encodeT r
  TMap k v -> "\x07\x60" <> encodeCT k <> encodeT v
  TBigMap k v -> "\x07\x61" <> encodeCT k <> encodeT v

encodeT' :: forall (t :: T). SingI t => LByteString
encodeT' = encodeT (demote @t)

encodeCT :: CT -> LByteString
encodeCT = ("\x03" <>) . \case
  CInt -> "\x5b"
  CNat -> "\x62"
  CString -> "\x68"
  CBytes -> "\x69"
  CMutez -> "\x6a"
  CBool -> "\x59"
  CKeyHash -> "\x5d"
  CTimestamp -> "\x6b"
  CAddress -> "\x6e"