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

-- | Utilities for resolving addresses and aliases.
module Morley.Client.TezosClient.Resolve
  ( ResolveError(..)
  , Resolve(..)
  , resolveAddress
  , resolveAddressMaybe
  , getAlias
  , getAliasMaybe
  , getTezosClientConfig
  , resolveAddressWithAlias
  , resolveAddressWithAliasMaybe
  ) where

import Data.Constraint ((\\))
import Fmt (pretty)

import Morley.Client.TezosClient.Class qualified as Class
import Morley.Client.TezosClient.Config
import Morley.Client.TezosClient.Types.Errors
import Morley.Client.Types
import Morley.Client.Types.AliasesAndAddresses
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Address.Kinds
import Morley.Util.Constrained

class Resolve addressOrAlias where
  type ResolvedAddress addressOrAlias :: Type
  type ResolvedAlias addressOrAlias :: Type
  type ResolvedAddressAndAlias addressOrAlias :: Type

  -- | Looks up the address associated with the given @addressOrAlias@.
  --
  -- When the alias is associated with __both__ an implicit and a contract address:
  --
  -- * The 'SomeAddressOrAlias' instance will return 'REAmbiguousAlias',
  --   unless the alias is prefixed with @implicit:@ or @contract:@ to disambiguate.
  -- * The 'AddressOrAlias' instance will return the address with the requested kind.
  resolveAddressEither
    :: (Class.HasTezosClient m)
    => addressOrAlias
    -> m (Either ResolveError (ResolvedAddress addressOrAlias))

  {- | Looks up the alias associated with the given @addressOrAlias@.

  When the alias is associated with __both__ an implicit and a contract address:

  * The 'SomeAddressOrAlias' instance will return 'REAmbiguousAlias',
    unless the alias is prefixed with @implicit:@ or @contract:@ to disambiguate.
  * The 'AddressOrAlias' instance will return the alias of the address with the requested kind.

  The primary (and probably only) reason this function exists is that
  @octez-client sign@ command only works with aliases. It was
  reported upstream: <https://gitlab.com/tezos/tezos/-/issues/836>.
  -}
  getAliasEither
    :: (Class.HasTezosClient m)
    => addressOrAlias
    -> m (Either ResolveError (ResolvedAlias addressOrAlias))

  -- | Resolve both address and alias at the same time
  resolveAddressWithAliasEither
    :: (Class.HasTezosClient m)
    => addressOrAlias
    -> m (Either ResolveError (ResolvedAddressAndAlias addressOrAlias))

instance L1AddressKind kind => Resolve (KindedAddress kind) where
  type ResolvedAddress (KindedAddress kind) = KindedAddress kind
  type ResolvedAlias (KindedAddress kind) = Alias kind
  type ResolvedAddressAndAlias (KindedAddress kind) = AddressWithAlias kind

  resolveAddressEither
    :: (Class.HasTezosClient m)
    => KindedAddress kind -> m (Either ResolveError (KindedAddress kind))
  resolveAddressEither = pure . Right

  getAliasEither
    :: (Class.HasTezosClient m)
    => KindedAddress kind -> m (Either ResolveError (Alias kind))
  getAliasEither addr
    = maybeToRight (REAddressNotFound addr)
    . lookupAlias addr <$> Class.getAliasesAndAddresses

  resolveAddressWithAliasEither addr = fmap (AddressWithAlias addr) <$> getAliasEither addr

instance Resolve (Alias kind) where
  type ResolvedAddress (Alias kind) = KindedAddress kind
  type ResolvedAlias (Alias kind) = Alias kind
  type ResolvedAddressAndAlias (Alias kind) = AddressWithAlias kind

  resolveAddressEither
    :: (Class.HasTezosClient m)
    => Alias kind -> m (Either ResolveError (KindedAddress kind))
  resolveAddressEither alias = do
    aas <- Class.getAliasesAndAddresses
    pure $ lookupAddr alias aas
      & maybeToRight (handleMissing aas)
    where
      handleMissing :: AliasesAndAddresses -> ResolveError
      handleMissing aas
        = maybe (REAliasNotFound $ pretty $ AddressAlias alias) (REWrongKind alias)
        $ case alias of
            -- notice kind is flipped
            ImplicitAlias aliasTxt -> Constrained <$> lookupAddr (ContractAlias aliasTxt) aas
            ContractAlias aliasTxt -> Constrained <$> lookupAddr (ImplicitAlias aliasTxt) aas

  getAliasEither
    :: (Class.HasTezosClient m)
    => Alias kind -> m (Either ResolveError (Alias kind))
  getAliasEither alias = ($> alias) <$> resolveAddressEither alias -- check if alias exists

  resolveAddressWithAliasEither alias =
    fmap (`AddressWithAlias` alias) <$> resolveAddressEither alias

instance Resolve (AddressOrAlias kind) where
  type ResolvedAddress (AddressOrAlias kind) = KindedAddress kind
  type ResolvedAlias (AddressOrAlias kind) = Alias kind
  type ResolvedAddressAndAlias (AddressOrAlias kind) = AddressWithAlias kind

  resolveAddressEither
    :: (Class.HasTezosClient m)
    => AddressOrAlias kind -> m (Either ResolveError (KindedAddress kind))
  resolveAddressEither = \case
    AddressResolved addr -> resolveAddressEither addr
    AddressAlias alias -> resolveAddressEither alias

  getAliasEither
    :: (Class.HasTezosClient m)
    => AddressOrAlias kind -> m (Either ResolveError (Alias kind))
  getAliasEither = \case
    AddressAlias alias -> getAliasEither alias
    AddressResolved addr -> getAliasEither addr

  resolveAddressWithAliasEither = \case
    AddressAlias alias -> resolveAddressWithAliasEither alias
    AddressResolved addr -> resolveAddressWithAliasEither addr

instance Resolve SomeAddressOrAlias where
  type ResolvedAddress SomeAddressOrAlias = L1Address
  type ResolvedAlias SomeAddressOrAlias = SomeAlias
  type ResolvedAddressAndAlias SomeAddressOrAlias = Constrained L1AddressKind AddressWithAlias

  resolveAddressEither
    :: (Class.HasTezosClient m)
    => SomeAddressOrAlias -> m (Either ResolveError L1Address)
  resolveAddressEither = \case
    SAOAKindUnspecified aliasText -> do
      aas <- Class.getAliasesAndAddresses
      let addrs = traverseConstrained (`lookupAddr` aas) `mapMaybe`
            [ Constrained $ mkAlias @'AddressKindContract aliasText
            , Constrained $ mkAlias @'AddressKindImplicit aliasText
            ]
      pure $ case addrs of
        [] -> Left $ REAliasNotFound aliasText
        [addr] -> Right addr
        as -> Left $ REAmbiguousAlias aliasText as
    SAOAKindSpecified aoa ->
      fmap Constrained <$> resolveAddressEither aoa \\ addressOrAliasKindSanity aoa

  getAliasEither
    :: (Class.HasTezosClient m)
    => SomeAddressOrAlias -> m (Either ResolveError SomeAlias)
  getAliasEither = \case
    SAOAKindSpecified aoa -> do
      fmap SomeAlias
        <$> getAliasEither aoa \\ addressOrAliasKindSanity aoa
    aoa@SAOAKindUnspecified{} -> runExceptT do
      -- Find out whether this alias is associated with an implicit address or a contract,
      -- and return an @Alias kind@ of the correct kind.
      ExceptT (resolveAddressWithAliasEither aoa) <&> foldConstrained (SomeAlias . awaAlias)

  resolveAddressWithAliasEither addr = runExceptT case addr of
    SAOAKindSpecified aoa -> do
      kaddr <- ExceptT $ resolveAddressEither aoa
      kalias <- ExceptT $ getAliasEither aoa
      pure $ Constrained (AddressWithAlias kaddr kalias) \\ addressOrAliasKindSanity aoa
    aoa@(SAOAKindUnspecified aliasText) -> do
      ExceptT (resolveAddressEither aoa) <&> foldConstrained
        \kaddr -> Constrained $ AddressWithAlias kaddr $ mkAlias aliasText \\ addressKindSanity kaddr

-- | Looks up the address and alias with the given @addressOrAlias@.
resolveAddressWithAlias
  :: forall addressOrAlias m
   . (Class.HasTezosClient m, MonadThrow m,Resolve addressOrAlias)
  => addressOrAlias
  -> m (ResolvedAddressAndAlias addressOrAlias)
resolveAddressWithAlias = resolveAddressWithAliasEither >=> either (throwM . ResolveError) pure

-- | Looks up the address and alias with the given @addressOrAlias@.
resolveAddressWithAliasMaybe
  :: forall addressOrAlias m
   . (Class.HasTezosClient m, Resolve addressOrAlias)
  => addressOrAlias
  -> m (Maybe (ResolvedAddressAndAlias addressOrAlias))
resolveAddressWithAliasMaybe = fmap rightToMaybe . resolveAddressWithAliasEither

-- | Looks up the address associated with the given @addressOrAlias@.
--
-- Will throw a 'TezosClientError' if @addressOrAlias@ is an alias and:
--
-- * the alias does not exist.
-- * the alias exists but its address is of the wrong kind.
--
-- When the alias is associated with __both__ an implicit and a contract address:
--
-- * The 'SomeAddressOrAlias' instance will throw a 'TezosClientError',
--   unless the alias is prefixed with @implicit:@ or @contract:@ to disambiguate.
-- * The 'AddressOrAlias' instance will return the address with the requested kind.
resolveAddress
  :: forall addressOrAlias m
   . (Class.HasTezosClient m, MonadThrow m, Resolve addressOrAlias)
  => addressOrAlias
  -> m (ResolvedAddress addressOrAlias)
resolveAddress = resolveAddressEither >=> either (throwM . ResolveError) pure

-- | Looks up the address associated with the given @addressOrAlias@.
--
-- Will return 'Nothing' if @addressOrAlias@ is an alias and:
--
-- * the alias does not exist.
-- * the alias exists but its address is of the wrong kind.
--
-- When the alias is associated with __both__ an implicit and a contract address:
--
-- * The 'SomeAddressOrAlias' instance will throw a 'TezosClientError',
--   unless the alias is prefixed with @implicit:@ or @contract:@ to disambiguate.
-- * The 'AddressOrAlias' instance will return the address with the requested kind.
resolveAddressMaybe
  :: forall addressOrAlias m
   . (Class.HasTezosClient m, MonadThrow m, Resolve addressOrAlias)
  => addressOrAlias
  -> m (Maybe (ResolvedAddress addressOrAlias))
resolveAddressMaybe aoa = resolveAddressEither aoa >>= either handleResolveError (pure . Just)

handleResolveError :: MonadThrow m => ResolveError -> m (Maybe a)
handleResolveError = \case
  e@REAmbiguousAlias{} -> throwM $ ResolveError e
  REAliasNotFound{} -> pure Nothing
  REAddressNotFound{} -> pure Nothing
  REWrongKind{} -> pure Nothing

{- | Looks up the alias associated with the given @addressOrAlias@.

Will throw a 'TezosClientError' if @addressOrAlias@:

  * is an address that is not associated with any alias.
  * is an alias that does not exist.
  * is an alias that exists but its address is of the wrong kind.

When the alias is associated with __both__ an implicit and a contract address:

  * The 'SomeAddressOrAlias' instance will throw a 'TezosClientError',
    unless the alias is prefixed with @implicit:@ or @contract:@ to disambiguate.
  * The 'AddressOrAlias' instance will return the alias.
-}
getAlias
  :: forall addressOrAlias m
   . (Class.HasTezosClient m, MonadThrow m, Resolve addressOrAlias)
  => addressOrAlias
  -> m (ResolvedAlias addressOrAlias)
getAlias = getAliasEither >=> either (throwM . ResolveError) pure

{- | Looks up the alias associated with the given @addressOrAlias@.

Will return 'Nothing' if @addressOrAlias@:

  * is an address that is not associated with any alias.
  * is an alias that does not exist.
  * is an alias that exists but its address is of the wrong kind.

When the alias is associated with __both__ an implicit and a contract address:

  * The 'SomeAddressOrAlias' instance will throw a 'TezosClientError',
    unless the alias is prefixed with @implicit:@ or @contract:@ to disambiguate.
  * The 'AddressOrAlias' instance will return the alias.
-}
getAliasMaybe
  :: forall addressOrAlias m
   . (Class.HasTezosClient m, MonadThrow m, Resolve addressOrAlias)
  => addressOrAlias
  -> m (Maybe (ResolvedAlias addressOrAlias))
getAliasMaybe aoa = getAliasEither aoa >>= either handleResolveError (pure . Just)