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

-- | Global blockchain state (emulated).

module Michelson.Runtime.GState
  (
    -- * Auxiliary types
    ContractState (..)
  , getTypedContract
  , getTypedStorage
  , SomeContractAndStorage (..)
  , getTypedContractAndStorage
  , AddressState (..)
  , asBalance

  -- * GState
  , GState (..)
  , gsChainIdL
  , gsAddressesL
  , 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 ((+|), (|+), (||+))
import Formatting.Buildable (Buildable(build))
import System.IO.Error (IOError, isDoesNotExistError)

import Michelson.TypeCheck
  (SomeContract(..), TCError, TcOriginatedContracts, typeCheckContract, typeCheckTopLevelType)
import Michelson.Typed (SomeValue, SomeValue'(..))
import qualified Michelson.Typed as T
import Michelson.Typed.Scope
import Michelson.Untyped (Contract, ParameterType, Value, contractParameter, contractStorage)
import Tezos.Address (Address(..), ContractHash)
import Tezos.Core (ChainId, Mutez, divModMutezInt, dummyChainId)
import Tezos.Crypto
import Util.Aeson
import Util.Lens

-- | State of a contract with code.
data ContractState = ContractState
  { csBalance :: Mutez
  -- ^ Amount of mutez owned by this contract.
  , csStorage :: Value
  -- ^ Storage value associated with this contract.
  , csContract :: Contract
  -- ^ Contract itself (untyped).
  , csTypedContract :: (Maybe SomeContract)
  , csTypedStorage :: (Maybe SomeValue)
  -- ^ We keep typed representation of contract code
  -- and storage in form, that hides their actual type
  -- in order to simplify the rest of the code
  -- (e.g. avoid type parameters for `ContractState` and so on).
  -- They are made optional in order to perform safe parsing
  -- from JSON (we simply return `Nothing` in this parser and use
  -- `getTypedStorage` or `getTypedContract` that optionally typecheck
  -- storage or contract code).
  }

deriving stock instance Show ContractState

instance ToJSON ContractState where
  toJSON ContractState{..} = object
    [ "balance" .= csBalance
    , "storage" .= csStorage
    , "contract" .= 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
    csBalance <- o .: "balance"
    csStorage <- o .: "storage"
    csContract <- o .: "contract"
    let csTypedContract = Nothing
    let csTypedStorage = Nothing
    return ContractState {..}

instance Buildable ContractState where
  build ContractState{..} =
    "Contractstate:\n csBalance: " +| csBalance |+
    "\n  csStorage: " +| csStorage |+
    "\n  csContract: " +| csContract |+
    "\n  csTypedContract: " +| csTypedContract ||+
    "\n  csTypedStorage: " +| csTypedStorage ||+ ""

-- | 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.
  } deriving stock (Show)

makeLensesWith postfixLFields ''GState

deriveJSON morleyAesonOptions ''GState

getTypedContract :: GState -> ContractState -> Either TCError SomeContract
getTypedContract gs ContractState{..} =
  typeCheckContract (extractAllContracts gs) csContract

getTypedStorage :: GState -> ContractState -> Either TCError SomeValue
getTypedStorage gs ContractState{..} =
  typeCheckTopLevelType (extractAllContracts gs) (contractStorage csContract) csStorage

-- [#36] TODO: try to get rid of this type, 'ContractState' should become
-- broader than it
data SomeContractAndStorage =
  forall cp st. (ParameterScope cp, StorageScope st) => SomeContractAndStorage
  { scsContract :: T.Contract cp st
  , scsStorage :: T.Value st
  }

getTypedContractAndStorage
  :: (TCError -> err)
  -> (TCError -> err)
  -> GState
  -> ContractState
  -> Either err SomeContractAndStorage
getTypedContractAndStorage liftContractErr liftStorageErr gs cs = do
  SomeContract (contract@T.Contract{} :: T.Contract cp st) <-
    first liftContractErr $ getTypedContract gs cs
  SomeValue (storage :: T.Value st') <-
    first liftStorageErr $ getTypedStorage gs cs
  Refl <- pure $ eqT @st @st'
              ?: error "Storage type does not match the contract in runtime state"
  return $ SomeContractAndStorage contract storage

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

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
  = GSAddAddress Address AddressState
  | GSSetStorageValue Address Value SomeValue
  | GSSetBalance Address Mutez
  deriving stock (Show)

instance Buildable GStateUpdate where
  build =
    \case
      GSAddAddress addr st ->
        "Add address " +| addr |+ " with state " +| st |+ ""
      GSSetStorageValue addr val _ ->
        "Set storage value of address " +| addr |+ " to " +| val |+ ""
      GSSetBalance addr balance ->
        "Set balance of address " +| addr |+ " to " +| balance |+ ""

data GStateUpdateError
  = GStateAddressExists Address
  | GStateUnknownAddress Address
  | GStateNotContract 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

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

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

-- | Updare storage value associated with given address.
setStorageValue ::
     Address -> Value -> SomeValue -> GState -> Either GStateUpdateError GState
setStorageValue addr newValue newTypedValue = updateAddressState addr modifier
  where
    modifier (ASSimple _) = Left (GStateNotContract addr)
    modifier (ASContract cs) = Right $ ASContract $
      cs { csStorage = newValue
         , csTypedStorage = Just newTypedValue
         }

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

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