module Michelson.TypeCheck.Types
    ( HST (..)
    , (-:&)
    , pattern (::&+)
    , SomeHST (..)
    , SomeInstrOut (..)
    , SomeInstr (..)
    , SomeContract (..)
    , BoundVars (..)
    , TcExtFrames
    , mapSomeContract
    , noBoundVars
    ) where

import Data.Singletons (SingI)
import Prelude hiding (EQ, GT, LT)
import qualified Text.Show
import qualified Data.Map.Lazy as Map

import Michelson.Untyped (Var, Type, noAnn)
import Michelson.Untyped.Annotation (VarAnn)
import Michelson.Typed
  (Notes(..), Sing(..), T(..), starNotes, notesT)
import qualified Michelson.Typed as T
import Michelson.Typed.Instr
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, Typeable x, SingI x)
        => (Notes x, VarAnn)
        -> HST xs
        -> HST (x ': xs)

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

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

infixr 7 ::&

instance Eq (HST ts) where
  SNil == SNil = True
  (n1, a1) ::& h1 == (n2, 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, Typeable x, SingI x)
  => Sing x
  -> HST xs
  -> HST (x ': xs)
_ -:& hst = (starNotes, noAnn) ::& hst
infixr 7 -:&

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

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

data SomeContract where
  SomeContract :: FullContract 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 $ mapFullContractCode f fc

deriving stock instance Show SomeContract

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