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

{-# OPTIONS_HADDOCK not-home #-}

module Morley.Michelson.Typed.Scope.Internal.WellTyped
  ( module Morley.Michelson.Typed.Scope.Internal.WellTyped
  ) where

import Data.Constraint (Dict(..))
import Data.Singletons (Sing, SingI(..), fromSing, withSingI)
import Data.Type.Equality ((:~:)(..))
import Fmt (Buildable(..), quoteOrIndentF, (++|), (|++^))

import Morley.Michelson.Typed.Scope.Internal.Comparable
import Morley.Michelson.Typed.Scope.Internal.ForbidT
import Morley.Michelson.Typed.Scope.Internal.Presence
import Morley.Michelson.Typed.Sing (SingT(..))
import Morley.Michelson.Typed.T (T(..))

{-$setup
>>> import Morley.Michelson.Typed
>>> import Fmt (pretty)
-}

data BadTypeForScope
  = BtNotComparable
  | BtIsOperation
  | BtHasBigMap
  | BtHasNestedBigMap
  | BtHasContract
  | BtHasTicket
  | BtHasSaplingState
  deriving stock (Show, Eq, Generic)
  deriving anyclass (NFData)

instance Buildable BadTypeForScope where
  build = \case
    BtNotComparable -> "is not comparable"
    BtIsOperation -> "has 'operation' type"
    BtHasBigMap -> "has 'big_map'"
    BtHasNestedBigMap -> "has nested 'big_map'"
    BtHasContract -> "has 'contract' type"
    BtHasTicket -> "has 'ticket' type"
    BtHasSaplingState -> "has 'sapling_state' type"

-- | This class encodes Michelson rules w.r.t where it requires comparable
-- types. Earlier we had a dedicated type for representing comparable types @CT@.
-- But then we integreated those types into @T@. This meant that some of the
-- types that could be formed with various combinations of @T@ would be
-- illegal as per Michelson typing rule. Using this class, we inductively
-- enforce that a type and all types it contains are well typed as per
-- Michelson's rules.
type WellTyped :: T -> Constraint
class (SingI t, RecSuperC WellTyped t, WellTypedConstraints t) => WellTyped t
instance (SingI t, RecSuperC WellTyped t, WellTypedConstraints t) => WellTyped t

-- | Additional (non-recursive) constraints on type arguments for specific types
type WellTypedConstraints :: T -> Constraint
type family WellTypedConstraints t where
  WellTypedConstraints ('TSet t)      = Comparable t
  WellTypedConstraints ('TContract t) = (ForbidOp t, ForbidNestedBigMaps t)
  WellTypedConstraints ('TTicket t)   = Comparable t
  WellTypedConstraints ('TMap k _)    = Comparable k
  WellTypedConstraints ('TBigMap k v) = (Comparable k, ForbidBigMap v, ForbidOp v)
  WellTypedConstraints _ = ()

-- | Error type for when a value is not well-typed.
data NotWellTyped = NotWellTyped
  { nwtBadType :: T
  , nwtCause :: BadTypeForScope
  }

instance Buildable NotWellTyped where
  build (NotWellTyped t c) =
    "Given type is not well typed because" ++| quoteOrIndentF t |++^ build c

{- | Given a type, provide evidence that it is well typed w.r.t to the Michelson
rules regarding where comparable types are required.

>>> either pretty print $ getWTP @'TUnit
Dict
>>> either pretty print $ getWTP @('TSet 'TOperation)
Given type is not well typed because 'operation' is not comparable
>>> type Pairs a = 'TPair a a
>>> type Pairs2 = Pairs (Pairs 'TUnit)
>>> either pretty print $ getWTP @('TSet ('TPair Pairs2 'TOperation))
Given type is not well typed because
  pair (pair (pair unit unit) unit unit) operation
is not comparable
-}
getWTP :: forall t. (SingI t) => Either NotWellTyped (Dict (WellTyped t))
getWTP = getWTP' sing

-- | Version of 'getWTP' that accepts 'Sing' at term-level.
getWTP' :: Sing t -> Either NotWellTyped (Dict (WellTyped t))
getWTP' = \case
  STKey -> Right Dict
  STUnit -> Right Dict
  STSignature -> Right Dict
  STChainId -> Right Dict
  STOption s -> do
    Dict <- getWTP' s
    pure Dict
  STList s -> do
    Dict <- getWTP' s
    pure Dict
  STSet s -> do
    Dict <- getWTP' s
    Dict <- eitherWellTyped BtNotComparable s comparabilityPresence
    pure Dict
  STOperation -> Right Dict
  STContract s -> do
    Dict <- getWTP' s
    Refl <- eitherWellTyped BtIsOperation s $ tAbsence SPSOp
    Refl <- eitherWellTyped BtHasNestedBigMap s $ tAbsence SPSNestedBigMaps
    pure Dict
  STTicket s -> do
    Dict <- getWTP' s
    Dict <- eitherWellTyped BtNotComparable s comparabilityPresence
    pure Dict
  STPair s1 s2 -> do
    Dict <- getWTP' s1
    Dict <- getWTP' s2
    pure Dict
  STOr s1 s2 -> do
    Dict <- getWTP' s1
    Dict <- getWTP' s2
    pure Dict
  STLambda s1 s2 -> do
    Dict <- getWTP' s1
    Dict <- getWTP' s2
    pure Dict
  STMap s1 s2 -> do
    Dict <- getWTP' s1
    Dict <- getWTP' s2
    Dict <- eitherWellTyped BtNotComparable s1 comparabilityPresence
    pure Dict
  STBigMap s1 s2 -> do
    Dict <- getWTP' s1
    Dict <- getWTP' s2
    Dict <- eitherWellTyped BtNotComparable s1 $ comparabilityPresence
    Refl <- eitherWellTyped BtIsOperation s2 $ tAbsence SPSOp
    Refl <- eitherWellTyped BtHasBigMap s2 $ tAbsence SPSBigMap
    pure Dict
  STInt -> Right Dict
  STNat -> Right Dict
  STString -> Right Dict
  STBytes -> Right Dict
  STMutez -> Right Dict
  STBool -> Right Dict
  STKeyHash -> Right Dict
  STBls12381Fr -> Right Dict
  STBls12381G1 -> Right Dict
  STBls12381G2 -> Right Dict
  STTimestamp -> Right Dict
  STAddress -> Right Dict
  STChest -> Right Dict
  STChestKey -> Right Dict
  STNever -> Right Dict
  STSaplingState s -> withSingI s $ Right Dict
  STSaplingTransaction s -> withSingI s $ Right Dict
  where
    eitherWellTyped
      :: BadTypeForScope
      -> SingT t
      -> (SingT t -> Maybe a)
      -> Either NotWellTyped a
    eitherWellTyped bt sng f = maybeToRight (NotWellTyped (fromSing sng) bt) $ f sng