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

{- | Lorentz contracts compilation.

Compilation in one scheme:

@
                    mkContract
                  mkContractWith
  ContractCode  -----------------→  Contract
 (Lorentz code)              (ready compiled contract)
        ↓                                ↑
        ↓                                ↑
defaultContractData            compileLorentzContract
   ContractData                          ↑
        ↓         ContractData           ↑
       (Lorentz code + compilation options)
@

-}
module Lorentz.Run
  ( Contract(..)
  , toMichelsonContract
  , defaultContract

  , CompilationOptions(..)
  , defaultCompilationOptions
  , intactCompilationOptions
  , coBytesTransformerL
  , coOptimizerConfL
  , coStringTransformerL

  , compileLorentz
  , compileLorentzWithOptions

  , mkContract
  , mkContractWith

  , ContractData(..)
  , ContractView(..)
  , defaultContractData
  , compileLorentzContract
  , mkView
  , setViews
  , setViewsRec
  , noViews
  , cdCodeL
  , coDisableInitialCastL
  , cdCompilationOptionsL

  , interpretLorentzInstr
  , interpretLorentzLambda

  , analyzeLorentz
  ) where

import Control.Lens.Type as Lens (Lens, Lens')
import Data.Constraint ((\\))
import Data.Default (def)
import Data.Vinyl.Core (Rec(..))
import Data.Vinyl.Functor qualified as Rec
import Data.Vinyl.Recursive qualified as Rec
import Fmt ((+|), (|+))

import Lorentz.Annotation
import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Constraints
import Lorentz.Doc
import Lorentz.Entrypoints
import Lorentz.Entrypoints.Doc
import Lorentz.ViewBase
import Morley.Michelson.Analyzer (AnalyzerRes, analyze)
import Morley.Michelson.Interpret
import Morley.Michelson.Optimizer (OptimizerConf, optimizeWithConf)
import Morley.Michelson.Text (MText)
import Morley.Michelson.TypeCheck (typeCheckingWith, typeVerifyContract, typeVerifyView)
import Morley.Michelson.Typed
  (Instr(..), IsoValue, IsoValuesStack(..), ToT, ToTs, convertContract, convertView, starParamNotes)
import Morley.Michelson.Typed qualified as M
import Morley.Michelson.Untyped qualified as U (canonicalEntriesOrder)
import Morley.Util.Lens
import Morley.Util.TypeLits
import Morley.Util.TypeTuple

-- | Options to control Lorentz to Michelson compilation.
data CompilationOptions = CompilationOptions
  { coOptimizerConf :: Maybe OptimizerConf
  -- ^ Config for Michelson optimizer.
  , coStringTransformer :: (Bool, MText -> MText)
  -- ^ Function to transform strings with. See 'transformStringsLorentz'.
  , coBytesTransformer :: (Bool, ByteString -> ByteString)
  -- ^ Function to transform byte strings with. See 'transformBytesLorentz'.
  , coDisableInitialCast :: Bool
  -- ^ Flag which defines whether compiled Michelson contract
  -- will have @CAST@ (which drops parameter annotations)
  -- as a first instruction. Note that when
  -- flag is false, there still may be no @CAST@ (in case
  -- when parameter type has no annotations).
  }

-- | Runs Michelson optimizer with default config and does not touch strings and bytes.
defaultCompilationOptions :: CompilationOptions
defaultCompilationOptions = CompilationOptions
  { coOptimizerConf = Just def
  , coStringTransformer = (False, id)
  , coBytesTransformer = (False, id)
  , coDisableInitialCast = False
  }

-- | Leave contract without any modifications. For testing purposes.
intactCompilationOptions :: CompilationOptions
intactCompilationOptions = CompilationOptions
  { coOptimizerConf = Nothing
  , coStringTransformer = (False, id)
  , coBytesTransformer = (False, id)
  , coDisableInitialCast = False
  }

-- | For use outside of Lorentz. Will use 'defaultCompilationOptions'.
compileLorentz :: (inp :-> out) -> Instr (ToTs inp) (ToTs out)
compileLorentz = compileLorentzWithOptions defaultCompilationOptions

-- | Compile Lorentz code, optionally running the optimizer, string and byte transformers.
compileLorentzWithOptions :: CompilationOptions -> (inp :-> out) -> Instr (ToTs inp) (ToTs out)
compileLorentzWithOptions CompilationOptions{..} =
  maybe id optimizeWithConf coOptimizerConf
  . iAnyCode
  . uncurry transformStringsLorentz coStringTransformer
  . uncurry transformBytesLorentz coBytesTransformer

-- | Construct and compile Lorentz contract.
--
-- This is an alias for 'mkContract'.
defaultContract
  :: (NiceParameterFull cp, NiceStorageFull st)
  => ContractCode cp st -> Contract cp st ()
defaultContract code =
  compileLorentzContract $ ContractData code mempty defaultCompilationOptions

-- | Construct and compile Lorentz contract.
--
-- Note that this accepts code with initial and final stacks unpaired for
-- simplicity.
mkContract
  :: (NiceParameterFull cp, NiceStorageFull st)
  => ContractCode cp st -> Contract cp st ()
mkContract = mkContractWith defaultCompilationOptions

-- | Version of 'mkContract' that accepts custom compilation options.
mkContractWith
  :: (NiceParameterFull cp, NiceStorageFull st)
  => CompilationOptions -> ContractCode cp st -> Contract cp st ()
mkContractWith opts code =
  compileLorentzContract $ ContractData code mempty opts

-- | Code for a contract along with compilation options for the Lorentz compiler.
--
-- It is expected that a 'Contract' is one packaged entity, wholly controlled by its author.
-- Therefore the author should be able to set all options that control contract's behavior.
--
-- This helps ensure that a given contract will be interpreted in the same way in all
-- environments, like production and testing.
--
-- Raw 'ContractCode' should not be used for distribution of contracts.
data ContractData cp st vd =
  (NiceParameterFull cp, NiceStorageFull st, NiceViewsDescriptor vd) =>
  ContractData
  { cdCode :: ContractCode cp st
  -- ^ The contract itself.
  , cdViews :: Rec (ContractView st) (RevealViews vd)
  -- ^ Contract views.
  , cdCompilationOptions :: CompilationOptions
  -- ^ General compilation options for the Lorentz compiler.
  }

-- | Single contract view.
data ContractView st (v :: ViewTyInfo) where
  ContractView
    :: ( KnownSymbol name, NiceViewable arg, NiceViewable ret
       , HasAnnotation arg, HasAnnotation ret
       )
    => ViewCode arg st ret
    -> ContractView st ('ViewTyInfo name arg ret)

-- | Construct a view.
--
-- > mkView @"add" @(Integer, Integer) do
-- >  car; unpair; add
mkView
  :: forall name arg ret st.
     ( KnownSymbol name, NiceViewable arg, NiceViewable ret
     , HasAnnotation arg, HasAnnotation ret
     , TypeHasDoc arg, TypeHasDoc ret
     )
  => ViewCode arg st ret
  -> ContractView st ('ViewTyInfo name arg ret)
mkView code = ContractView $
  docGroup (DView (demoteViewName @name)) $
    doc (DViewArg (Proxy @arg)) #
    doc (DViewRet (Proxy @ret)) #
    code

-- | Compile contract with 'defaultCompilationOptions'.
defaultContractData
  :: forall cp st. (NiceParameterFull cp, NiceStorageFull st)
  => ContractCode cp st -> ContractData cp st ()
defaultContractData code = ContractData
  { cdCode = code
  , cdViews = RNil
  , cdCompilationOptions = defaultCompilationOptions
  }

-- | Compile a whole contract to Michelson.
--
-- Note that compiled contract can be ill-typed in terms of Michelson code
-- when some of the compilation options are used (e.g. when 'coDisableInitialCast'
-- is @True@, resulted contract can be ill-typed).
-- However, compilation with 'defaultCompilationOptions' should be valid.
compileLorentzContract
  :: forall cp st vd. ContractData cp st vd -> Contract cp st vd
compileLorentzContract ContractData{..} = Contract{..}
  where
    verify = typeVerifyContract @(ToT cp) @(ToT st) (convertContract cMichelsonContractRaw)
    cMichelsonContract
      | cpNotes == starParamNotes || coDisableInitialCast cdCompilationOptions
      -- If contract parameter type has no annotations or explicitly asked, we drop CAST.
      = cMichelsonContractRaw
      | Right{} <- typeCheckingWith def verify = cMichelsonContractRaw
      -- Check if the contract typechecks
      | otherwise
      -- Perform CAST if it doesn't
      = cMichelsonContractRaw
          { M.cCode = compileLorentzWithOptions
              cdCompilationOptions (I CAST # cdCode :: ContractCode cp st) }
    cMichelsonContractRaw = M.Contract
      { cCode = compileLorentzWithOptions cdCompilationOptions cdCode
      , cParamNotes = cpNotes
      , cStoreNotes = getAnnotation @st NotFollowEntrypoint
      , cEntriesOrder = U.canonicalEntriesOrder
      , cViews = compileLorentzViews cdCompilationOptions cdViews
      } \\ niceParameterEvi @cp
        \\ niceStorageEvi @st

    cDocumentedCode =
      finalizeParamCallingDoc' (Proxy @cp) cdCode #
      foldr (#) (I Nop)
        ( Rec.recordToList $
           Rec.rmap (\(ContractView code) -> Rec.Const $ fakeCoercing code) cdViews
        )

    cpNotes = parameterEntrypointsToNotes @cp

-- | Compile multiple views, with the related checks.
compileLorentzViews
  :: forall vs st.
      ( HasCallStack
      , KnownValue st
      )
  => CompilationOptions
  -> Rec (ContractView st) vs
  -> M.ViewsSet (M.ToT st)
compileLorentzViews co views =
  let
    viewsList =
      Rec.recordToList $
      Rec.rmap (\v -> Rec.Const $ compileLorentzView co v) views
  in case M.mkViewsSet viewsList of
      Right viewsSet -> viewsSet
      Left e@M.DuplicatedViewName{} ->
        error $ "An impossble happened: " +| e |+ ""

-- | Compile a single view.
compileLorentzView
  :: forall st vt.
     (KnownValue st)
  => CompilationOptions
  -> ContractView st vt
  -> M.SomeView (M.ToT st)
compileLorentzView co (ContractView viewCode)
  | (_ :: Proxy ('ViewTyInfo name arg ret)) <- Proxy @vt =
    let argNotes = getAnnotation @arg NotFollowEntrypoint
        retNotes = getAnnotation @ret NotFollowEntrypoint
        stNotes = getAnnotation @ret NotFollowEntrypoint
        verify = typeVerifyView @(ToT arg) @(ToT st) @(ToT ret) stNotes (convertView viewRaw)
        viewRaw, viewWithNecessaryCast :: M.View (ToT arg) (ToT st) (ToT ret)
        viewRaw = M.View
          { M.vName = demoteViewName @name
          , M.vArgument = argNotes
          , M.vReturn = retNotes
          , M.vCode = compileLorentzWithOptions co viewCode
          }
          \\ niceViewableEvi @arg
          \\ niceViewableEvi @ret
        viewWithNecessaryCast
          | argNotes == M.starNotes || coDisableInitialCast co = viewRaw
          | Right{} <- typeCheckingWith def verify = viewRaw
          | otherwise = viewRaw {
              M.vCode = compileLorentzWithOptions co (I CAST # viewCode :: ViewCode arg st ret) }
    in M.SomeView viewWithNecessaryCast

{- | Set all the contract's views.

@
compileLorentzContract $
  defaultContractData do
    ...
  & setViews
    ( mkView @"myView" @() do
        ...
    , mkView @"anotherView" @Integer do
        ...
    )
@
-}
setViews
  :: forall vd cp st.
     ( RecFromTuple (Rec (ContractView st) (RevealViews vd))
     , NiceViewsDescriptor vd
     )
  => IsoRecTuple (Rec (ContractView st) (RevealViews vd))
  -> ContractData cp st ()
  -> ContractData cp st vd
setViews views = setViewsRec (recFromTuple views)

-- | Version of 'setViews' that accepts a 'Rec'.
--
-- May be useful if you have too many views or want to combine views sets.
setViewsRec
  :: forall vd cp st.
     (NiceViewsDescriptor vd)
  => Rec (ContractView st) (RevealViews vd)
  -> ContractData cp st ()
  -> ContractData cp st vd
setViewsRec views ContractData{..} =
  ContractData{ cdViews = views, .. }

-- | Restrict type of 'Contract', 'ContractData' or other similar type to
-- have no views.
noViews :: contract cp st () -> contract cp st ()
noViews = id

-- | Interpret a Lorentz instruction, for test purposes. Note that this does not run the
-- optimizer.
interpretLorentzInstr
  :: (IsoValuesStack inp, IsoValuesStack out)
  => ContractEnv
  -> inp :-> out
  -> Rec Identity inp
  -> Either MichelsonFailureWithStack (Rec Identity out)
interpretLorentzInstr env (compileLorentz -> instr) inp =
  fromValStack <$> interpretInstr env instr (toValStack inp)

-- | Like 'interpretLorentzInstr', but works on lambda rather than
-- arbitrary instruction.
interpretLorentzLambda
  :: (IsoValue inp, IsoValue out)
  => ContractEnv
  -> Lambda inp out
  -> inp
  -> Either MichelsonFailureWithStack out
interpretLorentzLambda env instr inp = do
  res <- interpretLorentzInstr env instr (Identity inp :& RNil)
  let Identity out :& RNil = res
  return out

instance ContainsDoc (ContractData cp st vd) where
  buildDocUnfinalized =
    buildDocUnfinalized . compileLorentzContract
instance ContainsUpdateableDoc (ContractData cp st vd) where
  modifyDocEntirely how c =
    c{ cdCode = modifyDocEntirely how (cdCode c) }

-- | Lorentz version of analyzer.
analyzeLorentz :: inp :-> out -> AnalyzerRes
analyzeLorentz = analyze . compileLorentz

makeLensesWith postfixLFields ''CompilationOptions

cdCodeL ::
  forall cp st vd cp1. NiceParameterFull cp1 =>
  Lens.Lens (ContractData cp st vd) (ContractData cp1 st vd) (ContractCode cp st) (ContractCode cp1 st)
cdCodeL f (ContractData code views options)
  = fmap (\code' -> ContractData code' views options) (f code)

cdCompilationOptionsL ::
  forall cp st vd.
  Lens.Lens' (ContractData cp st vd) CompilationOptions
cdCompilationOptionsL f (ContractData code views options)
  = fmap (\options' -> ContractData code views options') (f options)