-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Global blockchain state (emulated).

module Michelson.Runtime.GState
  (
    -- * Auxiliary types
    ContractState (..)
  , AddressState (..)
  , asBalance

  -- * GState
  , GState (..)
  , gsChainIdL
  , gsAddressesL
  , gsCounterL
  , genesisAddresses
  , genesisKeyHashes
  , genesisAddress
  -- * More genesisAddresses which can be used in tests
  , genesisAddress1
  , genesisAddress2
  , genesisAddress3
  , genesisAddress4
  , genesisAddress5
  , genesisAddress6
  , genesisKeyHash
  -- * Genesis secret keys
  , genesisSecretKey
  , genesisSecrets

  , initGState
  , readGState
  , writeGState

  -- * Operations on GState
  , GStateUpdate (..)
  , GStateUpdateError (..)
  , applyUpdate
  , applyUpdates
  , extractAllContracts
  ) where

import Control.Lens (at)
import Data.Aeson (FromJSON(..), ToJSON(..), object, withObject, (.:), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Encode.Pretty as Aeson
import Data.Aeson.TH (deriveJSON)
import qualified Data.ByteString.Lazy as LBS
import Data.List.NonEmpty ((!!))
import qualified Data.Map.Strict as Map
import Data.Typeable ((:~:)(..), eqT)
import Fmt (Buildable(build), (+|), (|+), (||+))
import System.IO.Error (IOError, isDoesNotExistError)

import Michelson.TypeCheck
  (SomeContractAndStorage(..), SomeParamType(..), TcOriginatedContracts,
  typeCheckContractAndStorage)
import qualified Michelson.Typed as T
import Michelson.Typed.Scope
import Michelson.Untyped (Contract, Value)
import Tezos.Address (Address(..), ContractHash, GlobalCounter(..))
import Tezos.Core (ChainId, Mutez, divModMutezInt, dummyChainId)
import Tezos.Crypto
import Util.Aeson
import Util.Lens

-- | State of a contract with code.
data ContractState =
  forall cp st. (ParameterScope cp, StorageScope st) => ContractState
  { csBalance :: Mutez
  -- ^ Amount of mutez owned by this contract.
  , csContract :: T.Contract cp st
  -- ^ Contract itself.
  , csStorage :: T.Value st
  -- ^ Storage value associated with this contract.
  }

deriving stock instance Show ContractState

instance ToJSON ContractState where
  toJSON ContractState{..} =
    object
    [ "balance" .= csBalance
    , "storage" .= T.untypeValue csStorage
    , "contract" .= T.convertContract csContract
    ]

-- These instance is a bit hacky because it is quite painful to
-- write proper JSON instances for typed `Instr` and `Value` so
-- we typecheck untyped representation instead of parsing.
instance FromJSON ContractState where
  parseJSON =
    withObject "contractstate" $ \o -> do
      (balance :: Mutez) <- o .: "balance"
      (uStorage :: Value) <- o .: "storage"
      (uContract :: Contract) <- o .: "contract"
      case typeCheckContractAndStorage uContract uStorage of
        Right (SomeContractAndStorage contract storage) ->
          pure $ ContractState balance contract storage
        Left err -> fail $ "Unable to parse `ContractState`: " <> (show err)

instance Buildable ContractState where
  build ContractState{..} =
    "Contractstate:\n Balance: " +| csBalance |+
    "\n  Storage: " +| T.untypeValue csStorage ||+
    "\n  Contract: " +| T.convertContract csContract ||+ ""

-- | State of an arbitrary address.
data AddressState
  = ASSimple Mutez
  -- ^ For contracts without code we store only its balance.
  | ASContract ContractState
  -- ^ For contracts with code we store more state represented by
  -- 'ContractState'.
  deriving stock (Show, Generic)

instance Buildable AddressState where
  build =
    \case
      ASSimple balance -> "Balance = " +| balance |+ ""
      ASContract cs -> build cs

deriveJSON morleyAesonOptions ''AddressState

-- | Extract balance from 'AddressState'.
asBalance :: AddressState -> Mutez
asBalance =
  \case
    ASSimple b -> b
    ASContract cs -> csBalance cs

-- | Persistent data passed to Morley contracts which can be updated
-- as result of contract execution.
data GState = GState
  { gsChainId :: ChainId
  -- ^ Identifier of chain.
  , gsAddresses :: Map Address AddressState
  -- ^ All known addresses and their state.
  , gsCounter :: GlobalCounter
  -- ^ Ever increasing operation counter.
  } deriving stock (Show)

makeLensesWith postfixLFields ''GState

deriveJSON morleyAesonOptions ''GState

-- | Number of genesis addresses.
genesisAddressesNum :: Word
genesisAddressesNum = 10

-- | Secrets from which genesis addresses are derived from.
genesisSecrets :: NonEmpty SecretKey
genesisSecrets = do
  i <- 1 :| [2 .. genesisAddressesNum]
  let seed = encodeUtf8 (show i :: Text)
  return $ detSecretKey seed

-- | KeyHash of genesis address.
genesisKeyHashes :: NonEmpty KeyHash
genesisKeyHashes = hashKey . toPublic <$> genesisSecrets

-- | Initially these addresses have a lot of money.
genesisAddresses :: NonEmpty Address
genesisAddresses = KeyAddress <$> genesisKeyHashes

-- | One of genesis key hashes.
genesisKeyHash :: KeyHash
genesisKeyHash = head genesisKeyHashes

-- | One of genesis addresses.
genesisAddress :: Address
genesisAddress = head genesisAddresses

-- | Secret key assotiated with 'genesisAddress'.
genesisSecretKey :: SecretKey
genesisSecretKey = head genesisSecrets

-- | More genesis addresses
--
-- We know size of @genesisAddresses@, so it is safe to use @!!@
genesisAddress1, genesisAddress2, genesisAddress3 :: Address
genesisAddress4, genesisAddress5, genesisAddress6 :: Address
genesisAddress1 = genesisAddresses !! 1
genesisAddress2 = genesisAddresses !! 2
genesisAddress3 = genesisAddresses !! 3
genesisAddress4 = genesisAddresses !! 4
genesisAddress5 = genesisAddresses !! 5
genesisAddress6 = genesisAddresses !! 6

-- | Initial 'GState'. It's supposed to be used if no 'GState' is
-- provided. It puts plenty of money on each genesis address.
initGState :: GState
initGState =
  GState
  { gsChainId = dummyChainId
  , gsAddresses = Map.fromList
    [ (genesis, ASSimple money)
    | let (money, _) = maxBound @Mutez `divModMutezInt` genesisAddressesNum
                    ?: error "Number of genesis addresses is 0"
    , genesis <- toList genesisAddresses
    ]
  , gsCounter = GlobalCounter 0
  }

data GStateParseError =
  GStateParseError String
  deriving stock (Show)

instance Exception GStateParseError where
  displayException (GStateParseError str) = "Failed to parse GState: " <> str

-- | Read 'GState' from a file.
readGState :: FilePath -> IO GState
readGState fp = (LBS.readFile fp >>= parseFile) `catch` onExc
  where
    parseFile :: LByteString -> IO GState
    parseFile lByteString =
      if null lByteString
      then pure initGState
      else (either (throwM . GStateParseError) pure . Aeson.eitherDecode') lByteString
    onExc :: IOError -> IO GState
    onExc exc
      | isDoesNotExistError exc = pure initGState
      | otherwise = throwM exc

-- | Write 'GState' to a file.
writeGState :: FilePath -> GState -> IO ()
writeGState fp gs = LBS.writeFile fp (Aeson.encodePretty' config gs)
  where
    config =
      Aeson.defConfig
      { Aeson.confTrailingNewline = True
      }

-- | Updates that can be applied to 'GState'.
data GStateUpdate where
  GSAddAddress :: Address -> AddressState -> GStateUpdate
  GSSetStorageValue :: StorageScope st => Address -> T.Value st -> GStateUpdate
  GSSetBalance :: Address -> Mutez -> GStateUpdate
  GSIncrementCounter :: GStateUpdate

deriving stock instance Show GStateUpdate

instance Buildable GStateUpdate where
  build =
    \case
      GSAddAddress addr st ->
        "Add address " +| addr |+ " with state " +| st |+ ""
      GSSetStorageValue addr tVal ->
        "Set storage value of address " +| addr |+ " to " +| T.untypeValue tVal |+ ""
      GSSetBalance addr balance ->
        "Set balance of address " +| addr |+ " to " +| balance |+ ""
      GSIncrementCounter ->
        "Increment internal counter after operation"

data GStateUpdateError
  = GStateAddressExists Address
  | GStateUnknownAddress Address
  | GStateNotContract Address
  | GStateStorageNotMatch Address
  deriving stock (Show)

instance Buildable GStateUpdateError where
  build =
    \case
      GStateAddressExists addr -> "Address already exists: " <> build addr
      GStateUnknownAddress addr -> "Unknown address: " <> build addr
      GStateNotContract addr -> "Address doesn't have contract: " <> build addr
      GStateStorageNotMatch addr ->
        "Storage type does not match the contract in run-time state\
        \ when updating new storage value to address: " <> build addr

-- | Apply 'GStateUpdate' to 'GState'.
applyUpdate :: GStateUpdate -> GState -> Either GStateUpdateError GState
applyUpdate =
  \case
    GSAddAddress addr st ->
      maybeToRight (GStateAddressExists addr) . addAddress addr st
    GSSetStorageValue addr newValue ->
      setStorageValue addr newValue
    GSSetBalance addr newBalance -> setBalance addr newBalance
    GSIncrementCounter -> Right . over gsCounterL (+1)

-- | Apply a list of 'GStateUpdate's to 'GState'.
applyUpdates :: [GStateUpdate] -> GState -> Either GStateUpdateError GState
applyUpdates = flip (foldM (flip applyUpdate))

-- | Add an address if it hasn't been added before.
addAddress :: Address -> AddressState -> GState -> Maybe GState
addAddress addr st gs
    | addr `Map.member` accounts = Nothing
    | otherwise = Just (gs {gsAddresses = accounts & at addr .~ Just st})
  where
    accounts = gsAddresses gs

-- | Update storage value associated with given address.
setStorageValue :: forall st. (StorageScope st) =>
     Address -> T.Value st -> GState -> Either GStateUpdateError GState
setStorageValue addr newValue = updateAddressState addr modifier
  where
    modifier :: AddressState -> Either GStateUpdateError AddressState
    modifier (ASSimple _) = Left (GStateNotContract addr)
    modifier (ASContract (ContractState b c (_ :: T.Value st') )) = do
      case (eqT @st @st') of
        Just Refl -> Right $ ASContract $ ContractState b c newValue
        _ -> Left $ GStateStorageNotMatch addr

-- | Update storage value associated with given address.
setBalance :: Address -> Mutez -> GState -> Either GStateUpdateError GState
setBalance addr newBalance = updateAddressState addr (Right . modifier)
  where
    modifier (ASSimple _) = ASSimple newBalance
    modifier (ASContract cs) = ASContract (cs {csBalance = newBalance})

updateAddressState ::
     Address
  -> (AddressState -> Either GStateUpdateError AddressState)
  -> GState
  -> Either GStateUpdateError GState
updateAddressState addr f gs =
  case addresses ^. at addr of
    Nothing -> Left (GStateUnknownAddress addr)
    Just as -> do
      newState <- f as
      return $ gs { gsAddresses = addresses & at addr .~ Just newState }
  where
    addresses = gsAddresses gs

-- | Retrieve all contracts stored in GState
extractAllContracts :: GState -> TcOriginatedContracts
extractAllContracts = Map.fromList . mapMaybe extractContract . toPairs . gsAddresses
 where
    extractContract
      :: (Address, AddressState) -> Maybe (ContractHash, SomeParamType)
    extractContract =
      \case (KeyAddress _, ASSimple {}) -> Nothing
            (KeyAddress _, _) -> error "broken GState"
            (ContractAddress ca, ASContract (ContractState{..})) ->
              Just (ca, SomeParamType sing $ T.cParamNotes $ csContract)
            (ContractAddress _, _) -> error "broken GState"