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

module Michelson.TypeCheck.Types
    ( HST (..)
    , (-:&)
    , pattern (::&+)
    , SomeHST (..)
    , SomeInstrOut (..)
    , SomeInstr (..)
    , SomeContract (..)
    , SomeContractAndStorage (..)
    , BoundVars (..)
    , TcExtFrames
    , NotWellTyped (..)
    , getWTP
    , withWTPm
    , unsafeWithWTP
    , mapSomeContract
    , mapSomeInstr
    , mapSomeInstrOut
    , noBoundVars
    ) where

import Data.Constraint (Dict(..))
import qualified Data.Map.Lazy as Map
import Data.Singletons (Sing, SingI(..), demote)
import Fmt (Buildable(..), Builder, (+|), (|+), pretty)
import Prelude hiding (EQ, GT, LT)
import qualified Text.Show

import Michelson.Typed (ParameterScope, StorageScope, Notes(..), T(..), SingT(..), notesT, starNotes)
import Michelson.Typed.Haskell.Value (WellTyped)
import qualified Michelson.Typed as T
import Michelson.Typed.Instr
import Michelson.Untyped (Type, Var, noAnn)
import Michelson.Untyped.Annotation (VarAnn)
import Util.Typeable

-- | Data type holding type information for stack (Heterogeneous Stack Type).
--
-- This data type is used along with instruction data type @Instr@
-- to carry information about its input and output stack types.
--
-- That is, if there is value @instr :: Instr inp out@, along with this
-- @instr@ one may carry @inpHST :: HST inp@ and @outHST :: HST out@ which will
-- contain whole information about input and output stack types for @instr@.
--
-- Data type @HST@ is very similar to @Data.Vinyl.Rec@,
-- but is specialized for a particular purpose.
-- In particular, definition of @HST (t1 ': t2 ': ... tn ': '[])@ requires
-- constraints @(Typeable t1, Typeable t2, ..., Typeable tn)@ as well as
-- constraints @(Typeable '[ t1 ], Typeable '[ t1, t2 ], ...)@.
-- These applications of @Typeable@ class are required for convenient usage
-- of type encoded by @HST ts@ with some functions from @Data.Typeable@.
--
-- Data type @HST@ (Heterogeneous Stack Type) is a heterogenuous list of tuples.
-- First element of tuple is a structure, holding field and type annotations
-- for a given type.
-- Second element of tuple is an optional variable annotation for the stack
-- element.
-- Additionally constructor keeps 'SingI' constraint for the current type.
data HST (ts :: [T])  where
  SNil :: HST '[]
  (::&) :: (Typeable xs, T.KnownT x)
        => (Notes x, Dict (WellTyped x), VarAnn)
        -> HST xs
        -> HST (x ': xs)

instance NFData (HST ts) where
  rnf (SNil) = ()
  rnf ((a, d, b) ::& hst) = rnf (a, d, b, hst)

instance Show (HST ts) where
  show SNil = "[]"
  show (r ::& rs) = "[ " <> showDo (r ::& rs) <> " ]"
    where
      showDo :: HST (t ': ts_) -> String
      showDo ((notesT -> t, Dict, _vn) ::& (b ::& c)) =
          show t <> ", " <> showDo (b ::& c)
      showDo ((notesT -> t, Dict, _vn) ::& SNil) = show t

instance Buildable (HST ts) where
  build SNil = "[]"
  build (r ::& rs) = "[ " +| showDo (r ::& rs) |+ " ]"
    where
      showDo :: HST (t ': ts_) -> Builder
      showDo ((notesT -> t, Dict, _vn) ::& (b ::& c)) =
          build t |+ ", " +| showDo (b ::& c)
      showDo ((notesT -> t, Dict, _vn) ::& SNil) = build t

infixr 7 ::&

instance Eq (HST ts) where
  SNil == SNil = True
  (n1, Dict, a1) ::& h1 == (n2, Dict, a2) ::& h2 =
    n1 == n2 && a1 == a2 && h1 == h2

-- | Append a type to 'HST', assuming that notes and annotations
-- for this type are unknown.
(-:&)
  :: (Typeable xs, WellTyped x)
  => Sing x
  -> HST xs
  -> HST (x ': xs)
_ -:& hst = (starNotes, Dict, noAnn) ::& hst
infixr 7 -:&

-- | Extended pattern-match - adds @Sing x@ argument.
infixr 7 ::&+
pattern (::&+)
  :: ()
  => ( ys ~ (x ': xs)
     , T.KnownT x, Typeable xs
     )
  => (Sing x, Notes x, Dict (WellTyped x), VarAnn)
  -> HST xs
  -> HST ys
pattern x ::&+ hst <- ((\(n, d, v) -> (T.notesSing n, n, d, v)) -> x) ::& hst
  where (_, n, d, v) ::&+ hst = (n, d, v) ::& hst

-- | No-argument type wrapper for @HST@ data type.
data SomeHST where
  SomeHST :: Typeable ts => HST ts -> SomeHST

deriving stock instance Show SomeHST
instance NFData SomeHST where
  rnf (SomeHST h) = rnf h

instance Eq SomeHST where
  SomeHST hst1 == SomeHST hst2 = hst1 `eqParam1` hst2

-- | This data type keeps part of type check result - instruction and
-- corresponding output stack.
data SomeInstrOut inp where
  -- | Type-check result with concrete output stack, most common case.
  --
  -- Output stack type is wrapped inside the type and @Typeable@
  -- constraint is provided to allow convenient unwrapping.
  (:::)
    :: (Typeable out)
    => Instr inp out
    -> HST out
    -> SomeInstrOut inp

  -- | Type-check result which matches against arbitrary output stack.
  -- Information about annotations in the output stack is absent.
  --
  -- This case is only possible when the corresponding code terminates
  -- with @FAILWITH@ instruction in all possible executions.
  -- The opposite may be not true though (example: you push always-failing
  -- lambda and immediatelly execute it - stack type is known).
  AnyOutInstr
    :: (forall out. Instr inp out)
    -> SomeInstrOut inp
infix 9 :::

instance Show (ExtInstr inp) => Show (SomeInstrOut inp) where
  show (i ::: out) = show i <> " :: " <> show out
  show (AnyOutInstr i) = show i <> " :: *"

-- | Data type keeping the whole type check result: instruction and
-- type representations of instruction's input and output.
data SomeInstr inp where
  (:/) :: HST inp -> SomeInstrOut inp -> SomeInstr inp
infix 8 :/

mapSomeInstrOut
  :: (forall out. Instr inp out -> Instr inp' out)
  -> SomeInstrOut inp
  -> SomeInstrOut inp'
mapSomeInstrOut f (i ::: out) = f i ::: out
mapSomeInstrOut f (AnyOutInstr i) = AnyOutInstr (f i)

mapSomeInstr
  :: (forall out. Instr inp out -> Instr inp out)
  -> SomeInstr inp
  -> SomeInstr inp
mapSomeInstr f (inp :/ instrAndOut) = inp :/ mapSomeInstrOut f instrAndOut

instance Show (ExtInstr inp) => Show (SomeInstr inp) where
  show (inp :/ out) = show inp <> " -> " <> show out

data SomeContract where
  SomeContract :: Contract cp st -> SomeContract

instance NFData SomeContract where
  rnf (SomeContract c) = rnf c

mapSomeContract ::
  (forall inp out. Instr inp out -> Instr inp out)
  -> SomeContract
  -> SomeContract
mapSomeContract f (SomeContract fc) = SomeContract $ mapContractCode f fc

deriving stock instance Show SomeContract

-- | Represents a typed contract & a storage value of the type expected by the contract.
data SomeContractAndStorage where
  SomeContractAndStorage
    :: forall cp st.
       (StorageScope st, ParameterScope cp)
       => Contract cp st
       -> T.Value st
       -> SomeContractAndStorage

deriving stock instance Show SomeContractAndStorage

-- | Set of variables defined in a let-block.
data BoundVars = BoundVars (Map Var Type) (Maybe SomeHST)

noBoundVars :: BoundVars
noBoundVars = BoundVars Map.empty Nothing

-- | State for type checking @nop@
type TcExtFrames = [BoundVars]

-- | Error type for when a value is not well-typed.
data NotWellTyped = NotWellTyped T

instance Buildable NotWellTyped where
  build (NotWellTyped t) =
    "Given type is not well typed because it has an non-comparable type in it: \
    \'" <> (build t) <> "', where a comparable type is required"

fromEDict :: Either e (Dict a) -> (a => Either e (Dict b)) -> Either e (Dict b)
fromEDict ma b = ma >>= (\Dict -> b)

-- | Given a type, provide evidence that it is well typed w.r.t to the
--  Michelson rules regarding where comparable types are required.
getWTP :: forall t. (SingI t) => Either NotWellTyped (Dict (WellTyped t))
getWTP = case sing @t of
  STKey -> Right Dict
  STUnit -> Right Dict
  STSignature -> Right Dict
  STChainId -> Right Dict
  STOption s -> fromEDict (getWTP_ s) (Right Dict)
  STList s -> fromEDict (getWTP_ s) (Right Dict)
  STSet (s :: Sing si) ->
    fromEDict (getWTP_ s) $
    fromEDict
        ( maybeToRight (NotWellTyped $ demote @si) $ T.getComparableProofS s
        ) $ Right Dict
  STOperation  -> Right Dict
  STContract s -> fromEDict (getWTP_ s) (Right Dict)
  STPair s1 s2 -> fromEDict (getWTP_ s1) $ fromEDict (getWTP_ s2) $ Right Dict
  STOr s1 s2 -> fromEDict (getWTP_ s1) $ fromEDict (getWTP_ s2) $ Right Dict
  STLambda s1 s2 -> fromEDict (getWTP_ s1) $ fromEDict (getWTP_ s2) $ Right Dict
  STMap (s1 :: Sing si) s2 ->
      fromEDict (getWTP_ s1) $
      fromEDict (getWTP_ s2) $
      fromEDict
        ( maybeToRight (NotWellTyped $ demote @si) $ T.getComparableProofS s1
        ) $ Right Dict
  STBigMap (s1 :: Sing si) s2 ->
      fromEDict (getWTP_ s1) $
      fromEDict (getWTP_ s2) $
      fromEDict
        ( maybeToRight (NotWellTyped $ demote @si) $ T.getComparableProofS s1
        ) $ Right Dict
  STInt -> Right Dict
  STNat -> Right Dict
  STString -> Right Dict
  STBytes -> Right Dict
  STMutez -> Right Dict
  STBool -> Right Dict
  STKeyHash -> Right Dict
  STTimestamp -> Right Dict
  STAddress -> Right Dict
  where
    getWTP_ :: forall t1. SingI t1 => Sing t1 -> Either NotWellTyped (Dict (WellTyped t1))
    getWTP_ _ = getWTP @t1

-- | Given a type and an action that requires evidence that the type is well typed,
--  generate the evidence and execute the action, or else fail with an error.
withWTPm :: forall t m a. (MonadFail m, SingI t) => (WellTyped t => m a) -> m a
withWTPm a = case getWTP @t of
  Right Dict -> a
  Left e -> fail (pretty e)

-- | Similar to @withWTPm@ but is mean to be used within tests.
unsafeWithWTP :: forall t a. SingI t => (WellTyped t => a) -> a
unsafeWithWTP fn = case getWTP @t of
  Right Dict -> fn
  Left e -> error $ pretty e