-- | By default we represent error tags using strings. This module
-- makes it possible to use numbers instead.
--
-- 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
  ( ErrorTagMap
  , gatherErrorTags
  , addNewErrorTags
  , buildErrorTagMap
  , applyErrorTagMap
  , useNumericErrors

  , errorFromValNumeric
  ) where

import Data.Bimap (Bimap)
import qualified Data.Bimap as Bimap
import Data.Default (def)
import qualified Data.HashSet as HS
import Data.Singletons (SingI, sing)
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

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

-- | 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 = iMapAnyCode (applyErrorTagMapT errorTagMap)

-- | 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.
applyErrorTagMapT ::
     HasCallStack
  => ErrorTagMap
  -> Instr inp out
  -> Instr inp out
applyErrorTagMapT errorTagMap instr =
  dfsModifyInstr dfsSettings step instr
  where
    dfsSettings :: DfsSettings ()
    dfsSettings = def
      { dsGoToValues = True
      }

    tagToNatValue :: HasCallStack => MText -> Value ('Tc 'CNat)
    tagToNatValue tag =
      case Bimap.lookupR tag errorTagMap of
        -- 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.
        Nothing -> error $ "Can't find a tag: " <> pretty tag
        Just n -> VC $ CvNat 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 ::
  (Typeable t, SingI t, IsError e) => ErrorTagMap -> Value t -> Either Text e
errorFromValNumeric errorTagMap v =
  case v of
    VC (CvNat tag)
      | Just textualTag <- Bimap.lookup tag errorTagMap ->
        errorFromVal . VC . CvString $ textualTag
    VPair (VC (CvNat tag), something)
      | Just textualTag <- Bimap.lookup tag errorTagMap
      , _ :: Value pair <- v ->
        case sing @pair of
          STPair {} ->
            errorFromVal $ VPair (VC $ CvString textualTag, something)
    _ -> errorFromVal v