module EVM.Transaction where

import EVM (initialContract, ceilDiv)
import EVM.FeeSchedule
import EVM.RLP
import EVM.Types
import EVM.Format (hexText)
import EVM.Sign
import qualified EVM.Expr as Expr

import Optics.Core hiding (cons)

import Data.Aeson (FromJSON (..))
import Data.Aeson qualified as JSON
import Data.Aeson.Types qualified as JSON
import Data.ByteString (ByteString, cons)
import Data.ByteString qualified as BS
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (fromMaybe, isNothing, fromJust)
import Data.Word (Word64)
import GHC.Generics (Generic)
import Numeric (showHex)
import Witch (into, unsafeInto)

data AccessListEntry = AccessListEntry {
  address :: Addr,
  storageKeys :: [W256]
} deriving (Show, Generic)

instance JSON.ToJSON AccessListEntry

data TxType = LegacyTransaction
            | AccessListTransaction
            | EIP1559Transaction
  deriving (Show, Eq, Generic)

instance JSON.ToJSON TxType where
  toJSON t = case t of
               EIP1559Transaction    -> "0x2" -- EIP1559
               LegacyTransaction     -> "0x1" -- EIP2718
               AccessListTransaction -> "0x1" -- EIP2930


data Transaction = Transaction {
    txdata            :: ByteString,
    gasLimit          :: Word64,
    gasPrice          :: Maybe W256,
    nonce             :: W256,
    r                 :: W256,
    s                 :: W256,
    toAddr            :: Maybe Addr,
    v                 :: W256,
    value             :: W256,
    txtype            :: TxType,
    accessList        :: [AccessListEntry],
    maxPriorityFeeGas :: Maybe W256,
    maxFeePerGas      :: Maybe W256,
    chainId           :: W256
} deriving (Show, Generic)

instance JSON.ToJSON Transaction where
  toJSON t = JSON.object [ ("input",             (JSON.toJSON (ByteStringS t.txdata)))
                         , ("gas",               (JSON.toJSON $ "0x" ++ showHex (into @Integer $ t.gasLimit) ""))
                         , ("gasPrice",          (JSON.toJSON $ show $ fromJust $ t.gasPrice))
                         , ("v",                 (JSON.toJSON $ show $ (t.v)-27))
                         , ("r",                 (JSON.toJSON $ show $ t.r))
                         , ("s",                 (JSON.toJSON $ show $ t.s))
                         , ("to",                (JSON.toJSON $ t.toAddr))
                         , ("nonce",             (JSON.toJSON $ show $ t.nonce))
                         , ("value",             (JSON.toJSON $ show $ t.value))
                         , ("type",              (JSON.toJSON $ t.txtype))
                         , ("accessList",        (JSON.toJSON $ t.accessList))
                         , ("maxPriorityFeePerGas", (JSON.toJSON $ show $ fromJust $ t.maxPriorityFeeGas))
                         , ("maxFeePerGas",      (JSON.toJSON $ show $ fromJust $ t.maxFeePerGas))
                         , ("chainId",           (JSON.toJSON $ show t.chainId))
                         ]

emptyTransaction :: Transaction
emptyTransaction = Transaction { txdata = mempty
                               , gasLimit = 0
                               , gasPrice = Nothing
                               , nonce = 0
                               , r = 0
                               , s = 0
                               , toAddr = Nothing
                               , v = 0
                               , value = 0
                               , txtype = EIP1559Transaction
                               , accessList = []
                               , maxPriorityFeeGas = Nothing
                               , maxFeePerGas = Nothing
                               , chainId = 1
                               }

-- | utility function for getting a more useful representation of accesslistentries
-- duplicates only matter for gas computation
txAccessMap :: Transaction -> Map Addr [W256]
txAccessMap tx = ((Map.fromListWith (++)) . makeTups) tx.accessList
  where makeTups = map (\ale -> (ale.address , ale.storageKeys ))

-- Given Transaction, it recovers the address that sent it
sender :: Transaction -> Maybe Addr
sender tx = ecrec v' tx.r  tx.s hash
  where hash = keccak' (signingData tx)
        v    = tx.v
        v'   = if v == 27 || v == 28 then v
               else 27 + v

sign :: Integer -> Transaction -> Transaction
sign sk tx = tx { v = into v, r = r, s = s}
  where
    hash = keccak' $ signingData tx
    (v, r, s) = EVM.Sign.sign hash sk

signingData :: Transaction -> ByteString
signingData tx =
  case tx.txtype of
    LegacyTransaction -> if v == (tx.chainId * 2 + 35) || v == (tx.chainId * 2 + 36)
      then eip155Data
      else normalData
    AccessListTransaction -> eip2930Data
    EIP1559Transaction -> eip1559Data
  where v          = tx.v
        to'        = case tx.toAddr of
          Just a  -> BS $ word160Bytes a
          Nothing -> BS mempty
        maxFee = fromJust tx.maxFeePerGas
        maxPrio = fromJust tx.maxPriorityFeeGas
        gasPrice = fromJust tx.gasPrice
        accessList = tx.accessList
        rlpAccessList = EVM.RLP.List $ map (\accessEntry ->
          EVM.RLP.List [BS $ word160Bytes accessEntry.address,
                        EVM.RLP.List $ map rlpWordFull accessEntry.storageKeys]
          ) accessList
        normalData = rlpList [rlpWord256 tx.nonce,
                              rlpWord256 gasPrice,
                              rlpWord256 (into tx.gasLimit),
                              to',
                              rlpWord256 tx.value,
                              BS tx.txdata]
        eip155Data = rlpList [rlpWord256 tx.nonce,
                              rlpWord256 gasPrice,
                              rlpWord256 (into tx.gasLimit),
                              to',
                              rlpWord256 tx.value,
                              BS tx.txdata,
                              rlpWord256 tx.chainId,
                              rlpWord256 0x0,
                              rlpWord256 0x0]
        eip1559Data = cons 0x02 $ rlpList [
          rlpWord256 tx.chainId,
          rlpWord256 tx.nonce,
          rlpWord256 maxPrio,
          rlpWord256 maxFee,
          rlpWord256 (into tx.gasLimit),
          to',
          rlpWord256 tx.value,
          BS tx.txdata,
          rlpAccessList]

        eip2930Data = cons 0x01 $ rlpList [
          rlpWord256 tx.chainId,
          rlpWord256 tx.nonce,
          rlpWord256 gasPrice,
          rlpWord256 (into tx.gasLimit),
          to',
          rlpWord256 tx.value,
          BS tx.txdata,
          rlpAccessList]

accessListPrice :: FeeSchedule Word64 -> [AccessListEntry] -> Word64
accessListPrice fs al =
    sum (map
      (\ale ->
        fs.g_access_list_address  +
        (fs.g_access_list_storage_key  * (unsafeInto . length) ale.storageKeys))
        al)

txGasCost :: FeeSchedule Word64 -> Transaction -> Word64
txGasCost fs tx =
  let calldata     = tx.txdata
      zeroBytes    = BS.count 0 calldata
      nonZeroBytes = BS.length calldata - zeroBytes
      baseCost     = fs.g_transaction
        + (if isNothing tx.toAddr then fs.g_txcreate + initcodeCost else 0)
        + (accessListPrice fs tx.accessList )
      zeroCost     = fs.g_txdatazero
      nonZeroCost  = fs.g_txdatanonzero
      initcodeCost = fs.g_initcodeword * unsafeInto (ceilDiv (BS.length calldata) 32)
  in baseCost + zeroCost * (unsafeInto zeroBytes) + nonZeroCost * (unsafeInto nonZeroBytes)

instance FromJSON AccessListEntry where
  parseJSON (JSON.Object val) = do
    accessAddress_ <- addrField val "address"
    accessStorageKeys_ <- (val JSON..: "storageKeys") >>= parseJSONList
    pure $ AccessListEntry accessAddress_ accessStorageKeys_
  parseJSON invalid =
    JSON.typeMismatch "AccessListEntry" invalid

instance FromJSON Transaction where
  parseJSON (JSON.Object val) = do
    tdata    <- hexText <$> (val JSON..: "data")
    gasLimit <- word64Field val "gasLimit"
    gasPrice <- fmap read <$> val JSON..:? "gasPrice"
    maxPrio  <- fmap read <$> val JSON..:? "maxPriorityFeePerGas"
    maxFee   <- fmap read <$> val JSON..:? "maxFeePerGas"
    nonce    <- wordField val "nonce"
    r        <- wordField val "r"
    s        <- wordField val "s"
    toAddr   <- addrFieldMaybe val "to"
    v        <- wordField val "v"
    value    <- wordField val "value"
    txType   <- fmap (read :: String -> Int) <$> (val JSON..:? "type")
    case txType of
      Just 0x00 -> pure $ Transaction tdata gasLimit gasPrice nonce r s toAddr v value LegacyTransaction [] Nothing Nothing 1
      Just 0x01 -> do
        accessListEntries <- (val JSON..: "accessList") >>= parseJSONList
        pure $ Transaction tdata gasLimit gasPrice nonce r s toAddr v value AccessListTransaction accessListEntries Nothing Nothing 1
      Just 0x02 -> do
        accessListEntries <- (val JSON..: "accessList") >>= parseJSONList
        pure $ Transaction tdata gasLimit gasPrice nonce r s toAddr v value EIP1559Transaction accessListEntries maxPrio maxFee 1
      Just _ -> fail "unrecognized custom transaction type"
      Nothing -> pure $ Transaction tdata gasLimit gasPrice nonce r s toAddr v value LegacyTransaction [] Nothing Nothing 1
  parseJSON invalid =
    JSON.typeMismatch "Transaction" invalid

accountAt :: Expr EAddr -> Getter (Map (Expr EAddr) Contract) Contract
accountAt a = (at a) % (to $ fromMaybe newAccount)

touchAccount :: Expr EAddr -> Map (Expr EAddr) Contract -> Map (Expr EAddr) Contract
touchAccount a = Map.insertWith (flip const) a newAccount

newAccount :: Contract
newAccount = initialContract (RuntimeCode (ConcreteRuntimeCode ""))

-- | Increments origin nonce and pays gas deposit
setupTx :: Expr EAddr -> Expr EAddr -> W256 -> Word64 -> Map (Expr EAddr) Contract -> Map (Expr EAddr) Contract
setupTx origin coinbase gasPrice gasLimit prestate =
  let gasCost = gasPrice * (into gasLimit)
  in (Map.adjust ((over #nonce   (fmap ((+) 1)))
               . (over #balance (`Expr.sub` (Lit gasCost)))) origin)
    . touchAccount origin
    . touchAccount coinbase $ prestate

-- | Given a valid tx loaded into the vm state,
-- subtract gas payment from the origin, increment the nonce
-- and pay receiving address
initTx :: VM t s -> VM t s
initTx vm = let
    toAddr   = vm.state.contract
    origin   = vm.tx.origin
    gasPrice = vm.tx.gasprice
    gasLimit = vm.tx.gaslimit
    coinbase = vm.block.coinbase
    value    = vm.state.callvalue
    toContract = initialContract vm.state.code
    preState = setupTx origin coinbase gasPrice gasLimit vm.env.contracts
    oldBalance = view (accountAt toAddr % #balance) preState
    creation = vm.tx.isCreate
    initState =
        ((Map.adjust (over #balance (`Expr.sub` value))) origin)
      . (Map.adjust (over #balance (Expr.add value))) toAddr
      . (if creation
         then Map.insert toAddr (toContract & (set #balance oldBalance))
         else touchAccount toAddr)
      $ preState
    in
      vm & #env % #contracts .~ initState
         & #tx % #txReversion .~ preState