module Plutarch.Pretty.Internal.Types (
  PrettyCursor (..),
  PrettyState (..),
  PrettyMonad,
  forkState,
  normalizeCursor,
  specializeCursor,
  memorizeName,
  insertName,
  insertBindings,
  builtinFunAtRef,
  nameOfRef,
) where

import Control.Monad ((<=<))
import Control.Monad.Reader (ReaderT)
import Control.Monad.ST (ST)
import Control.Monad.State (MonadState (get, put), StateT)
import Data.List (find)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as Txt
import System.Random.Stateful (STGenM, StdGen)

import PlutusCore qualified as PLC
import UntypedPlutusCore (DefaultFun, Index)

import Plutarch.Pretty.Internal.Config (forcedPrefix)

{- | Notifies the prettifier what "state" the cursor currently is, so it can decide
whether or not to wrap the target expression in parens.

Normal indicates no parens wrapping is necessary, even for complex expressions.

Special indicates complex expressions should be wrapped in parens.

Usually, "Special" just hints at one of three states:

1. Applying - The expression is being applied like a function.
2. Applied - The expression is being applied as a function argument.
3. Unary arg - The expression is being used as an argument to a high arity unary operator (~ and !).
-}
data PrettyCursor = Normal | Special
  deriving stock (PrettyCursor
PrettyCursor -> PrettyCursor -> Bounded PrettyCursor
forall a. a -> a -> Bounded a
$cminBound :: PrettyCursor
minBound :: PrettyCursor
$cmaxBound :: PrettyCursor
maxBound :: PrettyCursor
Bounded, Int -> PrettyCursor
PrettyCursor -> Int
PrettyCursor -> [PrettyCursor]
PrettyCursor -> PrettyCursor
PrettyCursor -> PrettyCursor -> [PrettyCursor]
PrettyCursor -> PrettyCursor -> PrettyCursor -> [PrettyCursor]
(PrettyCursor -> PrettyCursor)
-> (PrettyCursor -> PrettyCursor)
-> (Int -> PrettyCursor)
-> (PrettyCursor -> Int)
-> (PrettyCursor -> [PrettyCursor])
-> (PrettyCursor -> PrettyCursor -> [PrettyCursor])
-> (PrettyCursor -> PrettyCursor -> [PrettyCursor])
-> (PrettyCursor -> PrettyCursor -> PrettyCursor -> [PrettyCursor])
-> Enum PrettyCursor
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: PrettyCursor -> PrettyCursor
succ :: PrettyCursor -> PrettyCursor
$cpred :: PrettyCursor -> PrettyCursor
pred :: PrettyCursor -> PrettyCursor
$ctoEnum :: Int -> PrettyCursor
toEnum :: Int -> PrettyCursor
$cfromEnum :: PrettyCursor -> Int
fromEnum :: PrettyCursor -> Int
$cenumFrom :: PrettyCursor -> [PrettyCursor]
enumFrom :: PrettyCursor -> [PrettyCursor]
$cenumFromThen :: PrettyCursor -> PrettyCursor -> [PrettyCursor]
enumFromThen :: PrettyCursor -> PrettyCursor -> [PrettyCursor]
$cenumFromTo :: PrettyCursor -> PrettyCursor -> [PrettyCursor]
enumFromTo :: PrettyCursor -> PrettyCursor -> [PrettyCursor]
$cenumFromThenTo :: PrettyCursor -> PrettyCursor -> PrettyCursor -> [PrettyCursor]
enumFromThenTo :: PrettyCursor -> PrettyCursor -> PrettyCursor -> [PrettyCursor]
Enum, PrettyCursor -> PrettyCursor -> Bool
(PrettyCursor -> PrettyCursor -> Bool)
-> (PrettyCursor -> PrettyCursor -> Bool) -> Eq PrettyCursor
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PrettyCursor -> PrettyCursor -> Bool
== :: PrettyCursor -> PrettyCursor -> Bool
$c/= :: PrettyCursor -> PrettyCursor -> Bool
/= :: PrettyCursor -> PrettyCursor -> Bool
Eq, Int -> PrettyCursor -> ShowS
[PrettyCursor] -> ShowS
PrettyCursor -> [Char]
(Int -> PrettyCursor -> ShowS)
-> (PrettyCursor -> [Char])
-> ([PrettyCursor] -> ShowS)
-> Show PrettyCursor
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PrettyCursor -> ShowS
showsPrec :: Int -> PrettyCursor -> ShowS
$cshow :: PrettyCursor -> [Char]
show :: PrettyCursor -> [Char]
$cshowList :: [PrettyCursor] -> ShowS
showList :: [PrettyCursor] -> ShowS
Show)

data PrettyState = PrettyState
  { PrettyState -> Map Index Text
ps'nameMap :: Map Index Text
  , PrettyState -> Set Text
ps'names :: Set Text
  , PrettyState -> PrettyCursor
ps'cursor :: PrettyCursor
  }

type PrettyMonad s = ReaderT (STGenM StdGen s) (StateT PrettyState (ST s))

forkState :: MonadState s m => m b -> m b
forkState :: forall s (m :: Type -> Type) b. MonadState s m => m b -> m b
forkState m b
x = m s
forall s (m :: Type -> Type). MonadState s m => m s
get m s -> (s -> m b) -> m b
forall a b. m a -> (a -> m b) -> m b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\s
s -> m b
x m b -> m () -> m b
forall a b. m a -> m b -> m a
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f a
<* s -> m ()
forall s (m :: Type -> Type). MonadState s m => s -> m ()
put s
s)

normalizeCursor :: PrettyState -> PrettyState
normalizeCursor :: PrettyState -> PrettyState
normalizeCursor PrettyState
x = PrettyState
x {ps'cursor = Normal}

specializeCursor :: PrettyState -> PrettyState
specializeCursor :: PrettyState -> PrettyState
specializeCursor PrettyState
x = PrettyState
x {ps'cursor = Special}

memorizeName :: Text -> PrettyState -> PrettyState
memorizeName :: Text -> PrettyState -> PrettyState
memorizeName Text
n x :: PrettyState
x@PrettyState {Set Text
$sel:ps'names:PrettyState :: PrettyState -> Set Text
ps'names :: Set Text
ps'names} = PrettyState
x {ps'names = Set.insert n ps'names}

-- | Insert a fresh binding onto the name map, i.e a name at index 0 - incrementing all other indices.
insertName :: Text -> PrettyState -> PrettyState
insertName :: Text -> PrettyState -> PrettyState
insertName Text
name x :: PrettyState
x@PrettyState {Map Index Text
$sel:ps'nameMap:PrettyState :: PrettyState -> Map Index Text
ps'nameMap :: Map Index Text
ps'nameMap} =
  PrettyState
x
    { ps'nameMap = Map.mapKeys (+ 1) ps'nameMap <> Map.singleton 0 name
    }

insertBindings :: [Text] -> PrettyState -> PrettyState
insertBindings :: [Text] -> PrettyState -> PrettyState
insertBindings [Text]
names prst :: PrettyState
prst@PrettyState {Map Index Text
$sel:ps'nameMap:PrettyState :: PrettyState -> Map Index Text
ps'nameMap :: Map Index Text
ps'nameMap} =
  PrettyState
prst
    { ps'nameMap =
        Map.mapKeys (+ nameCount) ps'nameMap
          <> foldMap (uncurry Map.singleton) (zip [0 .. (nameCount - 1)] names)
    }
  where
    nameCount :: Index
nameCount = Int -> Index
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Index) -> Int -> Index
forall a b. (a -> b) -> a -> b
$ [Text] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Text]
names

builtinFunAtRef :: Map Index Text -> Index -> Maybe DefaultFun
builtinFunAtRef :: Map Index Text -> Index -> Maybe DefaultFun
builtinFunAtRef Map Index Text
nameMap = Text -> Maybe DefaultFun
builtinFunFromName (Text -> Maybe DefaultFun)
-> (Index -> Maybe Text) -> Index -> Maybe DefaultFun
forall (m :: Type -> Type) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< (Index -> Map Index Text -> Maybe Text)
-> Map Index Text -> Index -> Maybe Text
forall a b c. (a -> b -> c) -> b -> a -> c
flip Index -> Map Index Text -> Maybe Text
nameOfRef Map Index Text
nameMap

nameOfRef :: Index -> Map Index Text -> Maybe Text
nameOfRef :: Index -> Map Index Text -> Maybe Text
nameOfRef Index
ix = Index -> Map Index Text -> Maybe Text
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Index
ix Index -> Index -> Index
forall a. Num a => a -> a -> a
- Index
1)

builtinFunFromName :: Text -> Maybe DefaultFun
builtinFunFromName :: Text -> Maybe DefaultFun
builtinFunFromName Text
res =
  if Int -> Text -> Text
Txt.take Int
prefixLen Text
res Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
forcedPrefix
    then Text -> Maybe DefaultFun
helper (Text -> Maybe DefaultFun) -> Text -> Maybe DefaultFun
forall a b. (a -> b) -> a -> b
$ Int -> Text -> Text
Txt.drop Int
prefixLen Text
res
    else Text -> Maybe DefaultFun
helper Text
res
  where
    prefixLen :: Int
prefixLen = Text -> Int
Txt.length Text
forcedPrefix
    helper :: Text -> Maybe DefaultFun
helper Text
s = (DefaultFun -> Bool) -> [DefaultFun] -> Maybe DefaultFun
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Maybe a
find (\DefaultFun
e -> DefaultFun -> Text
forall a. Show a => a -> Text
showText DefaultFun
e Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
s) [DefaultFun]
builtinFunNames
    builtinFunNames :: [DefaultFun]
builtinFunNames = [Item [DefaultFun]
forall a. Bounded a => a
minBound .. Item [DefaultFun]
forall a. Bounded a => a
maxBound] :: [PLC.DefaultFun]

showText :: Show a => a -> Text
showText :: forall a. Show a => a -> Text
showText = [Char] -> Text
Txt.pack ([Char] -> Text) -> (a -> [Char]) -> a -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [Char]
forall a. Show a => a -> [Char]
show