{-# LANGUAGE DeriveGeneric #-}

module SymbolicState where

import SymbolicExpression
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Word
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import GHC.Generics (Generic)

newtype SymStack = SymStack [SymExpr]
  deriving (Show, Eq)

newtype SymMemory = SymMemory (Map Integer SymExpr)
  deriving (Show, Eq)

newtype SymStorage = SymStorage (Map SymExpr SymExpr)
  deriving (Show, Eq)

data ExecState = ExecState
  { stack       :: SymStack
  , memory      :: SymMemory
  , storage     :: SymStorage
  , pc          :: Int
  , gas         :: Integer
  , constraints :: [Constraint]
  , calldata    :: ByteString
  , returnData  :: Maybe SymExpr
  , reverted    :: Bool
  , halted      :: Bool
  , symbolicId  :: Int
  } deriving (Show)

initState :: ByteString -> ExecState
initState cd = ExecState
  { stack = emptyStack
  , memory = emptyMemory
  , storage = emptyStorage
  , pc = 0
  , gas = 1000000
  , constraints = []
  , calldata = cd
  , returnData = Nothing
  , reverted = False
  , halted = False
  , symbolicId = 0
  }

emptyStack :: SymStack
emptyStack = SymStack []

pushStack :: SymExpr -> SymStack -> Maybe SymStack
pushStack expr (SymStack xs)
  | length xs >= 1024 = Nothing
  | otherwise = Just (SymStack (expr : xs))

popStack :: SymStack -> Maybe (SymExpr, SymStack)
popStack (SymStack []) = Nothing
popStack (SymStack (x:xs)) = Just (x, SymStack xs)

popStackN :: Int -> SymStack -> Maybe ([SymExpr], SymStack)
popStackN n (SymStack xs)
  | n > length xs = Nothing
  | otherwise = Just (take n xs, SymStack (drop n xs))

peekStack :: Int -> SymStack -> Maybe SymExpr
peekStack n (SymStack xs)
  | n >= length xs = Nothing
  | otherwise = Just (xs !! n)

dupStack :: Int -> SymStack -> Maybe SymStack
dupStack n s = do
  expr <- peekStack (n - 1) s
  pushStack expr s

swapStack :: Int -> SymStack -> Maybe SymStack
swapStack n (SymStack xs)
  | n >= length xs = Nothing
  | otherwise = Just (SymStack (xs0 : (take (n-1) xs1) ++ [xn] ++ (drop n xs1)))
  where
    xs0 = head xs
    xs1 = tail xs
    xn = xs !! n

stackSize :: SymStack -> Int
stackSize (SymStack xs) = length xs

emptyMemory :: SymMemory
emptyMemory = SymMemory Map.empty

mstore :: Integer -> SymExpr -> SymMemory -> SymMemory
mstore offset value (SymMemory m) = SymMemory (Map.insert offset value m)

mstore8 :: Integer -> SymExpr -> SymMemory -> SymMemory
mstore8 = mstore

mload :: Integer -> SymMemory -> SymExpr
mload offset (SymMemory m) = Map.findWithDefault (Concrete 0) offset m

emptyStorage :: SymStorage
emptyStorage = SymStorage Map.empty

sstore :: SymExpr -> SymExpr -> SymStorage -> SymStorage
sstore key value (SymStorage s) = SymStorage (Map.insert key value s)

sload :: SymExpr -> SymStorage -> SymExpr
sload key (SymStorage s) = Map.findWithDefault (Concrete 0) key s

addConstraint :: Constraint -> ExecState -> ExecState
addConstraint c state = state { constraints = c : constraints state }

newSymbolic :: String -> ExecState -> (SymExpr, ExecState)
newSymbolic name state = 
  let sid = symbolicId state
      sym = Symbolic name sid
      state' = state { symbolicId = sid + 1 }
  in (sym, state')

prettyStack :: SymStack -> String
prettyStack (SymStack xs) = unlines $
  [ "[" ++ show i ++ "] " ++ prettyExpr expr 
  | (i, expr) <- zip [0..] xs ]

prettyMemory :: SymMemory -> String
prettyMemory (SymMemory m) = unlines $
  [ "  0x" ++ showHex offset ++ ": " ++ prettyExpr value
  | (offset, value) <- Map.toList m ]

prettyStorage :: SymStorage -> String
prettyStorage (SymStorage s) = unlines $
  [ "  " ++ prettyExpr key ++ " => " ++ prettyExpr value
  | (key, value) <- Map.toList s ]

showHex :: Integer -> String
showHex n = reverse $ take 4 $ reverse (go n) ++ repeat '0'
  where
    go 0 = "0"
    go n = let (q, r) = n `divMod` 16
               digit = "0123456789abcdef" !! fromIntegral r
           in if q == 0 then [digit] else go q ++ [digit]

prettyState :: ExecState -> String
prettyState state = unlines
  [ "PC: " ++ show (pc state)
  , "Gas: " ++ show (gas state)
  , "Halted: " ++ show (halted state)
  , "Reverted: " ++ show (reverted state)
  , "Stack:"
  , prettyStack (stack state)
  , "Constraints: " ++ show (length (constraints state))
  , unlines ["  " ++ prettyConstraint c | c <- reverse (constraints state)]
  ]