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

-- TODO [#606]: reconsider the implementation

-- | Functions to originate large smart contracts via @octez-client@ and node RPC.
--
-- This is based on a workaround leveraging the lack of gas cost limits on
-- internal transactions produced by @CREATE_CONTRACT@.
--
-- So, in brief, we cannot directly originate a contract that's too large, but
-- we can originate a small "originator" contract, progressively load a packed
-- lambda into it in chunks and finally unpack and execute it, which will
-- run the actual large contract origination.
module Morley.Client.Action.Origination.Large
  ( LargeOriginationData (..)
  , SomeLargeContractOriginator (..)
  , mkLargeOriginationData
  , mkSomeLargeContractOriginator

    -- * Originator contract
  , LargeOriginatorParam
  , LargeOriginatorStore
  , largeContractOriginator

    -- * Origination lambda
  , divideValueInChunks
  , mkOriginationLambda

    -- * Utilities
  , mkLargeOriginatorStore
  , mkLargeOriginatorData
  , mkLargeOriginatorTransactions
  , retrieveLargeContracts
  ) where

import Prelude hiding (concat, drop, swap)

import Data.ByteString.Lazy qualified as LBS

import Lorentz hiding (bytes)
import Morley.Client.Action.Common
import Morley.Client.RPC.Class
import Morley.Client.RPC.Error
import Morley.Client.RPC.Getters (getContractStorage)
import Morley.Client.TezosClient
import Morley.Micheline (fromExpression)
import Morley.Michelson.Interpret.Pack (packValue)
import Morley.Michelson.Parser (notes)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Instr
import Morley.Michelson.Typed.Scope
import Morley.Michelson.Typed.Util (PushableStorageSplit(..), splitPushableStorage)
import Morley.Michelson.Typed.Value
import Morley.Michelson.Untyped.Annotation (annQ, noAnn)
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Util.Constrained

-- | Just a utility type to hold 'SomeLargeContractOriginator' and its large
-- contract 'OriginationData'.
data LargeOriginationData = LargeOriginationData
  { largeOriginator :: SomeLargeContractOriginator
  , largeContractData :: OriginationData
  }

-- | Contains the 'Value heavy' with all the large contract @big_map@s
-- and @ticket@s, the 'largeContractOriginator' for it as well as the lambda
-- to use there.
data SomeLargeContractOriginator where
  SomeLargeContractOriginator
    :: forall heavy. StorageScope heavy
    => Value heavy
    -> T.Contract LargeOriginatorParam (LargeOriginatorStore heavy)
    -> Value (ToT (Lambda (Value heavy) (Address, [Operation])))
    -> SomeLargeContractOriginator

mkLargeOriginationData :: OriginationData -> LargeOriginationData
mkLargeOriginationData largeContractData@OriginationData{..} = LargeOriginationData{..}
  where
    largeOriginator = mkSomeLargeContractOriginator odStorage odContract odBalance

mkSomeLargeContractOriginator
  :: (ParameterScope param, StorageScope store)
  => Value store -- ^ initial storage of the large contract
  -> T.Contract param store -- ^ large contract
  -> Mutez -- ^ balance to tranfer during contract creation
  -> SomeLargeContractOriginator
mkSomeLargeContractOriginator store largeContract xtzs =
  case splitPushableStorage store of
    ConstantStorage val ->
      let origContract = largeContractOriginator
          origLambda = mkOriginationLambda (DROP `Seq` PUSH val) largeContract xtzs
      in SomeLargeContractOriginator VUnit origContract origLambda

    PushableValueStorage instr ->
      let origContract = largeContractOriginator
          origLambda = mkOriginationLambda (DROP `Seq` instr) largeContract xtzs
      in SomeLargeContractOriginator VUnit origContract origLambda

    PartlyPushableStorage val instr ->
      let origContract = largeContractOriginator
          origLambda = mkOriginationLambda instr largeContract xtzs
      in SomeLargeContractOriginator val origContract origLambda


--------------------------------------------------------------------------------
-- Originator contract
--------------------------------------------------------------------------------

-- | Parameter of the originator contract.
type LargeOriginatorParam = 'T.TOr 'T.TBytes 'T.TUnit

-- | Storage of the originator contract.
type LargeOriginatorStore heavy =
  'T.TPair 'T.TAddress ('T.TOr 'T.TAddress ('T.TPair 'T.TBytes heavy))

-- | Large Originator contract.
--
-- Only keeps track of the "owner" address and either
-- - the heavy entries and packed lambda to do the generation (if still loading), or
-- - the resulting address of the originated large contract.
--
-- If the large contract was originated any call will result in a failure containing
-- its address.
-- Any call from an address that's not the "owner" will result in a failure.
largeContractOriginator
  :: StorageScope heavy
  => T.Contract LargeOriginatorParam (LargeOriginatorStore heavy)
largeContractOriginator = T.Contract{..}
  where
    epsNotes = [notes|or (bytes %load_lambda) (unit %run_lambda)|]
    cParamNotes = fromRight T.starParamNotes $ T.mkParamNotes epsNotes noAnn

    stateNotes = T.NTOr noAnn [annQ|originated|] [annQ|loading|] T.starNotes T.starNotes
    cStoreNotes = T.NTPair noAnn [annQ|owner|] noAnn noAnn noAnn T.starNotes stateNotes

    cEntriesOrder = def

    cViews = def
    cCode = T.mkContractCode $
      UNPAIR `Seq`
      -- make checks on the storage
      DIP
        ( UNPAIR `Seq` SWAP `Seq` IF_LEFT
            -- if the large contract has already been originated, fails with its address
            FAILWITH
            -- otherwise, check for the sender
            ( SWAP `Seq` DUP `Seq` SENDER `Seq` COMPARE `Seq` T.EQ `Seq` IF
                ( SWAP `Seq` UNPAIR )
                ( PUSH (VString [mt|sender is not originator owner|]) `Seq` FAILWITH )
            )
        ) `Seq`
      -- stack at this point:
      -- parameter : packed lambda : heavy : owner address : []
      IF_LEFT
        -- if still loading lambda just concat to the exising 'bytes'
        ( CONCAT `Seq` PAIR `Seq` RIGHT `Seq` NIL)
        -- otherwise extract and run the origination lambda
        ( DROP `Seq` UNPACK `Seq`
          IF_NONE
            ( PUSH (VString [mt|failed to unpack lambda|]) `Seq` FAILWITH )
            ( SWAP `Seq` EXEC `Seq` UNPAIR `Seq` LEFT `Seq` SWAP
            )
        ) `Seq`
      -- reconstruct the overall storage
      DIP (SWAP `Seq` PAIR) `Seq`
      -- pair with operations and return
      PAIR

--------------------------------------------------------------------------------
-- Origination lambda
--------------------------------------------------------------------------------

-- | Returns bytes that fit into transaction limits from 'mkOriginationLambda'.
--
-- Note: these have the original order, meaning they should be given to the
-- originator contract from last to first.
divideValueInChunks :: ConstantScope val => Value val -> [ByteString]
divideValueInChunks = divideInChunks . packValue

-- | Returns strict bytes chunks that fit into transaction limits of the input.
divideInChunks :: LByteString -> [ByteString]
divideInChunks bytes
  | LBS.null bytes = []
  | otherwise =
    -- Note: the size is quite a bit below 16k for safety:
    let (chunk, rest) = LBS.splitAt 14000 bytes
    in LBS.toStrict chunk : divideInChunks rest

-- | Generates the lambda to originate a large contract.
mkOriginationLambda
  :: (ParameterScope param, StorageScope store, StorageScope heavy)
  => Instr '[heavy] '[store] -- ^ instruction to recreate the initial storage
  -> T.Contract param store -- ^ large contract
  -> Mutez -- ^ balance to tranfer during contract creation
  -> Value (ToT (Lambda (Value heavy) (Address, [Operation])))
mkOriginationLambda instr largeContract xtzs = mkVLam $ RfNormal $
  instr `Seq` PUSH (VMutez xtzs) `Seq` NONE `Seq`
  CREATE_CONTRACT largeContract `Seq`
  NIL `Seq` SWAP `Seq` CONS `Seq` SWAP `Seq` PAIR


--------------------------------------------------------------------------------
-- Utilities
--------------------------------------------------------------------------------

-- | Helper to create a 'LargeOriginatorStore' 'Value'.
mkLargeOriginatorStore
  :: StorageScope heavy
  => Value heavy
  -> ImplicitAddress
  -> Value (LargeOriginatorStore heavy)
mkLargeOriginatorStore heavyVal owner =
  let vAddr = toVal (MkAddress owner) in
  VPair (vAddr, VOr $ Right $ VPair (VBytes mempty, heavyVal))

-- | Makes 'OriginationData' of the 'largeContractOriginator' that will generate
-- the large contract of the given 'OriginationData' for the sender
-- t'ImplicitAddress'.
mkLargeOriginatorData
  :: ImplicitAddress
  -> LargeOriginationData
  -> OriginationData
mkLargeOriginatorData sender' LargeOriginationData{..} = case largeOriginator of
  SomeLargeContractOriginator heavyVal origContract _origLambda -> OriginationData
    { odAliasBehavior = odAliasBehavior largeContractData
    , odName = ContractAlias $ "largeOriginator." <> unAlias (odName largeContractData)
    , odBalance = zeroMutez
    -- Note ^ we don't transfer any balance here, we instead do it as part of the
    -- last transaction (where it will be transferred to the large contract)
    , odContract = origContract
    , odStorage = mkLargeOriginatorStore heavyVal sender'
    , odDelegate = odDelegate largeContractData
    , odMbFee = odMbFee largeContractData
    }

-- | Makes all the 'TransactionData' to feed the origination lambda into a
-- 'largeContractOriginator' from the t'ContractAddress' of the latter.
mkLargeOriginatorTransactions
  :: ContractAddress -> LargeOriginationData
  -> [TransactionData]
mkLargeOriginatorTransactions originatorAddr LargeOriginationData{..} =
  case largeContractData of
    OriginationData{..} -> case largeOriginator of
      SomeLargeContractOriginator _ _ origLambda ->
        let lambdaChunks = divideValueInChunks origLambda
            doRunLambda = TransactionData @'T.TUnit $ TD
              { tdReceiver = Constrained originatorAddr
              , tdAmount   = odBalance
              , tdEpName   = eprName $ Call @"run_lambda"
              , tdParam    = VUnit
              , tdMbFee    = odMbFee
              }
            mkLoadLambda bytes = TransactionData @'T.TBytes $ TD
              { tdReceiver = Constrained originatorAddr
              , tdAmount   = zeroMutez
              , tdEpName   = eprName $ Call @"load_lambda"
              , tdParam    = VBytes bytes
              , tdMbFee    = odMbFee
              }
        in foldl' (\lst bytes -> mkLoadLambda bytes : lst) [doRunLambda] lambdaChunks

-- | Fetches back the t'ContractAddress' of the large contract generated by a
-- completed 'largeContractOriginator' process.
--
-- It also uses the large contract 'OriginationData' to associate it to the
-- expected alias.
retrieveLargeContracts
  :: (HasTezosRpc m, HasTezosClient m)
  => ContractAddress -> OriginationData -> m ContractAddress
retrieveLargeContracts originatorAddr OriginationData{..} = do
  expr <- getContractStorage originatorAddr
  -- note: for simplicity here we convert the "wrong" value
  -- because we cannot convert from a value with a big_map in its *type*
  -- and also something went wrong if this has a big_map or ticket in its *value*
  let completedStore = fromExpression @(Value (LargeOriginatorStore 'T.TUnit)) expr
  case completedStore of
    Right (VPair (_, VOr (Left largeVAddr))) ->
      withConstrained (fromVal @Address largeVAddr) \case
        largeAddr@ContractAddress{} -> do
          rememberContract odAliasBehavior largeAddr odName
          pure largeAddr
        _ -> error "impossible"
    _ -> throwM RpcOriginatedNoContracts