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

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

import Data.Constraint (Dict(..))
import qualified Data.Map.Lazy as Map
import Data.Singletons (Sing, SingI(..), fromSing)
import Fmt (Buildable(..), pretty)
import Prelude hiding (EQ, GT, LT)
import Text.PrettyPrint.Leijen.Text hiding (pretty)
import qualified Text.Show


import Morley.Michelson.Printer.Util
import Morley.Michelson.Typed (Notes(..), SingT(..), SomeContract(..), T(..), notesT, starNotes)
import qualified Morley.Michelson.Typed as T
import Morley.Michelson.Typed.Contract
import Morley.Michelson.Typed.Haskell.Value (WellTyped)
import Morley.Michelson.Typed.Instr
import Morley.Michelson.Untyped (Ty, Var, noAnn)
import Morley.Michelson.Untyped.Annotation (VarAnn)
import Morley.Util.Sing (eqParamSing)

-- | 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 '[]
  (::&) :: (T.SingI x, T.SingI xs)
        => (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 = buildRenderDocExtended

instance RenderDoc (HST ts) where
  renderDoc _ SNil = "[]"
  renderDoc context (r ::& rs) = "[" <+> doRender (r ::& rs) <+> "]"
    where
      doRender :: HST (t ': ts_) -> Doc
      doRender ((notesT -> t, Dict, _vn) ::& (b ::& c)) =
          renderDoc context t <> "," <+> doRender (b ::& c)
      doRender ((notesT -> t, Dict, _vn) ::& SNil) = renderDoc context 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.
(-:&)
  :: (WellTyped x, SingI xs)
  => 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)
     , SingI x, SingI 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 :: SingI 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 `eqParamSing` 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.
  (:::)
    :: SingI 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

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

-- | Set of variables defined in a let-block.
data BoundVars = BoundVars (Map Var Ty) (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
  { nwtBadType :: T
  , nwtCause :: T.BadTypeForScope
  }

instance Buildable NotWellTyped where
  build (NotWellTyped t c) =
    "Given type is not well typed because '" <> (build 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.
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 T.BtNotComparable T.getComparableProofS s
    pure Dict
  STOperation -> Right Dict
  STContract s -> do
    Dict <- getWTP' s
    Dict <- eitherWellTyped T.BtIsOperation T.opAbsense s
    pure Dict
  STTicket s -> do
    Dict <- getWTP' s
    Dict <- eitherWellTyped T.BtNotComparable T.getComparableProofS s
    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 T.BtNotComparable T.getComparableProofS s1
    pure Dict
  STBigMap s1 s2 -> do
    Dict <- getWTP' s1
    Dict <- getWTP' s2
    Dict <- eitherWellTyped T.BtNotComparable T.getComparableProofS s1
    Dict <- eitherWellTyped T.BtIsOperation T.opAbsense s2
    Dict <- eitherWellTyped T.BtHasBigMap T.bigMapAbsense s2
    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
  where
    eitherWellTyped
      :: T.BadTypeForScope
      -> (Sing (ty :: T) -> Maybe a)
      -> Sing (ty :: T)
      -> Either NotWellTyped a
    eitherWellTyped bt sf sng = maybeToRight (NotWellTyped (fromSing sng) bt) $ sf sng

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