-- 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
  , 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.Typed (Instr(..), IsoValue, IsoValuesStack(..), ToTs)
import Morley.Michelson.Typed qualified as M
import Morley.Michelson.Typed.Contract (giveNotInView)
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'.
  }

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

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

-- | 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)
  => (IsNotInView => '[(cp, st)] :-> ContractOut st) -> Contract cp st ()
defaultContract code =
  compileLorentzContract $ defaultContractData code

-- | 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)
  => (IsNotInView => '[(cp, st)] :-> ContractOut st) -> ContractData cp st ()
defaultContractData code = ContractData
  { cdCode = mkContractCode 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.
-- However, compilation with 'defaultCompilationOptions' should be valid.
compileLorentzContract
  :: forall cp st vd. ContractData cp st vd -> Contract cp st vd
compileLorentzContract ContractData{..} = Contract{..}
  where
    cMichelsonContract = M.Contract
      { cCode = M.ContractCode $
          compileLorentzWithOptions cdCompilationOptions $ unContractCode cdCode
      , cParamNotes = cpNotes
      , cStoreNotes = getAnnotation @st NotFollowEntrypoint
      , cEntriesOrder = U.canonicalEntriesOrder
      , cViews = compileLorentzViews cdCompilationOptions cdViews
      } \\ niceParameterEvi @cp
        \\ niceStorageEvi @st

    cDocumentedCode = case cdCode of
      ContractCode x -> ContractCode $
        finalizeParamCallingDoc' (Proxy @cp) x #
        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 impossible 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 =
      M.SomeView M.View
        { M.vName = demoteViewName @name
        , M.vArgument = getAnnotation @arg NotFollowEntrypoint
        , M.vReturn = getAnnotation @ret NotFollowEntrypoint
        , M.vCode = compileLorentzWithOptions co viewCode
        }
        \\ niceViewableEvi @arg
        \\ niceViewableEvi @ret

{- | 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
  -> (IsNotInView => inp :-> out)
  -> Rec Identity inp
  -> Either MichelsonFailureWithStack (Rec Identity out)
interpretLorentzInstr env instr inp =
  fromValStack <$> interpretInstr env (compileLorentz $ giveNotInView instr) (toValStack inp)

-- | Like 'interpretLorentzInstr', but works on singleton input and output
-- stacks.
interpretLorentzLambda
  :: (IsoValue inp, IsoValue out)
  => ContractEnv
  -> (IsNotInView => Fn 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 = case cdCode c of
        ContractCode x -> ContractCode $ modifyDocEntirely how x }

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