-- | Address in Tezos.

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

  -- * Formatting
  , ParseAddressError (..)
  , ParseContractAddressError
  , formatAddress
  , mformatAddress
  , parseAddress
  , 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 qualified Data.ByteString as BS
import Fmt (fmt, hexF, pretty)
import qualified Formatting.Buildable as Buildable
import Test.QuickCheck (Arbitrary(..), oneof, vector)

import Michelson.Text
import Tezos.Crypto

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

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

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

-- | Smart constructor for 'ContractAddress'. Its argument is
-- serialized origination operation.
--
-- Note: it's quite unsafe to pass 'ByteString', because we can pass
-- some garbage which is not a serialized origination operation, but
-- this operation includes contract itself and necessary types are
-- defined in 'Michelson.*'. So we have to serialize this data outside
-- this module and pass it here as a 'ByteString'. Alternatively we
-- could add some constraint, but it would be almost as unsafe as
-- passing a 'ByteString'. For this reason we add `Raw` suffix to this
-- function and provide a safer function in 'Michelson.Untyped.Instr'.
-- We may reconsider it later.
mkContractAddressRaw :: ByteString -> Address
mkContractAddressRaw = ContractAddress . mkContractHashRaw

-- | Create a dummy 'ContractHash' value.
mkContractHashRaw :: ByteString -> ContractHash
mkContractHashRaw = ContractHash . blake2b160 . blake2b160

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

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.Buildable Address where
  build = Buildable.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)

instance Buildable.Buildable ParseAddressError where
  build =
    \case
      ParseAddressWrongBase58Check -> "Wrong base58check format"
      ParseAddressBothFailed pkErr contractErr ->
        mconcat
        [ "Address is neither `KeyAddress` ("
        , Buildable.build pkErr
        , "), nor `ContractAddress` ("
        , Buildable.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)

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

data ParseContractAddressError
  = ParseContractAddressWrongBase58Check
  | ParseContractAddressWrongTag ByteString
  | ParseContractAddressWrongSize Int
  deriving stock (Show, Eq)

instance Buildable.Buildable ParseContractAddressError where
  build =
    \case
      ParseContractAddressWrongBase58Check ->
        "Wrong base58check format"
      ParseContractAddressWrongTag tag ->
        "Wrong tag for a contract address: " <> fmt (hexF tag)
      ParseContractAddressWrongSize s ->
        "Wrong size for a contract address: " <> Buildable.build s

parseContractHash :: Text -> Either ParseContractAddressError ContractHash
parseContractHash text =
  case decodeBase58CheckWithPrefix contractAddressPrefix text of
    Left (B58CheckWithPrefixWrongPrefix prefix) ->
      Left (ParseContractAddressWrongTag prefix)
    Left B58CheckWithPrefixWrongEncoding ->
      Left ParseContractAddressWrongBase58Check
    -- We know that the length must be 20.
    -- Currently it's hardcoded here, later we'll probably have a `Hash` type.
    Right bs | length bs == 20 -> Right (ContractHash bs)
             | otherwise -> Left $ ParseContractAddressWrongSize (length bs)

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

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

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

instance ToJSON ContractHash where
  toJSON = Aeson.String . formatContractHash
  toEncoding = Aeson.text . formatContractHash

instance ToJSONKey ContractHash where
  toJSONKey = AesonTypes.toJSONKeyText formatContractHash

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

instance FromJSONKey ContractHash where
  fromJSONKey =
    AesonTypes.FromJSONKeyTextParser
      (either (fail . pretty) pure . parseContractHash)


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)

----------------------------------------------------------------------------
-- Arbitrary
----------------------------------------------------------------------------

instance Arbitrary Address where
  arbitrary = oneof [genKeyAddress, genContractAddress]
    where
      genKeyAddress = KeyAddress <$> arbitrary
      genContractAddress = ContractAddress . ContractHash . BS.pack <$> vector 20