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

-- | Michelson contract in untyped model.

module Morley.Michelson.Untyped.Contract
  ( EntriesOrder (..)
  , canonicalEntriesOrder
  , entriesOrderToInt
  , mapEntriesOrdered

  , ContractBlock (..)
  , orderContractBlock

  , Contract' (..)
  , View' (..)
  , Storage
  , mapContractCode
  ) where

import Data.Aeson.TH (deriveJSON)
import Data.Data (Data(..))
import Data.Default (Default(..))
import Fmt (Buildable(build))
import Text.PrettyPrint.Leijen.Text (indent, nest, semi, text, (<$$>), (<+>))

import Morley.Michelson.Printer.Util
  (Prettier(..), RenderDoc(..), assertParensNotNeeded, buildRenderDoc, needsParens, renderOpsList)
import Morley.Michelson.Untyped.Type (ParameterType(..), Ty(..))
import Morley.Michelson.Untyped.View
import Morley.Util.Aeson

-- TODO [#698]: views are yet handled in a very hacky and not fully working way

-- | Top-level entries order of the contract.
-- This is preserved due to the fact that it affects
-- the output of pretty-printing and serializing contract.

-- Each constructors is created by the order of the first letter of
-- @parameter@, @storage@, and @code@.
--
-- For example, @PSC@ would be @parameter@, @storage@ and @code@,
-- @CPS@ would be @code@, @parameter@,and @storage@, and so on.
data EntriesOrder
  = PSC
  | PCS
  | SPC
  | SCP
  | CSP
  | CPS
  deriving stock (Bounded, Data, Enum, Eq, Generic, Show)

instance Default EntriesOrder where
  def = canonicalEntriesOrder
instance NFData EntriesOrder

-- | The canonical entries order which is ordered as follow:
-- @parameter@, @storage@, and @code@.
canonicalEntriesOrder :: EntriesOrder
canonicalEntriesOrder = PSC

-- | @(Int, Int, Int)@ is the positions of @parameter@, @storage@, and @code@
-- respectively.
entriesOrderToInt :: EntriesOrder -> (Int, Int, Int)
entriesOrderToInt = \case
  PSC -> (0, 1, 2)
  PCS -> (0, 2, 1)
  SPC -> (1, 0, 2)
  SCP -> (1, 2, 0)
  CSP -> (2, 1, 0)
  CPS -> (2, 0, 1)

-- | Contract block, convenient when parsing
data ContractBlock op
  = CBParam ParameterType
  | CBStorage Ty
  | CBCode [op]
  | CBView (View' op)
  deriving stock (Eq, Show, Functor)

-- | Construct a contract representation from the contract blocks (i.e. parameters,
-- storage, code blocks, etc.) in arbitrary order.
-- This makes sure that unique blocks like @code@ do not duplicate, and saves the
-- order in the contract so that it can print the contract blocks in the same
-- order it was parsed. TODO [#698]: this is not fully true now.
orderContractBlock :: [ContractBlock op] -> Maybe (Contract' op)
orderContractBlock blocks =
  let
    vs = mapMaybe (\case CBView v -> Just v; _ -> Nothing) blocks
    plain = filter (\case CBView{} -> False; _ -> True) blocks
  in case plain of
  [CBParam   p, CBStorage s, CBCode    c] -> Just $ Contract p s c PSC vs
  [CBParam   p, CBCode    c, CBStorage s] -> Just $ Contract p s c PCS vs
  [CBStorage s, CBParam   p, CBCode    c] -> Just $ Contract p s c SPC vs
  [CBStorage s, CBCode    c, CBParam   p] -> Just $ Contract p s c SCP vs
  [CBCode    c, CBStorage s, CBParam   p] -> Just $ Contract p s c CSP vs
  [CBCode    c, CBParam   p, CBStorage s] -> Just $ Contract p s c CPS vs
  _                                       -> Nothing

instance Buildable (ContractBlock op) where
  build (CBParam{}) = "parameter"
  build (CBStorage{}) = "storage"
  build (CBCode{}) = "code"
  build (CBView{}) = "view"

-- | Map each contract fields by the given function and sort the output
-- based on the 'EntriesOrder'.
mapEntriesOrdered
  :: Contract' op
  -> (ParameterType -> a)
  -> (Storage -> a)
  -> ([op] -> a)
  -> (View' op -> a)
  -> [a]
mapEntriesOrdered Contract{..} fParam fStorage fCode fView =
  mconcat
    [ fmap snd $ sortWith fst
        [ (paramPos, fParam contractParameter)
        , (storagePos, fStorage contractStorage)
        , (codePos, fCode contractCode)
        ]
    , fmap fView contractViews
    ]
  where
    (paramPos, storagePos, codePos) = entriesOrderToInt entriesOrder

-- | Convenience synonym for 'Ty' representing the storage type
type Storage = Ty

-- | General untyped contract representation.
data Contract' op = Contract
  { contractParameter :: ParameterType
    -- ^ Contract parameter type
  , contractStorage :: Storage
    -- ^ Contract storage type
  , contractCode :: [op]
    -- ^ Contract code as a list of operations
  , entriesOrder :: EntriesOrder
    -- ^ Original order of contract blocks, so that we can print them
    -- in the same order they were read
  , contractViews :: [View' op]
    -- ^ Contract views
  } deriving stock (Eq, Show, Functor, Data, Generic)

instance NFData op => NFData (Contract' op)

instance (RenderDoc op) => RenderDoc (Contract' op) where
  renderDoc pn contract =
    assertParensNotNeeded pn
      $ foldr (<$$>) (text "")
      $ mapEntriesOrdered contract
        (\parameter -> "parameter" <+> renderDoc needsParens (Prettier parameter) <> semi)
        (\storage -> "storage" <+> renderDoc needsParens (Prettier storage) <> semi)
        (\code -> "code" <+> nest (length ("code {" :: Text)) (renderOpsList False code <> semi))
        (\View{..} -> "view" <+> renderViewName viewName
           <+> renderDoc needsParens viewArgument
           <+> renderDoc needsParens viewReturn
           <$$> indent 5  -- 5 is forced by Michelson
                 (renderOpsList False viewCode)
           <> semi
          )

instance RenderDoc op => Buildable (Contract' op) where
  build = buildRenderDoc

-- | Map all the instructions appearing in the contract.
mapContractCode :: (op -> op) -> Contract' op -> Contract' op
mapContractCode f (Contract param st code o vs) =
  Contract
    param
    st
    (code <&> f)
    o
    (vs <&> \v -> v{ viewCode = viewCode v <&> f })

deriveJSON morleyAesonOptions ''EntriesOrder
deriveJSON morleyAesonOptions ''Contract'