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

-- | Address in Tezos.

{-# LANGUAGE DeriveLift #-}

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

  , GlobalCounter(..)
  , mkContractHashHack
  , parseConstrainedAddress

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

  -- * Utilities
  , addressKindSanity
  , usingImplicitOrContractKind
  , unImplicitAddress
  , addressKindTag
  ) 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.Char (toUpper)
import Data.Constraint (Bottom(..), Dict(..), (\\))
import Data.Constraint.Extras (has)
import Data.Constraint.Extras.TH (deriveArgDict)
import Data.GADT.Compare.TH (deriveGCompare, deriveGEq)
import Data.List.Singletons (SList(..))
import Data.Singletons (Sing, SingI(..), demote)
import Data.Some (Some(..))
import Data.Text (strip)
import Data.Type.Equality (testEquality, (:~:)(..))
import Fmt (Buildable(build), hexF, nameF, pretty, unwordsF, (<+>))
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 Options.Applicative (ReadM)
import Prettyprinter (backslash, dquotes)

import Morley.Michelson.Text
import Morley.Tezos.Address.Kinds
import Morley.Tezos.Crypto
import Morley.Util.Binary
import Morley.Util.CLI
import Morley.Util.Constrained
import Morley.Util.Sing
import Morley.Util.TH
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
  SmartRollupAddress :: SmartRollupHash -> KindedAddress 'AddressKindSmartRollup

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

deriveGADTNFData ''KindedAddress
deriveGEq ''KindedAddress
deriveGCompare ''KindedAddress
deriveArgDict ''KindedAddress

-- | 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'SmartRollupAddress'
type SmartRollupAddress = KindedAddress 'AddressKindSmartRollup

-- | Data type corresponding to @address@ structure in Tezos.
type Address = Constrained NullConstraint KindedAddress

-- | 'Constrained' specialized to 'Address'
pattern MkAddress :: KindedAddress kind -> Address
pattern MkAddress x = Constrained x
{-# COMPLETE MkAddress #-}

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 :: [AddressKind] -> AddressKind -> Constraint
class ( CheckConstrainAddressKindError k (ConstrainAddressKindHelper ks k)
  , ConstrainAddressKindHelper ks k ~ 'True) => ConstrainAddressKind ks k
instance ( CheckConstrainAddressKindError k (ConstrainAddressKindHelper ks k)
  , ConstrainAddressKindHelper ks k ~ 'True) => ConstrainAddressKind ks k

-- | An existential of 'KindedAddress' constrained by its type argument.
type ConstrainedAddress (ks :: [AddressKind]) =
  Constrained (ConstrainAddressKind ks) KindedAddress

-- | 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@ or @sr1@ addresses,
-- identifying respectively transaction rollups and smart rollups. While they
-- are technically also level-1 (level-2 being @tx_rollup_l2_address@), in
-- practice It's level-1 identifiers for bundles 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 = ConstrainAddressKind '[ 'AddressKindImplicit, 'AddressKindContract ]

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

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

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

unImplicitAddress :: ImplicitAddress -> KeyHash
unImplicitAddress (ImplicitAddress kh) = kh

-- | 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
    SmartRollupAddress h -> formatHash h

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

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

-- | Errors that can happen during address parsing.
data ParseAddressError
  = ParseAddressCryptoError CryptoParseError
  -- ^ The address parsers failed with some error.
  | ParseAddressWrongKind [AddressKind] Address
  -- ^ The parsed address is of wrong kind
  deriving stock (Show, Eq, Generic)

instance NFData ParseAddressError

instance Buildable ParseAddressError where
  build = \case
    ParseAddressCryptoError pkErr -> nameF "Address failed to parse" $ build pkErr
    ParseAddressWrongKind expected (Constrained a) -> unwordsF
      [ "Expected address of kind", renderAddressKinds expected
      , ", but got", build a
      ]
    where
      renderAddressKinds as = mconcat $ intersperse ", " (build <$> as)

-- | 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 = do
  Constrained a <- parseConstrainedAddress @'[kind] addressText
  castSing a \\ addressKindSanity a &
    maybeToRight (ParseAddressWrongKind [demote @kind] $ Constrained a)

-- | Parse an 'ConstrainedAddress' of the given kinds from its human-readable textual
--  representation. Maybe fail with a 'ParseAddressWrongKind' in case the address parsed
--  is of wrong kind.
parseConstrainedAddress
  :: forall kinds . (SingI kinds)
  => Text -> Either ParseAddressError (ConstrainedAddress kinds)
parseConstrainedAddress addressText =
  parseAddress addressText >>= castConstrainedAddress (demote @kinds) (sing @kinds)

castConstrainedAddress
  :: [AddressKind]
  -> SList kinds
  -> Address
  -> Either ParseAddressError (ConstrainedAddress kinds)
castConstrainedAddress allowed = \case
  SNil -> Left . ParseAddressWrongKind allowed
  SCons kind ks -> \case
    Constrained (a :: KindedAddress kind')
      | Just Refl <- testEquality kind (sing @kind') \\ addressKindSanity a
      -> Right (Constrained a)
    a -> recastAddress kind <$> castConstrainedAddress allowed ks a

recastAddress
  :: forall xs x. Sing x -> ConstrainedAddress xs -> ConstrainedAddress (x ': xs)
recastAddress sx (Constrained (x :: KindedAddress k)) =
  Constrained x \\ proofAddressCast @xs (sing @k) sx \\ addressKindSanity x

proofAddressCast
  :: forall ks k x. ConstrainAddressKind ks k
  => Sing k -> Sing x -> Dict (ConstrainAddressKind (x ': ks) k)
proofAddressCast = $(forEachAddressKind $ forEachAddressKind [|Dict|])

-- | Parse an address of arbitrary kind from its human-readable textual
-- representation, or fail if it's invalid.
parseAddress :: Text -> Either ParseAddressError Address
parseAddress a = first ParseAddressCryptoError $ parseSomeHashBase58 a <&> \case
  Some h@(Hash hk _) -> case hk of
    HashKey{}     -> Constrained $ ImplicitAddress h
    HashContract  -> Constrained $ ContractAddress h
    HashSR        -> Constrained $ SmartRollupAddress h

data ParseAddressRawError
  = ParseAddressRawWrongSize ByteString
  -- ^ Raw bytes representation of an address has invalid length.
  | ParseAddressRawInvalidPrefix Word8
  -- ^ Raw bytes representation of an address has incorrect prefix.
  | ParseAddressRawUnsupportedPrefix Text Word8
  -- ^ Unsupported address type.
  | ParseAddressRawMalformedSeparator Word8
  -- ^ Raw bytes representation of an address does not end with "\00".
  | ParseAddressRawBinaryError Text
  -- ^ General binary decoding error.
  | ParseAddressRawCryptoError CryptoParseError
  -- ^ Crypto error in parsing key hash.
  deriving stock (Eq, Show, Generic)

instance NFData ParseAddressRawError

instance Buildable ParseAddressRawError where
  build =
    \case
      ParseAddressRawInvalidPrefix prefix ->
        "Invalid prefix for raw address" <+> dquotes (hexF prefix) <+> "provided"
      ParseAddressRawUnsupportedPrefix name prefix ->
        "Unsupported raw address prefix type" <+> build name
          <+> dquotes (hexF prefix) <+> "found"
      ParseAddressRawWrongSize addr -> "Given raw address+" <+>
        hexF addr <+> "has invalid length" <+> build (length addr)
      ParseAddressRawMalformedSeparator addr -> "Given raw address" <+> dquotes (hexF addr) <+>
        "does not end with" <+> dquotes (backslash <> "00")
      ParseAddressRawBinaryError err -> "Binary error during decoding address:" <+> build err
      ParseAddressRawCryptoError err -> "Key hash decoding error:" <+> build err

-- | 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) decoders
  where
    decoders = txr1error : map mkDecoder [minBound..]
    txr1error = 0x02 ##: throwError (ParseAddressRawUnsupportedPrefix "txr1" 0x02)
    mkDecoder ak = addressKindTag ak ##: case ak of
      AddressKindImplicit    -> MkAddress . ImplicitAddress <$> keyHash
      AddressKindContract    -> MkAddress . ContractAddress <$> sepHash HashContract
      AddressKindSmartRollup -> MkAddress . SmartRollupAddress <$> sepHash HashSR
    sep = lift Get.getWord8 >>= \case
      0x00 -> pass
      x -> throwError $ ParseAddressRawMalformedSeparator x
    keyHash = mapExceptT (fmap $ first ParseAddressRawCryptoError) decodeKeyHash
    sepHash :: HashTag kind -> ExceptT ParseAddressRawError Get.Get (Hash kind)
    sepHash kind = Hash kind <$> lift (getByteStringCopy hashLengthBytes) <* sep

addressKindTag :: AddressKind -> Word8
addressKindTag = \case
  AddressKindImplicit    -> 0x00
  AddressKindContract    -> 0x01
  -- 0x02 is txr1 which we do not support
  AddressKindSmartRollup -> 0x03

-- | 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
    ( Bottom
    , TypeError ('Text "There is no instance defined for (IsString Address)" ':$$:
                 'Text "Consider using QuasiQuotes: `[ta|some text...|]`"
                )) =>
    IsString (KindedAddress kind) where
  fromString = no


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

instance SingI kind => HasCLReader (KindedAddress kind) where
  getMetavar = toUpper <$> pretty (demote @kind) <> " ADDRESS"
  getReader = getAddressReader parseKindedAddress

instance SingI ks => HasCLReader (ConstrainedAddress ks) where
  getMetavar = intercalate " OR " (fmap toUpper . pretty <$> demote @ks) <> " ADDRESS"
  getReader = getAddressReader parseConstrainedAddress

instance HasCLReader Address where
  getMetavar = "ADDRESS"
  getReader = getAddressReader parseAddress

getAddressReader :: (Text -> Either ParseAddressError addr) -> ReadM addr
getAddressReader parser = eitherReader $
  first (mappend "Failed to parse address: " . pretty) . parser . toText

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

instance SingI kinds => FromJSON (ConstrainedAddress kinds) where
  parseJSON = Aeson.withText "Address" $
    either (fail . pretty) pure . parseConstrainedAddress

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 (Constrained c KindedAddress) where
  toJSON = foldConstrained toJSON
  toEncoding = foldConstrained toEncoding

instance ToJSONKey (Constrained c KindedAddress) where
  toJSONKey = AesonTypes.toJSONKeyText $ foldConstrained 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)