-- 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
  , AnyParsableAddress(..)
  , Constrained(.., MkConstrainedAddress, MkAddress)

  , GlobalCounter(..)
  , mkContractHashHack
  , parseConstrainedAddress
  , parseAnyAddress

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

  -- * Utilities
  , addressKindSanity
  , usingImplicitOrContractKind
  , unImplicitAddress
  ) 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 (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, 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.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
  -- | @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)

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'TxRollupAddress'
type TxRollupAddress = KindedAddress 'AddressKindTxRollup

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

pattern MkConstrainedAddress
  :: forall ks. ()
  => forall kind. ConstrainAddressKind ks kind
  => KindedAddress kind -> ConstrainedAddress ks
pattern MkConstrainedAddress a = Constrained a
{-# COMPLETE MkConstrainedAddress #-}
{-# DEPRECATED MkConstrainedAddress "Use Constrained instead" #-}

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

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

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

-- | A sum type for any known parsable address (@tz1@ - @tz4@, @KT1@, @txr1@)
data AnyParsableAddress
  = AnyParsableL2Address TxRollupL2Address
  | AnyParsableAddress Address
  deriving stock (Show, Eq, Generic)

instance NFData AnyParsableAddress

-- | 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
    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
  = ParseAddressCryptoError CryptoParseError
  -- ^ The address parsers failed with some error.
  | ParseAddressWrongKind [AddressKind] AnyParsableAddress
  -- ^ The parsed address is of wrong kind
  deriving stock (Show, Eq, Generic)

instance NFData ParseAddressError

instance Buildable ParseAddressError where
  build = buildRenderDoc

instance RenderDoc ParseAddressError where
  renderDoc context =
    \case
      ParseAddressCryptoError pkErr -> "Address failed to parse: " <> renderDoc context pkErr
      ParseAddressWrongKind expected (AnyParsableAddress (Constrained a)) -> mconcat
        [ "Expected address of kind ", renderAddressKinds expected
        , ", but got ", renderAnyBuildable a
        ]
      ParseAddressWrongKind expected (AnyParsableL2Address a) -> mconcat
        [ "Expected address of kind ", renderAddressKinds expected
        , ", but got unexpected TxRollupL2Address ", renderAnyBuildable a
        ]
    where
      renderAddressKinds as = mconcat $ intersperse ", " (renderAnyBuildable <$> 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] . AnyParsableAddress $ 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 . AnyParsableAddress
  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
  :: ConstrainAddressKind ks k
  => Sing k -> Sing x -> Dict (ConstrainAddressKind (x ': ks) k)
proofAddressCast k x = case x of
  SAddressKindImplicit -> case k of
    SAddressKindImplicit -> Dict
    SAddressKindContract -> Dict
    SAddressKindTxRollup -> Dict
  SAddressKindContract -> case k of
    SAddressKindImplicit -> Dict
    SAddressKindContract -> Dict
    SAddressKindTxRollup -> Dict
  SAddressKindTxRollup -> case k of
    SAddressKindImplicit -> Dict
    SAddressKindContract -> Dict
    SAddressKindTxRollup -> 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 x = parseAnyAddress x >>= \case
  a@AnyParsableL2Address{} -> Left $ ParseAddressWrongKind [minBound..] a
  AnyParsableAddress a -> Right a

-- | Parse any kind of known address from it's human-readable textual
-- representation, or fail if it's invalid.
parseAnyAddress :: Text -> Either ParseAddressError AnyParsableAddress
parseAnyAddress a = first ParseAddressCryptoError $ parseSomeHashBase58 a <&> \case
  Some h@(Hash HashBLS _) -> AnyParsableL2Address $ TxRollupL2Address h
  Some h@(Hash hk _) -> AnyParsableAddress $ case hk of
    HashKey KeyTypeEd25519   -> Constrained $ ImplicitAddress h
    HashKey KeyTypeSecp256k1 -> Constrained $ ImplicitAddress h
    HashKey KeyTypeP256      -> Constrained $ ImplicitAddress h
    HashContract  -> Constrained $ ContractAddress h
    HashTXR       -> Constrained $ TxRollupAddress 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.
  | 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 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
      ParseAddressRawCryptoError 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 ParseAddressRawCryptoError) 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
    ( 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
  getReader = eitherReader parseAddrDo
    where
      parseAddrDo addr =
        first (mappend "Failed to parse address: " . pretty) $
        parseKindedAddress $ toText addr
  getMetavar = "ADDRESS"

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