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

module 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(..), demote, withSingI)
import Fmt (Buildable(..), pretty)
import Prelude hiding (EQ, GT, LT)
import qualified Text.Show
import Text.PrettyPrint.Leijen.Text hiding (pretty)


import Michelson.Typed (Notes(..), SingT(..), SomeContract(..), T(..), notesT, starNotes)
import qualified Michelson.Typed as T
import Michelson.Typed.Haskell.Value (WellTyped)
import Michelson.Typed.Instr
import Michelson.Printer.Util
import Michelson.Untyped (Ty, Var, noAnn)
import Michelson.Untyped.Annotation (VarAnn)
import 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)

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 -> withSingI s $ fromEDict (getWTP_ s) (Right Dict)
  STList s -> withSingI s $ fromEDict (getWTP_ s) (Right Dict)
  STSet (s :: Sing si) ->
    withSingI s $
    fromEDict (getWTP_ s) $
    fromEDict
        ( maybeToRight (NotWellTyped (demote @si) T.BtNotComparable) $ T.getComparableProofS s
        ) $ Right Dict
  STOperation  -> Right Dict
  STContract (s :: Sing si) ->
    withSingI s $ fromEDict (getWTP_ s) $
    fromEDict
        ( maybeToRight (NotWellTyped (demote @si) T.BtIsOperation) $ T.opAbsense s
        ) $ Right Dict

  STTicket (s :: Sing si) ->
    withSingI s $ fromEDict (getWTP_ s) $
    fromEDict
        ( maybeToRight (NotWellTyped (demote @si) T.BtNotComparable) $ T.getComparableProofS s
        ) $ Right Dict
  STPair s1 s2 ->
    withSingI s1 $
    withSingI s2 $
    fromEDict (getWTP_ s1) $
    fromEDict (getWTP_ s2) $ Right Dict
  STOr s1 s2 ->
    withSingI s1 $
    withSingI s2 $
    fromEDict (getWTP_ s1) $
    fromEDict (getWTP_ s2) $ Right Dict
  STLambda s1 s2 ->
    withSingI s1 $
    withSingI s2 $
    fromEDict (getWTP_ s1) $ fromEDict (getWTP_ s2) $ Right Dict
  STMap (s1 :: Sing si) s2 ->
      withSingI s1 $
      withSingI s2 $
      fromEDict (getWTP_ s1) $
      fromEDict (getWTP_ s2) $
      fromEDict
        ( maybeToRight (NotWellTyped (demote @si) T.BtNotComparable) $ T.getComparableProofS s1
        ) $ Right Dict
  STBigMap (s1 :: Sing si) (s2 :: Sing si2) ->
      withSingI s1 $
      withSingI s2 $
      fromEDict (getWTP_ s1) $
      fromEDict (getWTP_ s2) $
      fromEDict
      ( maybeToRight (NotWellTyped (demote @si) T.BtNotComparable) $ T.getComparableProofS s1
      ) $
      fromEDict
      ( maybeToRight (NotWellTyped (demote @si2) T.BtIsOperation) $ T.opAbsense s2
      ) $
      fromEDict
      ( maybeToRight (NotWellTyped (demote @si2) T.BtHasBigMap) $ T.bigMapAbsense s2
      ) $ Right 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
  STNever -> Right Dict

-- | Version of 'getWTP' that accepts 'Sing' at term-level.
getWTP_ :: forall t. Sing t -> Either NotWellTyped (Dict (WellTyped t))
getWTP_ s = withSingI s $ getWTP @t

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