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

-- | Address in Tezos.

{-# LANGUAGE DeriveLift #-}

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

  , GlobalCounter(..)
  , mkContractHashHack

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

import Data.Aeson (FromJSON(..), FromJSONKey, ToJSON(..), ToJSONKey)
import Data.Aeson qualified as Aeson
import Data.Aeson.Encoding qualified as Aeson
import Data.Aeson.Types qualified as AesonTypes
import Data.ByteString qualified as BS
import Data.Text (strip)
import Fmt (Buildable(build), hexF, pretty)
import Instances.TH.Lift ()
import Language.Haskell.TH.Quote qualified as TH
import Language.Haskell.TH.Syntax (Lift)
import Language.Haskell.TH.Syntax qualified as TH
import Text.PrettyPrint.Leijen.Text (backslash, dquotes, int, parens, (<+>))

import Morley.Michelson.Printer.Util (RenderDoc(..), buildRenderDoc, renderAnyBuildable)
import Morley.Michelson.Text
import Morley.Tezos.Crypto
import Morley.Util.CLI
import Morley.Util.TypeLits

-- 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, Lift)

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, Lift)

instance NFData Address

-- | Returns @True@ if given address is implicit.
isKeyAddress :: Address -> Bool
isKeyAddress = \case
  KeyAddress _      -> True
  ContractAddress _ -> False

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

-- | Represents the network's global counter.
--
-- We store the current value of this counter in the operation at the time of its creation
-- for the following reasons:
-- * to guarantee the uniqueness of contract addresses upon origination
--   (see 'Morley.Michelson.Typed.Operation.mkContractAddress)
-- * to prevent replay attacks by checking that an operation with the same counter value
--   con't be performed twice.
--
-- The counter is incremented after every operation execution and interpretation of instructions
-- @CREATE_CONTRACT@ and @TRANSFER_TOKENS@, and thus ensures that these addresses are unique
-- (i.e. origination of identical contracts with identical metadata will result in
-- different addresses.)
--
-- 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, Buildable, Hashable)

-- | 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 = unsafe . mkMText . 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 = buildRenderDoc

instance RenderDoc ParseAddressError where
  renderDoc context =
    \case
      ParseAddressWrongBase58Check -> "Wrong base58check format"
      ParseAddressBothFailed pkErr contractErr ->
        mconcat
        [ "Address is neither `KeyAddress` "
        , parens $ renderDoc context pkErr
        , ", nor `ContractAddress` "
        , parens $ renderDoc context 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 RenderDoc ParseAddressRawError where
  renderDoc _ =
    \case
      ParseAddressRawInvalidPrefix prefix ->
        "Invalid prefix for raw address" <+> (dquotes $ renderAnyBuildable $ hexF prefix) <+> "provided"
      ParseAddressRawWrongSize addr -> "Given raw address+" <+>
        (renderAnyBuildable $ hexF addr) <+> "has invalid length" <+> int (length addr)
      ParseAddressRawMalformedSeparator addr -> "Given raw address+" <+> (renderAnyBuildable $ hexF addr) <+>
        "does not end with" <+> dquotes (backslash <> "00")

instance Buildable ParseAddressRawError where
  build = buildRenderDoc

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

instance RenderDoc ParseContractAddressError where
  renderDoc _ =
    \case
      ParseContractAddressWrongBase58Check ->
        "Wrong base58check format"
      ParseContractAddressWrongSize bs ->
        "Wrong size for a contract address:" <+> (renderAnyBuildable $ hexF bs) <+>
        (parens $ int (length bs))
      ParseContractAddressWrongPrefix prefix ->
        "Invalid prefix" <+> (dquotes $ renderAnyBuildable $ 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

-- | QuasyQuoter for constructing Tezos addresses.
--
-- Validity of result will be checked at compile time.
ta :: TH.QuasiQuoter
ta = TH.QuasiQuoter
  { TH.quoteExp = \s ->
      case parseAddress . strip $ toText s of
        Left   err -> fail $ pretty err
        Right addr -> TH.lift addr
  , TH.quotePat = \_ ->
      fail "Cannot use this QuasyQuotation at pattern position"
  , TH.quoteType = \_ ->
      fail "Cannot use this QuasyQuotation at type position"
  , TH.quoteDec = \_ ->
      fail "Cannot use this QuasyQuotation at declaration position"
  }


instance
    TypeError ('Text "There is no instance defined for (IsString Address)" ':$$:
               'Text "Consider using QuasiQuotes: `[ta|some text...|]`"
              ) =>
    IsString Address where
  fromString = error "impossible"


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

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)