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

-- | Michelson contract in untyped model.

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

  , ContractBlock (..)
  , ContractBlockError (..)
  , orderContractBlock

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

import Control.Lens (Prism', makePrisms)
import Data.Aeson (FromJSON, FromJSONKey(..), ToJSON, ToJSONKey(..))
import Data.Aeson.TH (deriveJSON)
import Data.Aeson.Types qualified as AesonTypes
import Data.Bitraversable (bitraverse)
import Data.Data (Data(..))
import Data.Default (Default(..))
import Data.Map qualified as Map
import Data.Text (stripPrefix)
import Fmt (Buildable(build), listF, nameF, pretty, (<$$>), (<+>))
import Fmt.Operators qualified as PP
import Prettyprinter (semi)
import Prettyprinter qualified as PP

import Morley.Michelson.Printer.Util
import Morley.Michelson.Untyped.Type (ParameterType(..), Ty(..))
import Morley.Michelson.Untyped.View
import Morley.Util.Aeson

data Entry
  = EntryParameter
  | EntryStorage
  | EntryCode
  | EntryView ViewName
  deriving stock (Eq, Ord, Show, Data, Generic)

instance Buildable Entry where
  build = \case
    EntryParameter -> "parameter"
    EntryStorage -> "storage"
    EntryCode -> "code"
    EntryView name -> nameF "view" $ build name

deriveJSON morleyAesonOptions ''Entry

instance ToJSONKey Entry where
  toJSONKey = AesonTypes.toJSONKeyText $ \case
    EntryParameter -> "parameter"
    EntryStorage -> "storage"
    EntryCode -> "code"
    EntryView name -> "view:" <> unViewName name

instance FromJSONKey Entry where
  fromJSONKey = AesonTypes.FromJSONKeyTextParser $ \case
    "parameter" -> pure EntryParameter
    "storage" -> pure EntryStorage
    "code" -> pure EntryCode
    x | Just name <- stripPrefix "view:" x
      -> either (fail . pretty) (pure . EntryView) $ mkViewName name
    _ -> fail $ "Unexpected Entry value"

instance NFData Entry

-- | 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.
newtype EntriesOrder = EntriesOrder { unEntriesOrder :: Map Entry Word }
  deriving stock (Eq, Ord, Show, Data)
  deriving newtype (NFData, FromJSON, ToJSON)

-- | Helper to construct 'EntriesOrder' from an ordered list of entires.
-- Duplicate entires are ignored.
mkEntriesOrder :: [Entry] -> EntriesOrder
mkEntriesOrder = EntriesOrder . Map.fromListWith const . (`zip` [0..])

instance Default EntriesOrder where
  def = canonicalEntriesOrder

-- | The canonical entries order which is ordered as follow:
-- @parameter@, @storage@, and @code@.
canonicalEntriesOrder :: EntriesOrder
canonicalEntriesOrder = mkEntriesOrder [EntryParameter, EntryStorage, EntryCode]

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

makePrisms ''ContractBlock
makePrisms ''Entry

data ContractBlockError
  = CBEDuplicate (NonEmpty Entry)
  | CBEMissing Entry
  deriving stock (Eq, Show)

instance Buildable ContractBlockError where
  build = \case
    CBEDuplicate (b :| []) -> nameF "Duplicate contract field" $ build b
    CBEDuplicate bs -> nameF "Duplicate contract fields" $ listF bs
    CBEMissing entry -> nameF "Missing contract field" $ build entry

newtype Validation e a = Validation { unValidation :: Either e a }
  deriving stock (Show, Eq)
  deriving newtype (Functor)

instance Semigroup e => Applicative (Validation e) where
  pure = Validation . Right
  Validation (Left e1) <*> Validation r = Validation $ Left $ either (e1 <>) (const e1) r
  Validation (Right f) <*> x = f <$> x

-- | 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.
orderContractBlock
  :: forall op. [ContractBlock op]
  -> Either (NonEmpty ContractBlockError) (Contract' op)
orderContractBlock blocks =
  let
    blockToBlockType = \case
      CBParam{} -> EntryParameter
      CBStorage{} -> EntryStorage
      CBCode{} -> EntryCode
      CBView View{..} -> EntryView viewName
    duplicates :: [Entry]
    blockMap :: Map Entry (Word, ContractBlock op)
    (duplicates, blockMap) =
      foldr go (mempty, mempty) (zip blocks [0..])
      where
        go (b, n) (!dups, !seen) =
          let blockType = blockToBlockType b
              dups' | Map.member blockType seen = blockType:dups
                    | otherwise = dups
          in (dups', Map.insert blockType (n, b) seen)
    entriesOrder = EntriesOrder $ fst <$> blockMap
    getBlock :: Entry -> Prism' (ContractBlock op) a -> Validation (NonEmpty ContractBlockError) a
    getBlock ty prism = Validation $ maybeToRight (one $ CBEMissing ty) $
      Map.lookup ty blockMap >>= preview (_2 . prism)
    contractViews = ViewsSet
      . Map.fromDistinctAscList -- we know that keys are still unique and sorted
      . mapMaybe (bitraverse (preview _EntryView) (preview $ _2 . _CBView))
      $ toPairs blockMap
  in unValidation do
    Validation $ nonEmpty duplicates & maybeToLeft () . fmap (one . CBEDuplicate)
    contractParameter <- getBlock EntryParameter _CBParam
    contractStorage <- getBlock EntryStorage _CBStorage
    contractCode <- getBlock EntryCode _CBCode
    pure Contract{..}

instance Buildable (ContractBlock op) where
  build CBParam{} = "parameter"
  build CBStorage{} = "storage"
  build CBCode{} = "code"
  build (CBView View{..}) = "view \"" <> build viewName <> "\""

-- | 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 = snd <$> sortWith fst elements
  where
    getElemOrder ty = fromMaybe maxBound $ Map.lookup ty $ unEntriesOrder entriesOrder
    elements
      = first getElemOrder
      <$> [ (EntryParameter, fParam contractParameter)
          , (EntryStorage, fStorage contractStorage)
          , (EntryCode, fCode contractCode)]
      <>  (toList contractViews <&> \v@View{..} -> (EntryView viewName, fView v))

-- | 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 :: ViewsSet 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 (<$$>) mempty
      $ mapEntriesOrdered contract
        (\parameter -> "parameter" <+> renderDoc needsParens parameter <> semi)
        (\storage -> "storage" <+> renderDoc needsParens storage <> semi)
        (\code -> "code" <+> PP.align (renderDoc doesntNeedParens code <> semi))
        (\View{..} -> PP.group $ "view" <+> PP.align (
          PP.sep
            [ renderViewName viewName
            , renderDoc needsParens viewArgument
            , renderDoc needsParens viewReturn
            ]
          PP.<$> renderDoc doesntNeedParens viewCode <> semi
        ))

deriveJSON morleyAesonOptions ''Contract'