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

-- | Global blockchain state (emulated).

module Morley.Michelson.Runtime.GState
  (
    -- * Auxiliary types
    ContractState (..)
  , AddressState (..)
  , asBalance
  , VotingPowers (..)
  , vpPick
  , vpTotal
  , mkVotingPowers
  , mkVotingPowersFromMap
  , dummyVotingPowers
  , BigMapCounter(..)
  , bigMapCounter

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

  , initGState
  , readGState
  , writeGState

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

import Control.Lens (at, makeLenses)
import Data.Aeson (FromJSON(..), ToJSON(..), object, withObject, (.:), (.:?), (.=))
import Data.Aeson qualified as Aeson
import Data.Aeson.Encode.Pretty qualified as Aeson
import Data.Aeson.TH (deriveJSON)
import Data.ByteString.Lazy qualified as LBS
import Data.Default (def)
import Data.Map.Strict qualified as Map
import Data.Type.Equality ((:~:)(..))
import Fmt (Buildable(build), pretty, (+|), (|+))
import System.IO.Error (IOError, isDoesNotExistError)

import Morley.Michelson.TypeCheck
  (SomeParamType(..), TcOriginatedContracts, typeCheckContractAndStorage, typeCheckingWith)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Existential (SomeContractAndStorage(..))
import Morley.Michelson.Typed.Scope
import Morley.Michelson.Untyped (Contract, Value)
import Morley.Tezos.Address (Address(..), ContractHash, GlobalCounter(..))
import Morley.Tezos.Core (ChainId, Mutez, divModMutezInt, dummyChainId)
import Morley.Tezos.Crypto
import Morley.Util.Aeson
import Morley.Util.Lens
import Morley.Util.Peano
import Morley.Util.Sing (eqI, eqParamSing, eqParamSing2)
import Morley.Util.SizedList qualified as SL
import Morley.Util.SizedList.Types

-- | 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.
  , csDelegate :: Maybe KeyHash
  -- ^ Delegate associated with this contract.
  }

deriving stock instance Show ContractState

instance Eq ContractState where
  ContractState b1 c1 s1 d1 == ContractState b2 c2 s2 d2 =
    b1 == b2
    && eqParamSing2 c1 c2
    && eqParamSing s1 s2
    && d1 == d2

instance ToJSON ContractState where
  toJSON ContractState{..} =
    object .
    maybe id ((:) . ("delegate" .=)) csDelegate $
    [ "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"
      (delegate :: Maybe KeyHash) <- o .:? "delegate"
      case typeCheckingWith def $ typeCheckContractAndStorage uContract uStorage of
        Right (SomeContractAndStorage contract storage) ->
          pure $ ContractState balance contract storage delegate
        Left err -> fail $ "Unable to parse `ContractState`: " <> pretty err

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

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

-- | Distribution of voting power among the contracts.
--
-- Voting power reflects the ability of bakers to accept, deny or pass new
-- proposals for protocol updates. I.e. each baker has its vote weight.
--
-- This datatype definition may change in future, so its internals are not
-- exported.
newtype VotingPowers
  = VotingPowers (Map KeyHash Natural)
  deriving stock (Show, Eq)

deriveJSON morleyAesonOptions ''VotingPowers

-- | Get voting power of the given address.
vpPick :: KeyHash -> VotingPowers -> Natural
vpPick key (VotingPowers distr) = Map.lookup key distr ?: 0

-- | Get total voting power.
vpTotal :: VotingPowers -> Natural
vpTotal (VotingPowers distr) = sum distr

-- | Create voting power distribution from map.
mkVotingPowersFromMap :: Map KeyHash Natural -> VotingPowers
mkVotingPowersFromMap = VotingPowers

-- | Create voting power distribution.
--
-- If some key is encountered multiple times, voting power will be summed up.
mkVotingPowers :: [(KeyHash, Natural)] -> VotingPowers
mkVotingPowers = mkVotingPowersFromMap . Map.fromListWith (+)

-- | All big_maps stored in a chain have a globally unique ID.
--
-- We use this counter to keep track of how many big_maps have been created so far,
-- and to generate new IDs whenever a new big_map is created.
newtype BigMapCounter = BigMapCounter { _bigMapCounter :: Natural }
  deriving stock (Show, Eq, Generic)
  deriving anyclass (NFData)
  deriving newtype (ToJSON, FromJSON, Num, Buildable)

makeLenses ''BigMapCounter

-- | 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.
  , gsVotingPowers :: VotingPowers
  -- ^ Voting power distribution.
  , gsCounter :: GlobalCounter
  -- ^ Ever increasing operation counter.
  , gsBigMapCounter :: BigMapCounter
  } deriving stock (Show, Eq)

makeLensesWith postfixLFields ''GState

deriveJSON morleyAesonOptions ''GState

-- | Number of genesis addresses, type-level
type GenesisAddressesNum = 10

-- | Convenience synonym
type GenesisList a = SizedList GenesisAddressesNum a

-- | Number of genesis addresses, term-level
genesisAddressesNum :: Natural
genesisAddressesNum = natVal @GenesisAddressesNum Proxy

-- | Secrets from which genesis addresses are derived from.
genesisSecrets :: GenesisList SecretKey
genesisSecrets = SL.generate @GenesisAddressesNum $ \i ->
    let seed = encodeUtf8 (show (i + 1) :: Text)
    in detSecretKey seed

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

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

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

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

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

-- | More genesis addresses
genesisAddress1, genesisAddress2, genesisAddress3 :: Address
genesisAddress4, genesisAddress5, genesisAddress6 :: Address
_ :< genesisAddress1 :< genesisAddress2 :< genesisAddress3
  :< genesisAddress4 :< genesisAddress5 :< genesisAddress6
  :< _ = genesisAddresses

{-# DEPRECATED genesisAddress4, genesisAddress5, genesisAddress6
               "Consider using 'genesisAddressN' instead" #-}

-- | More genesis addresses, via a type-level natural
--
-- > genesisAddressN @7
--
-- Note that @'genesisAddress' == genesisAddressN \@0@, @'genesisAddress1' == genesisAddressN \@1@,
-- etc.
genesisAddressN :: forall n. (SingIPeano n, ToPeano GenesisAddressesNum > ToPeano n ~ 'True) => Address
genesisAddressN = SL.index @n genesisAddresses

-- | Dummy 'VotingPowers'. We give all the voting power to two genesis addreses
-- as the addresses holding lot of money. Only two addresses are involved for
-- simplicity.
dummyVotingPowers :: VotingPowers
dummyVotingPowers = case genesisKeyHashes of
  k1 :< k2 :< _ -> mkVotingPowers [(k1, 50), (k2, 50)]

-- | 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
    ]
  , gsVotingPowers = dummyVotingPowers
  , gsCounter = GlobalCounter 0
  , gsBigMapCounter = BigMapCounter 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
  GSUpdateCounter :: GlobalCounter -> GStateUpdate
  GSSetBigMapCounter :: BigMapCounter -> GStateUpdate
  GSSetDelegate :: Address -> Maybe KeyHash -> 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"
      GSUpdateCounter v ->
        "Set internal counter to " +| v |+ " after interpreting " <>
        "several 'CREATE_CONTRACT' instructions"
      GSSetBigMapCounter inc ->
        "Increment internal big_map counter by: " +| build inc
      GSSetDelegate addr key ->
        "Set delegate for contract " +| addr |+ " to " +| maybe "<nobody>" build key

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)
    GSUpdateCounter newCounter -> Right . set gsCounterL newCounter
    GSSetBigMapCounter bmCounter -> Right . set gsBigMapCounterL bmCounter
    GSSetDelegate addr key -> setDelegate addr key

-- | 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{csStorage = _ :: T.Value st', ..}) = do
      case eqI @st @st' of
        Just Refl -> Right $ ASContract $ ContractState{csStorage = newValue, ..}
        _ -> Left $ GStateStorageNotMatch addr

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

-- | Set delegate for a given address
setDelegate :: Address -> Maybe KeyHash -> GState -> Either GStateUpdateError GState
setDelegate addr key = updateAddressState addr \case
  ASSimple _ -> Left $ GStateNotContract addr
  ASContract cs -> Right $ ASContract cs{csDelegate = key}

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 $ T.cParamNotes $ csContract)
            (ContractAddress _, _) -> error "broken GState"