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

{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Foundation of Lorentz development.
module Lorentz.Base
  ( (:->) (..)
  , type (%>)
  , type (&)
  , (#)
  , (##)
  , pattern I
  , pattern FI
  , iGenericIf
  , iAnyCode
  , iNonFailingCode
  , iMapAnyCode
  , iForceNotFail
  , iWithVarAnnotations

  , parseLorentzValue
  , transformStringsLorentz
  , transformBytesLorentz
  , optimizeLorentz
  , optimizeLorentzWithConf
  , MapLorentzInstr (..)

  , ContractOut
  , ContractCode
  , SomeContractCode (..)
  , Lambda
  ) where

import Data.Default (def)
import qualified Data.Kind as Kind
import qualified Data.List.NonEmpty as NE (fromList)
import Fmt (Buildable(..))

import Lorentz.Constraints
import Michelson.ErrorPos (InstrCallStack)
import Michelson.Optimizer (OptimizerConf, optimizeWithConf)
import Michelson.Parser (ParserException, parseExpandValue)
import Michelson.Preprocess (transformBytes, transformStrings)
import Michelson.Text (MText)
import Michelson.TypeCheck (TCError, runTypeCheckIsolated, typeCheckValue)
import Michelson.Typed
  (Instr(..), IsoValue(..), Operation, RemFail(..), ToT, ToTs, Value, rfAnyInstr, rfMapAnyInstr,
  rfMerge)
import qualified Michelson.Untyped as U

-- | Alias for instruction which hides inner types representation via 'T'.
newtype (inp :: [Kind.Type]) :-> (out :: [Kind.Type]) = LorentzInstr
  { unLorentzInstr :: RemFail Instr (ToTs inp) (ToTs out)
  } deriving newtype (Show, Eq)
infixr 1 :->

instance Semigroup (s :-> s) where
  (<>) = (#)
instance Monoid (s :-> s) where
  mempty = I Nop

pattern I :: Instr (ToTs inp) (ToTs out) -> inp :-> out
pattern I i = LorentzInstr (RfNormal i)

pattern FI :: (forall out'. Instr (ToTs inp) out') -> inp :-> out
pattern FI i = LorentzInstr (RfAlwaysFails i)

{-# COMPLETE I, FI #-}

iGenericIf
  :: (forall s'. Instr (ToTs a) s' -> Instr (ToTs b) s' -> Instr (ToTs c) s')
  -> (a :-> s) -> (b :-> s) -> (c :-> s)
iGenericIf merger (LorentzInstr instr1) (LorentzInstr instr2) =
  LorentzInstr (rfMerge merger instr1 instr2)

iAnyCode :: inp :-> out -> Instr (ToTs inp) (ToTs out)
iAnyCode = rfAnyInstr . unLorentzInstr

iNonFailingCode :: HasCallStack => inp :-> out -> Instr (ToTs inp) (ToTs out)
iNonFailingCode (I i) = i
iNonFailingCode (FI _) = error "Always failing code cannot appear here"

iMapAnyCode
  :: (forall o'. Instr (ToTs i1) o' -> Instr (ToTs i2) o')
  -> (i1 :-> o)
  -> (i2 :-> o)
iMapAnyCode f (LorentzInstr i) = LorentzInstr $ rfMapAnyInstr f i

iForceNotFail :: (i :-> o) -> (i :-> o)
iForceNotFail = I . iAnyCode

-- | Wrap Lorentz instruction with variable annotations, @annots@ list has to be
-- non-empty, otherwise this function raises an error.
iWithVarAnnotations :: HasCallStack => [Text] -> inp :-> out -> inp :-> out
iWithVarAnnotations annots (LorentzInstr i) = case i of
  RfNormal instr -> LorentzInstr $ RfNormal $
    InstrWithVarNotes (NE.fromList $ map U.AnnotationUnsafe annots) instr
  RfAlwaysFails instr -> LorentzInstr $ RfAlwaysFails $
    InstrWithVarNotes (NE.fromList $ map U.AnnotationUnsafe annots) instr

-- There is also @instance IsoValue (i :-> o)@ in a separate module.

-- | Alias for ':->', seems to make signatures more readable sometimes.
--
-- Let's someday decide which one of these two should remain.
type (%>) = (:->)
infixr 1 %>

type ContractOut st = '[([Operation], st)]
type ContractCode cp st = '[(cp, st)] :-> ContractOut st

data SomeContractCode where
  SomeContractCode
    :: (NiceParameterFull cp, NiceStorage st)
    => ContractCode cp st
    -> SomeContractCode

type (&) (a :: Kind.Type) (b :: [Kind.Type]) = a ': b
infixr 2 &

(#) :: (a :-> b) -> (b :-> c) -> a :-> c
I l # I r = I (l `Seq` r)
I l # FI r = FI (l `Seq` r)
FI l # _ = FI l
infixl 8 #

-- | Version of '#' which performs some optimizations immediately.
--
-- In particular, this avoids glueing @Nop@s.
(##) :: (a :-> b) -> (b :-> c) -> (a :-> c)
l ## r =
  -- We are very verbose about cases to avoid
  -- significant compilation time increase
  case l of
    I Nop -> case r of
      I x -> I x
      _ -> l # r
    I x -> case r of
      I Nop -> I x
      _ -> l # r
    _ -> l # r


type Lambda i o = '[i] :-> '[o]

-- | Errors that can happen during parsing into a Lorentz value.
data ParseLorentzError
  = ParseLorentzParseError ParserException
  | ParseLorentzTypecheckError TCError
  deriving stock (Show, Eq)

instance Buildable ParseLorentzError where
  build =
    \case
      ParseLorentzParseError e -> build e
      ParseLorentzTypecheckError e -> build e

-- | Parse textual representation of a Michelson value and turn it
-- into corresponding Haskell value.
--
-- Note: it won't work in some complex cases, e. g. if there is a
-- lambda which uses an instruction which depends on current
-- contract's type. Obviously it can not work, because we don't have
-- any information about a contract to which this value belongs (there
-- is no such contract at all).
parseLorentzValue
  :: forall v. KnownValue v
  => Text
  -> Either ParseLorentzError v
parseLorentzValue =
  fmap fromVal . (toTyped <=< first ParseLorentzParseError . parseExpandValue)
  where
    toTyped :: U.Value -> Either ParseLorentzError (Value (ToT v))
    toTyped =
      first ParseLorentzTypecheckError .
      runTypeCheckIsolated .
      usingReaderT (def @InstrCallStack) .
      typeCheckValue

-- | Lorentz version of 'transformStrings'.
transformStringsLorentz ::
  Bool -> (MText -> MText) -> inp :-> out -> inp :-> out
transformStringsLorentz goToValues f =
  iMapAnyCode $ transformStrings goToValues f

-- | Lorentz version of 'transformBytes'.
transformBytesLorentz ::
  Bool -> (ByteString -> ByteString) -> inp :-> out -> inp :-> out
transformBytesLorentz goToValues f =
  iMapAnyCode $ transformBytes goToValues f

optimizeLorentzWithConf
  :: OptimizerConf
  -> inp :-> out
  -> inp :-> out
optimizeLorentzWithConf conf =
  -- Optimizer can produce dead code.
  -- Example: @push True # if_ failWith nop # ...@ will fold to @failWith # ...@.
  -- But let's not care about this case for now until need in it fires.
  iMapAnyCode (optimizeWithConf conf)

optimizeLorentz
  :: inp :-> out
  -> inp :-> out
optimizeLorentz = optimizeLorentzWithConf def

-- | Applicable for wrappers over Lorentz code.
class MapLorentzInstr instr where
  -- | Modify all the code under given entity.
  mapLorentzInstr :: (forall i o. (i :-> o) -> (i :-> o)) -> instr -> instr

instance MapLorentzInstr (i :-> o) where
  mapLorentzInstr f = f