-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | By default we represent error tags using strings. This module
-- makes it possible to use numbers instead. It introduces new [error format].
--
-- There are two possible ways to use it:
-- 1. If you have just one Lorentz instruction (potentially a big one),
-- just use 'useNumericErrors' function. It will change error representation
-- there and return a map that can be used to interpret new error codes.
-- 2. If your contract consists of multiple parts, start with gathering all
-- error tags ('gatherErrorTags'). Then build 'ErrorTagMap' using
-- 'addNewErrorTags'. Pass empty map if you are building from scratch
-- (you can use 'buildErrorTagMap' shortcut) or an existing
-- map if you have one (e. g. you are upgrading a contract).

module Lorentz.Errors.Numeric.Contract
  ( ErrorTagMap
  , ErrorTagExclusions
  , gatherErrorTags
  , addNewErrorTags
  , buildErrorTagMap
  , excludeErrorTags
  , applyErrorTagMap
  , applyErrorTagMapWithExclusions
  , useNumericErrors

  , errorFromValNumeric
  , errorToValNumeric
  ) where

import Data.Bimap (Bimap)
import qualified Data.Bimap as Bimap
import Data.Default (def)
import qualified Data.HashSet as HS
import Fmt (pretty)

import Lorentz.Base
import Lorentz.Errors
import Michelson.Analyzer
import Michelson.FailPattern
import Michelson.Text (MText)
import Michelson.Typed

-- | This is a bidirectional map with correspondence between numeric
-- and textual error tags.
type ErrorTagMap = Bimap Natural MText

-- | Tags excluded from map.
type ErrorTagExclusions = HashSet MText

-- | Find all textual error tags that are used in typical
-- @FAILWITH@ patterns within given instruction.
-- Map them to natural numbers.
gatherErrorTags :: inp :-> out -> HashSet MText
gatherErrorTags = HS.fromMap . void . arErrorTags . analyze . iAnyCode

-- | Add more error tags to an existing 'ErrorTagMap'. It is useful when
-- your contract consists of multiple parts (e. g. in case of contract
-- upgrade), you have existing map for some part and want to add tags
-- from another part to it.
-- You can pass empty map as existing one if you just want to build
-- 'ErrorTagMap' from a set of textual tags. See 'buildErrorTagMap'.
addNewErrorTags :: ErrorTagMap -> HashSet MText -> ErrorTagMap
addNewErrorTags existingMap newTags =
  foldl' (flip $ uncurry Bimap.tryInsert) existingMap newItems
  where
    firstUnusedNumeric
      | Bimap.null existingMap = 0
      | otherwise = fst (Bimap.findMax existingMap) + 1

    newItems :: [(Natural, MText)]
    newItems = zip [firstUnusedNumeric .. ] (toList newTags)

-- | Build 'ErrorTagMap' from a set of textual tags.
buildErrorTagMap :: HashSet MText -> ErrorTagMap
buildErrorTagMap = addNewErrorTags Bimap.empty

-- | Remove some error tags from map.
-- This way you say to remain these string tags intact, while others will be
-- converted to numbers when this map is applied.
--
-- Note that later you have to apply this map using
-- 'applyErrorTagMapWithExclusions', otherwise an error would be raised.
excludeErrorTags
  :: HasCallStack
  => ErrorTagExclusions -> ErrorTagMap -> ErrorTagMap
excludeErrorTags toExclude errMap =
  foldl' (flip deleteExistingR) errMap toExclude
  where
    deleteExistingR k m = case Bimap.lookupR k m of
      Just _ -> Bimap.deleteR k m
      Nothing ->
        error $ "Tag " <> show k <> " does not appear in the contract"

-- | For each typical 'FAILWITH' that uses a string to represent error
-- tag this function changes error tag to be a number using the
-- supplied conversion map.
-- It assumes that supplied map contains all such strings
-- (and will error out if it does not).
-- It will always be the case if you gather all error tags using
-- 'gatherErrorTags' and build 'ErrorTagMap' from them using 'addNewErrorTags'.
applyErrorTagMap :: HasCallStack => ErrorTagMap -> inp :-> out -> inp :-> out
applyErrorTagMap errorTagMap = applyErrorTagMapWithExclusions errorTagMap mempty

-- | Similar to 'applyErrorTagMap', but for case when you have excluded some
-- tags from map via 'excludeErrorTags'.
-- Needed, because both 'excludeErrorTags' and this function do not tolerate
-- unknown errors in contract code (for your safety).
applyErrorTagMapWithExclusions
  :: HasCallStack
  => ErrorTagMap -> ErrorTagExclusions -> inp :-> out -> inp :-> out
applyErrorTagMapWithExclusions errorTagMap exclusions =
  iMapAnyCode (applyErrorTagMapWithExcT errorTagMap exclusions)

-- | This function implements the simplest scenario of using this
-- module's functionality:
-- 1. Gather all error tags from a single instruction.
-- 2. Turn them into error conversion map.
-- 3. Apply this conversion.
useNumericErrors ::
  HasCallStack => inp :-> out -> (inp :-> out, ErrorTagMap)
useNumericErrors instr = (applyErrorTagMap errorTagMap instr, errorTagMap)
  where
    errorTagMap = buildErrorTagMap $ gatherErrorTags instr

-- This function works with 'Michelson.Typed' representation, not with Lorentz.
applyErrorTagMapWithExcT ::
     HasCallStack
  => ErrorTagMap
  -> ErrorTagExclusions
  -> Instr inp out
  -> Instr inp out
applyErrorTagMapWithExcT errorTagMap exclusions instr =
  dfsModifyInstr dfsSettings step instr
  where
    dfsSettings :: DfsSettings ()
    dfsSettings = def
      { dsGoToValues = True
      }

    tagToNatValue :: HasCallStack => MText -> SomeConstrainedValue ConstantScope'
    tagToNatValue tag =
      case (HS.member tag exclusions, Bimap.lookupR tag errorTagMap) of
        (True, _) -> SomeConstrainedValue (VString tag)
        -- It will be applied to textual tags detected by 'modifyTypicalFailWith'.
        -- Here we assume that all of them are discovered by the analyzer.
        -- If this error ever happens, it means that someone used
        -- 'applyErrorTagMap' with incomplete 'ErrorTagMap' or there is an
        -- internal bug somewhere.
        (False, Nothing) -> error $ "Can't find a tag: " <> pretty tag
        (False, Just n) -> SomeConstrainedValue (VNat n)

    step :: HasCallStack => Instr inp out -> Instr inp out
    step = modifyTypicalFailWith tagToNatValue

-- | If you apply numeric error representation in your contract, 'errorFromVal'
-- will stop working because it doesn't know about this
-- transformation.
-- This function takes this transformation into account.
-- If a number is used as a tag, but it is not found in the passed
-- map, we conservatively preserve that number (because this whole
-- approach is rather a heuristic).
errorFromValNumeric ::
  (KnownT t, IsError e) => ErrorTagMap -> Value t -> Either Text e
errorFromValNumeric errorTagMap v =
  case v of
    VNat tag
      | Just textualTag <- Bimap.lookup tag errorTagMap ->
        errorFromVal $ VString textualTag
    (VPair (VNat tag, something) :: Value pair)
      | Just textualTag <- Bimap.lookup tag errorTagMap ->
        case sing @pair of
          STPair {} ->
            errorFromVal $ VPair (VString textualTag, something)
    _ -> errorFromVal v

-- | If you apply numeric error representation in your contract, 'errorToVal'
-- will stop working because it doesn't know about this
-- transformation.
-- This function takes this transformation into account.
-- If a string is used as a tag, but it is not found in the passed
-- map, we conservatively preserve that string (because this whole
-- approach is rather a heuristic).
errorToValNumeric ::
  IsError e => ErrorTagMap -> e -> (forall t. ErrorScope t => Value t -> r) -> r
errorToValNumeric errorTagMap e cont =
  errorToVal e $ \case
    VString textualTag
      | Just tag  <- Bimap.lookupR textualTag errorTagMap -> cont (VNat tag)

    (VPair (VString textualTag, something) :: Value pair)
      | Just tag <- Bimap.lookupR textualTag errorTagMap ->
          case sing @pair of
            STPair {} ->
              cont (VPair (VNat tag, something))
    v -> cont v