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

-- This module provides data types for representing partially typed
-- instructions.
module Morley.Michelson.TypeCheck.TypeCheckedOp
  ( TypeCheckedInstr
  , TypeCheckedOp(..)
  , IllTypedInstr(..)
  , Nesting
  , SomeSingInstr(..)
  , someInstrToOp
  , someViewToOp
  , withSomeTcInstr
  ) where

import Data.Singletons (SingI)
import Fmt (Doc)
import Prettyprinter (enclose)

import Morley.Michelson.Printer.Util
  (RenderDoc(..), doesntNeedParens, renderOpsList, renderOpsListNoBraces)
import Morley.Michelson.TypeCheck.Types (HST(..), SomeTcInstr(..), SomeTcInstrOut(..))
import Morley.Michelson.Typed.Aliases
import Morley.Michelson.Typed.Annotation
import Morley.Michelson.Typed.Convert (instrToOps)
import Morley.Michelson.Typed.Instr (Instr, castInstr)
import Morley.Michelson.Typed.View (SomeView'(..), View'(..))
import Morley.Michelson.Untyped qualified as U
import Morley.Michelson.Untyped.Instr (InstrAbstract(..))

-- | Represents a root of a partially typed operation tree.
type TypeCheckedInstr op = InstrAbstract [] (TypeCheckedOp op)

newtype Nesting = Nesting Word
  deriving newtype (Eq, Ord, Enum, Bounded, Num, NFData)
  deriving stock Show

data SomeSingInstr where
  SomeSingInstr :: (SingI inp, SingI out) => Instr inp out -> SomeSingInstr

deriving stock instance Show SomeSingInstr

instance Eq SomeSingInstr where
  SomeSingInstr i1 == SomeSingInstr i2 = castInstr i1 == Just i2

instance NFData SomeSingInstr where
  rnf (SomeSingInstr x) = rnf x

-- | Represents nodes of a partially typed operation tree.
data TypeCheckedOp op where
  -- | Constructs well-typed node.
  WellTypedOp :: SomeSingInstr -> TypeCheckedOp op
  -- | Constructs ill-typed node which might in turn contain well-typed and
  -- non-typed operations.
  IllTypedOp :: [IllTypedInstr op] -> TypeCheckedOp op
  -- | Partially typed sequence of operations. Used exclusively for error-reporting.
  -- 'Nesting' argument exists because we can't mix typed and untyped
  -- instructions, so we need a way to represent brace nesting of @{ <typed
  -- seq>; <untyped seq> }@
  MixedOp :: Nesting -> SomeSingInstr -> [IllTypedInstr op] -> TypeCheckedOp op
  deriving stock (Eq, Functor, Show)

-- | Represents a non-well-typed operation
data IllTypedInstr op
  = SemiTypedInstr (TypeCheckedInstr op) -- ^ Constructs a partially typed operation.
  | NonTypedInstr op  -- ^ Constructs a completely untyped operation.
  | IllTypedNest [IllTypedInstr op] -- ^ Nested sequence of ill-typed operations.
  deriving stock (Show, Eq, Generic, Functor)

instance RenderDoc op => RenderDoc (IllTypedInstr op) where
  renderDoc pn = \case
    SemiTypedInstr instr -> renderDoc pn instr
    NonTypedInstr op -> renderDoc pn op
    IllTypedNest ops -> renderOpsList False ops

deriving anyclass instance (NFData (TypeCheckedOp op), NFData op) => NFData (IllTypedInstr op)

instance RenderDoc op => RenderDoc (TypeCheckedOp op) where
  renderDoc _ (WellTypedOp (SomeSingInstr instr)) = renderOpsListNoBraces False (instrToOps instr)
  renderDoc _ (MixedOp nesting (SomeSingInstr instr) rest) =
    nestBraces nesting $ renderOpsListNoBraces False $ doRender (instrToOps instr) <> doRender rest
    where
      doRender :: RenderDoc a => [a] -> [Doc]
      doRender = fmap (renderDoc doesntNeedParens) . filter isRenderable
  renderDoc _ (IllTypedOp ops) = renderOpsListNoBraces False ops

nestBraces :: Nesting -> Doc -> Doc
nestBraces 0 = id
nestBraces n = enclose "{ " " }" . nestBraces (pred n)

-- | Makes a well-typed node out of `SomeTcInstr`
someInstrToOp :: SomeTcInstr inp -> TypeCheckedOp op
someInstrToOp hst = withSomeTcInstr hst $ WellTypedOp . SomeSingInstr

-- | Extracts 'Instr' from 'SomeTcInstr' and passes it to a continuation.
withSomeTcInstr
  :: SomeTcInstr inp
  -> (forall out. (SingI inp, SingI out) => Instr inp out -> r)
  -> r
withSomeTcInstr (hst :/ rest) = case hst of
  SNil -> go rest
  (::&){} -> go rest
  where
    go :: SingI inp => SomeTcInstrOut inp
       -> (forall out. (SingI inp, SingI out) => Instr inp out -> r)
       -> r
    go (i ::: _) cont = cont i
    go (AnyOutInstr i) cont = cont (i @'[])

-- | Makes takes a typed view and converts it into an untyped one with
-- typechecked code.
someViewToOp :: SomeView st -> U.View' (TypeCheckedOp op)
someViewToOp (SomeView View{..}) = U.View
  { U.viewCode = WellTypedOp $ SomeSingInstr vCode
  , U.viewName = vName
  , U.viewArgument = mkUType vArgument
  , U.viewReturn = mkUType vReturn
  }

instance NFData op => NFData (TypeCheckedOp op) where
  rnf = \case
    WellTypedOp x -> rnf x
    IllTypedOp x -> rnf x
    MixedOp x y z -> rnf x `seq` rnf y `seq` rnf z