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

-- | Address in Tezos.

{-# LANGUAGE DeriveLift #-}

module Morley.Tezos.Address
  ( ContractHash
  , KindedAddress (..)
  , TxRollupL2Address (..)
  , mkKeyAddress
  , detGenKeyAddress
  , isImplicitAddress
  , ImplicitAddress
  , ContractAddress
  , TxRollupAddress
  , L1Address
  , L1AddressKind
  , ConstrainAddressKind
  , Address
  , ConstrainedAddress(.., MkAddress)

  , GlobalCounter(..)
  , mkContractHashHack

  -- * Formatting
  , ParseAddressError (..)
  , ParseAddressRawError (..)
  , formatAddress
  , mformatAddress
  , parseAddressRaw
  , parseKindedAddress
  , parseAddress
  , ta

  -- * Utilities
  , addressKindSanity
  , usingImplicitOrContractKind
  ) where

import Control.Monad.Except (mapExceptT, throwError)
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.Binary.Get qualified as Get
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as LBS
import Data.Constraint (Dict(..), (\\))
import Data.Singletons (SingI(..))
import Data.Text (strip)
import Data.Type.Equality ((:~:)(..))
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, (<+>))

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

-- | A "kinded" address. This type carries 'AddressKind' on the type-level.
-- Useful in the internal API, not as much when we have to interact with the
-- network. See 'Address' for a type that is isomorphic to a Michelson
-- @address@.
data KindedAddress (kind :: AddressKind) where
  -- | @tz1@, @tz2@ or @tz3@ address which is a hash of a public key.
  ImplicitAddress :: KeyHash -> KindedAddress 'AddressKindImplicit
  -- | @KT1@ address which corresponds to a callable contract.
  ContractAddress :: ContractHash -> KindedAddress 'AddressKindContract
  -- | @txr1@ address which corresponds to a transaction rollup.
  TxRollupAddress :: TxRollupHash -> KindedAddress 'AddressKindTxRollup

deriving stock instance Show (KindedAddress kind)
deriving stock instance Eq (KindedAddress kind)
deriving stock instance Ord (KindedAddress kind)
deriving stock instance Lift (KindedAddress kind)

instance NFData (KindedAddress kind) where
  rnf = \case
    ImplicitAddress x -> rnf x
    ContractAddress x -> rnf x
    TxRollupAddress x -> rnf x

-- | A type only allowing v'ImplicitAddress'
type ImplicitAddress = KindedAddress 'AddressKindImplicit

-- | A type only allowing v'ContractAddress'
type ContractAddress = KindedAddress 'AddressKindContract

-- | A type only allowing v'TxRollupAddress'
type TxRollupAddress = KindedAddress 'AddressKindTxRollup

-- | Data type corresponding to @address@ structure in Tezos.
type Address =
  ConstrainedAddress '[ 'AddressKindImplicit, 'AddressKindContract, 'AddressKindTxRollup ]

type family ConstrainAddressKindHelper (ks :: [AddressKind]) kind where
  ConstrainAddressKindHelper (x ': _) x = 'True
  ConstrainAddressKindHelper (_ ': xs) x = ConstrainAddressKindHelper xs x
  ConstrainAddressKindHelper '[] _ = 'False

type family CheckConstrainAddressKindError k b :: Constraint where
  CheckConstrainAddressKindError _ 'True = ()
  CheckConstrainAddressKindError k 'False =
    TypeError ('ShowType k ':<>: 'Text "is forbidden in this context")

-- | Constrain address kind to be one of the kinds in the list.
type ConstrainAddressKind ks k =
  ( CheckConstrainAddressKindError k (ConstrainAddressKindHelper ks k)
  , ConstrainAddressKindHelper ks k ~ 'True)

-- | An existential of 'KindedAddress' constrained by its type argument.
data ConstrainedAddress (ks :: [AddressKind]) =
  forall kind. ConstrainAddressKind ks kind => MkConstrainedAddress (KindedAddress kind)

-- | A convenience synonym for 'ConstrainedAddress' allowing only implicit and
-- contract addresses.
--
-- 'L1Address' is named as such because in addition to implicit and contract
-- addresses, Michelson's @address@ type can contain @txr1@ addresses,
-- identifying transaction rollups. While @txr1@ are technically also level-1
-- (level-2 being @tx_rollup_l2_address@, i.e. @tz4@), in practice it's a
-- level-1 identifier for a bundle of level-2 operations. Hence, to keep type
-- names concise, we use 'L1Address'.
type L1Address =
  ConstrainedAddress '[ 'AddressKindImplicit, 'AddressKindContract ]

-- | Convenience synonym for 'ConstrainAddressKind' allowing only implicit and
-- contract addresses.
--
-- For a note on the naming convention, refer to 'L1Address'.
type L1AddressKind kind =
  ConstrainAddressKind '[ 'AddressKindImplicit, 'AddressKindContract ] kind

-- | A trick to avoid bogus redundant constraint warnings
usingImplicitOrContractKind :: forall kind a. L1AddressKind kind => a -> a
usingImplicitOrContractKind = id
  where _ = Dict :: Dict (L1AddressKind kind)

-- | 'MkConstrainedAddress' specialized to 'Address'
pattern MkAddress :: KindedAddress kind -> Address
pattern MkAddress x <- MkConstrainedAddress x
  where MkAddress x = case x of
          ImplicitAddress{} -> MkConstrainedAddress x
          ContractAddress{} -> MkConstrainedAddress x
          TxRollupAddress{} -> MkConstrainedAddress x

{-# COMPLETE MkAddress #-}

deriving stock instance Show (ConstrainedAddress c)
instance Eq (ConstrainedAddress c) where
  MkConstrainedAddress (addr1 :: KindedAddress kind1)
    == MkConstrainedAddress (addr2 :: KindedAddress kind2) =
    case eqI @kind1 @kind2 \\ addressKindSanity addr1 \\ addressKindSanity addr2 of
      Just Refl -> addr1 == addr2
      Nothing -> False
instance Ord (ConstrainedAddress c) where
  compare (MkConstrainedAddress (a1 :: KindedAddress kind1))
          (MkConstrainedAddress (a2 :: KindedAddress kind2)) =
    case (a1, a2) of
      (k1, k2)
        | Just Refl <- eqI @kind1 @kind2 \\ addressKindSanity k1 \\ addressKindSanity k2
        -> compare k1 k2
      (ImplicitAddress{}, _) -> LT
      (ContractAddress{}, ImplicitAddress{}) -> GT
      (ContractAddress{}, _) -> LT
      (TxRollupAddress{}, ImplicitAddress{}) -> GT
      (TxRollupAddress{}, ContractAddress{}) -> GT
      (TxRollupAddress{}, _) -> LT
deriving stock instance Lift (ConstrainedAddress c)

instance NFData (ConstrainedAddress c) where
  rnf (MkConstrainedAddress x) = rnf x

instance Buildable (ConstrainedAddress c) where
  build (MkConstrainedAddress x) = build x

-- | Given any (non-bottom) 'KindedAddress', prove that @kind@ is well-defined
-- (i.e. has a 'SingI' instance)
addressKindSanity :: KindedAddress kind -> Dict (SingI kind)
addressKindSanity = \case
  ImplicitAddress{} -> Dict
  ContractAddress{} -> Dict
  TxRollupAddress{} -> Dict

-- | @tz4@ level-2 public key hash address, used with transaction rollups, corresponds
-- to @tx_rollup_l2_address@ Michelson type.
newtype TxRollupL2Address = TxRollupL2Address KeyHashL2
  deriving stock (Show, Eq, Ord, Generic, Lift)
  deriving newtype NFData

-- | Checks if the provided 'KindedAddress' is an implicit address and returns
-- proof of the fact if it is.
isImplicitAddress :: KindedAddress kind -> Maybe (kind :~: 'AddressKindImplicit)
isImplicitAddress = \case
  ImplicitAddress{} -> Just Refl
  _ -> Nothing

-- | Smart constructor for t'ImplicitAddress'.
mkKeyAddress :: PublicKey -> ImplicitAddress
mkKeyAddress = ImplicitAddress . hashKey

-- | Deterministically generate a random t'ImplicitAddress' and discard its
-- secret key.
detGenKeyAddress :: ByteString -> ImplicitAddress
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 = Hash HashContract . blake2b160

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

formatAddress :: KindedAddress kind -> Text
formatAddress =
  \case
    ImplicitAddress h -> formatHash h
    ContractAddress h -> formatHash h
    TxRollupAddress h -> formatHash h

mformatAddress :: KindedAddress kind -> MText
mformatAddress = unsafe . mkMText . formatAddress

instance Buildable (KindedAddress kind) where
  build = build . formatAddress

instance Buildable TxRollupL2Address where
  build (TxRollupL2Address kh) = build $ formatHash kh

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

instance Semigroup ParseAddressError where
  ParseAddressWrongBase58Check <> _ =  ParseAddressWrongBase58Check
  _ <> ParseAddressWrongBase58Check =  ParseAddressWrongBase58Check
  ParseAddressAllFailed xs <> ParseAddressAllFailed ys = ParseAddressAllFailed $ xs <> ys

instance NFData ParseAddressError

instance Buildable ParseAddressError where
  build = buildRenderDoc

instance RenderDoc ParseAddressError where
  renderDoc context =
    \case
      ParseAddressWrongBase58Check -> "Wrong base58check format"
      ParseAddressAllFailed pkErr ->
        mconcat $ "Address failed to parse: " : intersperse ", "
          (toList $ renderDoc context <$> pkErr)

-- | Parse an address of a particular kind from its human-readable textual
-- representation used by Tezos (e. g. "tz1faswCTDciRzE4oJ9jn2Vm2dvjeyA9fUzU").
-- Or fail if it's invalid.
parseKindedAddress :: forall kind. SingI kind => Text -> Either ParseAddressError (KindedAddress kind)
parseKindedAddress addressText = case sing @kind of
  SAddressKindContract -> tryParse ContractAddress
  SAddressKindImplicit -> tryParse ImplicitAddress
  SAddressKindTxRollup -> tryParse TxRollupAddress
  where
    handleCrypto = \case
      CryptoParseWrongBase58Check -> ParseAddressWrongBase58Check
      x -> ParseAddressAllFailed $ pure x
    tryParse :: AllHashTags hkind => (Hash hkind -> b) -> Either ParseAddressError b
    tryParse ctor = bimap handleCrypto ctor $ parseHash addressText

-- | Parse an address of arbitrary kind from its human-readable textual
-- representation, or fail if it's invalid.
parseAddress :: Text -> Either ParseAddressError Address
parseAddress x =
          (MkAddress <$> parseKindedAddress @'AddressKindImplicit x)
  `merge` (MkAddress <$> parseKindedAddress @'AddressKindContract x)
  `merge` (MkAddress <$> parseKindedAddress @'AddressKindTxRollup x)
  where
    merge :: Semigroup a => Either a b -> Either a b -> Either a b
    merge (Left xs) (Left ys) = Left $ xs <> ys
    merge r@Right{} _ = r
    merge _ r@Right{} = r

data ParseAddressRawError
  = ParseAddressRawWrongSize ByteString
  -- ^ Raw bytes representation of an address has invalid length.
  | ParseAddressRawInvalidPrefix Word8
  -- ^ Raw bytes representation of an address has incorrect prefix.
  | ParseAddressRawMalformedSeparator Word8
  -- ^ Raw bytes representation of an address does not end with "\00".
  | ParseAddressRawBinaryError Text
  -- ^ General binary decoding error.
  | ParseAddressCryptoError CryptoParseError
  -- ^ Crypto error in parsing key hash.
  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")
      ParseAddressRawBinaryError err -> "Binary error during decoding address:" <+> renderAnyBuildable err
      ParseAddressCryptoError err -> "Key hash decoding error:" <+> renderAnyBuildable err

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 bytes
  -- NB: conveniently, the byte count is the same for 'KeyAddress',
  -- 'ContractAddress' and 'TransactionRollupAddress'. However, with
  -- 'KeyAddress' it's two tag bytes, while with the other two it's one tag byte
  -- and one separator byte.
  | BS.length bytes /= hashLengthBytes + 2 = Left $ ParseAddressRawWrongSize bytes
  | otherwise
  = either (Left . ParseAddressRawBinaryError . fromString . view _3) (view _3)
  $ flip Get.runGetOrFail (LBS.fromStrict bytes) $ runExceptT
  $ decodeWithTagM "address" (throwError . ParseAddressRawInvalidPrefix)
      [ 0x00 ##: MkAddress . ImplicitAddress <$> keyHash
      , 0x01 ##: MkAddress . ContractAddress <$> sepHash HashContract
      , 0x02 ##: MkAddress . TxRollupAddress <$> sepHash HashTXR
      ]
  where
    sep = lift Get.getWord8 >>= \case
      0x00 -> pass
      x -> throwError $ ParseAddressRawMalformedSeparator x
    keyHash = mapExceptT (fmap $ first ParseAddressCryptoError) decodeKeyHash
    sepHash :: HashTag kind -> ExceptT ParseAddressRawError Get.Get (Hash kind)
    sepHash kind = Hash kind <$> lift (getByteStringCopy hashLengthBytes) <* sep

-- | QuasiQuoter 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 (MkAddress addr) -> TH.lift addr
  , TH.quotePat = \_ ->
      fail "Cannot use this QuasiQuotation at pattern position"
  , TH.quoteType = \_ ->
      fail "Cannot use this QuasiQuotation at type position"
  , TH.quoteDec = \_ ->
      fail "Cannot use this QuasiQuotation at declaration position"
  }


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


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

instance SingI kind => HasCLReader (KindedAddress kind) where
  getReader = eitherReader parseAddrDo
    where
      parseAddrDo addr =
        first (mappend "Failed to parse address: " . pretty) $
        parseKindedAddress $ toText addr
  getMetavar = "ADDRESS"

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

instance ToJSON (KindedAddress kind) where
  toJSON = Aeson.String . formatAddress
  toEncoding = Aeson.text . formatAddress

instance ToJSONKey (KindedAddress kind) where
  toJSONKey = AesonTypes.toJSONKeyText formatAddress

instance SingI kind => FromJSON (KindedAddress kind) where
  parseJSON =
    Aeson.withText "Address" $
    either (fail . pretty) pure . parseKindedAddress

instance SingI kind => FromJSONKey (KindedAddress kind) where
  fromJSONKey =
    AesonTypes.FromJSONKeyTextParser
      (either (fail . pretty) pure . parseKindedAddress)

instance ToJSON (ConstrainedAddress c) where
  toJSON (MkConstrainedAddress addr) = toJSON addr
  toEncoding (MkConstrainedAddress addr) = toEncoding addr

instance ToJSONKey (ConstrainedAddress c) where
  toJSONKey = AesonTypes.toJSONKeyText \(MkConstrainedAddress addr) -> formatAddress addr

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)