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

-- | Address in Tezos.

module Tezos.Address
  ( ContractHash (..)
  , Address (..)
  , mkKeyAddress
  , detGenKeyAddress

  , OperationHash (..)
  , OriginationIndex (..)
  , GlobalCounter(..)
  , mkContractAddress
  , mkContractHashHack

  -- * Formatting
  , ParseAddressError (..)
  , ParseAddressRawError (..)
  , ParseContractAddressError (..)
  , formatAddress
  , mformatAddress
  , parseAddressRaw
  , parseContractHash
  , parseAddress
  , unsafeParseAddressRaw
  , unsafeParseAddress
  , unsafeParseContractHash
  ) where

import Data.Aeson (FromJSON(..), FromJSONKey, ToJSON(..), ToJSONKey)
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Encoding as Aeson
import qualified Data.Aeson.Types as AesonTypes
import Data.Binary.Put (putInt32be, putWord64be, runPut)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Fmt (Buildable(build), hexF, pretty)

import Michelson.Text
import Tezos.Crypto
import Util.CLI

-- TODO: we should probably have a `Hash` type.
-- | Hash of origination command for some contract.
newtype ContractHash = ContractHash ByteString
  deriving stock (Show, Eq, Ord, Generic)

instance NFData ContractHash

-- We know that the length must be 20.
-- Currently it's hardcoded here, later we'll probably have a `Hash` type.
-- | Length of contract hash in bytes (only hash itself, no tags, checksums
-- or anything).
contractHashLengthBytes :: Integral n => n
contractHashLengthBytes = 20

-- | Data type corresponding to address structure in Tezos.
data Address
  = KeyAddress KeyHash
  -- ^ `tz` address which is a hash of a public key.
  | ContractAddress ContractHash
  -- ^ `KT` address which corresponds to a callable contract.
  deriving stock (Show, Eq, Ord, Generic)

instance NFData Address

-- | Smart constructor for 'KeyAddress'.
mkKeyAddress :: PublicKey -> Address
mkKeyAddress = KeyAddress . hashKey

-- | Deterministically generate a random 'KeyAddress' and discard its
-- secret key.
detGenKeyAddress :: ByteString -> Address
detGenKeyAddress = mkKeyAddress . toPublic . detSecretKey

newtype OperationHash = OperationHash
  { unOperationHash :: ByteString
  }
  deriving stock (Show, Eq, Ord, Generic)
  deriving anyclass (NFData)

-- | Represents the network's global counter.
--
-- When a new contract is created (either via a "global" origination operation or
-- via a @CREATE_CONTRACT@ instruction), this counter is used to create a new address for it
-- (see 'mkContractAddress').
--
-- The counter is incremented after every operation, and thus ensures that these
-- addresses are unique (i.e. origination of identical contracts with identical metadata will
-- result in different addresses.)
--
-- In Tezos each operation has a special field called @counter@, see here:
-- https://gitlab.com/tezos/tezos/-/blob/397dd233a10cc6df0df959e2a624c7947997dd0c/src/proto_006_PsCARTHA/lib_protocol/operation_repr.ml#L113-120
--
-- This counter seems to be a part of global state of Tezos network. In fact, it may be observed
-- in raw JSON representation of the operation in the network explorer.
--
-- Our counter is represented as 'Word64', while in Tezos it is unbounded. We believe that
-- for our interpreter it should not matter.
newtype GlobalCounter = GlobalCounter { unGlobalCounter :: Word64 }
  deriving stock (Show, Eq, Generic)
  deriving anyclass (NFData)
  deriving newtype (ToJSON, FromJSON, Num)

-- | When a transfer operation triggers multiple @CREATE_CONTRACT@
-- instructions, using 'GlobalCounter' to compute those contracts' addresses
-- is not enough to ensure their uniqueness.
--
-- For that reason, we also keep track of an 'OriginationIndex' that starts out as 0
-- when a transfer is initiated, and is incremented every time a @CREATE_CONTRACT@
-- instruction is interpreted.
--
-- See 'mkContractAddress'.
newtype OriginationIndex = OriginationIndex { unOriginationIndex :: Int32 }
  deriving stock (Show, Eq, Ord, Generic)
  deriving anyclass (NFData)

-- | Compute address of a contract from its origination operation, origination index and global counter.
--
-- However, in real Tezos encoding of the operation is more than just 'OriginationOperation'.
-- There an Operation has several more meta-fields plus a big sum-type of all possible operations.
--
-- See here: https://gitlab.com/tezos/tezos/-/blob/f57c50e3a657956d69a1699978de9873c98f0018/src/proto_006_PsCARTHA/lib_protocol/operation_repr.ml#L78
--
-- What is important is that one (big) Operation may lead to origination of multiple contracts. That
-- is why contract address is constructed from hash of the operation that originated and of index
-- of the contract's origination in the execution of that operation.
--
-- In other words, contract hash is calculated as the blake2b160 (20-byte) hash of
-- origination operation hash + int32 origination index + word64 global counter.
--
-- In Morley we do not yet support full encoding of Tezos Operations, therefore we choose
-- to generate contract addresses in a simplified manner.
--
-- Namely, we encode 'OriginationOperation' as we can and concat it with the origination index
-- and the global counter.
-- Then we take 'blake2b160' hash of the resulting bytes and consider it to be the contract's
-- address.
mkContractAddress
  :: OperationHash
  -> OriginationIndex
  -> GlobalCounter
  -> Address
mkContractAddress (OperationHash opHash) (OriginationIndex idx) (GlobalCounter counter) =
  ContractAddress
  $ ContractHash
  $ blake2b160
  $ opHash <> BSL.toStrict (runPut $ putInt32be idx >> putWord64be counter)

-- | Create a dummy 'ContractHash' value by hashing given 'ByteString'.
--
-- Use in tests **only**.
mkContractHashHack :: ByteString -> ContractHash
mkContractHashHack = ContractHash . blake2b160

----------------------------------------------------------------------------
-- Formatting/parsing
----------------------------------------------------------------------------

-- It's a magic constant used by Tezos to encode a contract address.
-- It was deduced empirically.
contractAddressPrefix :: ByteString
contractAddressPrefix = "\2\90\121"

formatContractHash :: ContractHash -> Text
formatContractHash (ContractHash bs) =
  encodeBase58Check (contractAddressPrefix <> bs)

formatAddress :: Address -> Text
formatAddress =
  \case
    KeyAddress h -> formatKeyHash h
    ContractAddress h -> formatContractHash h

mformatAddress :: Address -> MText
mformatAddress = mkMTextUnsafe . formatAddress

instance Buildable Address where
  build = build . formatAddress

-- | Errors that can happen during address parsing.
data ParseAddressError
  = ParseAddressWrongBase58Check
  -- ^ Address is not in Base58Check format.
  | ParseAddressBothFailed CryptoParseError ParseContractAddressError
  -- ^ Both address parsers failed with some error.
  deriving stock (Show, Eq, Generic)

instance NFData ParseAddressError

instance Buildable ParseAddressError where
  build =
    \case
      ParseAddressWrongBase58Check -> "Wrong base58check format"
      ParseAddressBothFailed pkErr contractErr ->
        mconcat
        [ "Address is neither `KeyAddress` ("
        , build pkErr
        , "), nor `ContractAddress` ("
        , build contractErr
        , ")"
        ]

-- | Parse an address from its human-readable textual representation
-- used by Tezos (e. g. "tz1faswCTDciRzE4oJ9jn2Vm2dvjeyA9fUzU"). Or
-- fail if it's invalid.
parseAddress :: Text -> Either ParseAddressError Address
parseAddress addressText =
  case parseKeyHash addressText of
    Left CryptoParseWrongBase58Check -> Left ParseAddressWrongBase58Check
    Left keyAddrErr -> first (ParseAddressBothFailed keyAddrErr) $
      ContractAddress <$> parseContractHash addressText
    Right keyHash -> Right (KeyAddress keyHash)

data ParseAddressRawError
  = ParseAddressRawWrongSize ByteString
  -- ^ Raw bytes representation of an address has invalid length.
  | ParseAddressRawInvalidPrefix ByteString
  -- ^ Raw bytes representation of an address has incorrect prefix.
  | ParseAddressRawMalformedSeparator ByteString
  -- ^ Raw bytes representation of an address does not end with "\00".
  deriving stock (Eq, Show, Generic)

instance NFData ParseAddressRawError

instance Buildable ParseAddressRawError where
  build =
    \case
      ParseAddressRawInvalidPrefix prefix ->
        "Invalid prefix for raw address \"" <> hexF prefix <> "\" provided"
      ParseAddressRawWrongSize addr -> "Given raw address " <>
        hexF addr <> " has invalid length " <> show (length addr)
      ParseAddressRawMalformedSeparator addr -> "Given raw address " <> hexF addr <>
        " does not end with \"\\00\""

-- | Parse the given address in its raw byte form used by Tezos
-- (e.g "01521139f84791537d54575df0c74a8084cc68861c00")) . Or fail otherwise
-- if it's invalid.
parseAddressRaw :: ByteString -> Either ParseAddressRawError Address
parseAddressRaw (BS.splitAt 2 -> (prefix, address)) =
  case BS.take 1 prefix of
    "\00" -> parseKeyAddressRaw (BS.tail prefix) address
    "\01" -> parseContractAddressRaw prefix address
    _ -> Left $ ParseAddressRawInvalidPrefix prefix
  where
    parseKeyAddressRaw keyPrefix keyAddress
      | length keyAddress /= keyHashLengthBytes
          = Left $ ParseAddressRawWrongSize keyAddress
      | otherwise = do
        matchedPrefix <- case keyPrefix of
          "\00" -> Right KeyHashEd25519
          "\01" -> Right KeyHashSecp256k1
          "\02" -> Right KeyHashP256
          _ -> Left $ ParseAddressRawInvalidPrefix keyPrefix
        pure $ KeyAddress (KeyHash matchedPrefix keyAddress)

    parseContractAddressRaw contractPrefix contractAddress
      | length contractAddress /= contractHashLengthBytes
          = Left $ ParseAddressRawWrongSize contractAddress
      | BS.last contractAddress /= 0x00
          = Left $ ParseAddressRawMalformedSeparator contractAddress
      | otherwise = do
        let contractAddress' = BS.drop 1 contractPrefix <> BS.init contractAddress -- drop last "\00" of contract address
        pure $ ContractAddress (ContractHash contractAddress')

data ParseContractAddressError
  = ParseContractAddressWrongBase58Check
  | ParseContractAddressWrongSize ByteString
  | ParseContractAddressWrongPrefix ByteString
  deriving stock (Show, Eq, Generic)

instance NFData ParseContractAddressError

instance Buildable ParseContractAddressError where
  build =
    \case
      ParseContractAddressWrongBase58Check ->
        "Wrong base58check format"
      ParseContractAddressWrongSize bs ->
        "Wrong size for a contract address: " <> hexF bs <>
        " (" <> build (length bs) <> ")"
      ParseContractAddressWrongPrefix prefix ->
        "Invalid prefix \"" <> hexF prefix <> "\" provided"

parseContractHash :: Text -> Either ParseContractAddressError ContractHash
parseContractHash text =
  case decodeBase58CheckWithPrefix contractAddressPrefix text of
    Left (B58CheckWithPrefixWrongPrefix prefix) ->
      Left (ParseContractAddressWrongPrefix prefix)
    Left B58CheckWithPrefixWrongEncoding ->
      Left ParseContractAddressWrongBase58Check
    Right bs | length bs == contractHashLengthBytes -> Right (ContractHash bs)
             | otherwise -> Left $ ParseContractAddressWrongSize bs

----------------------------------------------------------------------------
-- Unsafe
----------------------------------------------------------------------------

-- | Parse a @KT1@ contract address, fail if address does not match
-- the expected format.
unsafeParseContractHash :: HasCallStack => Text -> ContractHash
unsafeParseContractHash = either (error . pretty) id . parseContractHash

-- | Partial version of 'parseAddress' which assumes that the address
-- is correct. Can be used in tests.
unsafeParseAddress :: HasCallStack => Text -> Address
unsafeParseAddress = either (error . pretty) id . parseAddress

-- | Partially parse raw bytes representation of an address and assume that
-- it is correct from the beginning. Can be used in tests.
unsafeParseAddressRaw :: ByteString -> Address
unsafeParseAddressRaw = either (error . pretty) id . parseAddressRaw

instance HasCLReader Address where
  getReader = eitherReader parseAddrDo
    where
      parseAddrDo addr =
        either (Left . mappend "Failed to parse address: " . pretty) Right $
        parseAddress $ toText addr
  getMetavar = "ADDRESS"

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

instance ToJSON Address where
  toJSON = Aeson.String . formatAddress
  toEncoding = Aeson.text . formatAddress

instance ToJSONKey Address where
  toJSONKey = AesonTypes.toJSONKeyText formatAddress

instance FromJSON Address where
  parseJSON =
    Aeson.withText "Address" $
    either (fail . pretty) pure . parseAddress

instance FromJSONKey Address where
  fromJSONKey =
    AesonTypes.FromJSONKeyTextParser
      (either (fail . pretty) pure . parseAddress)