{-# LANGUAGE LambdaCase #-}

module Executor where

import SymbolicExpression
import SymbolicState
import Opcode
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Maybe (fromMaybe)
import Control.Monad (foldM)
import Prelude hiding (LT, GT, EQ)

data StepResult
  = Continue ExecState
  | Branch ExecState ExecState
  | Halt ExecState
  | Error String
  deriving (Show)

executeOpcode :: ByteString -> Opcode -> ExecState -> StepResult
executeOpcode code op state = case op of
  ADD -> binOp mkAdd state
  MUL -> binOp mkMul state
  SUB -> binOp mkSub state
  DIV -> binOp mkDiv state
  MOD -> binOp mkMod state
  
  LT -> binOp mkLt state
  GT -> binOp mkGt state
  EQ -> binOp mkEq state
  ISZERO -> unOp mkIsZero state
  
  AND -> binOp mkAnd state
  OR -> binOp mkOr state
  XOR -> binOp mkXor state
  NOT -> unOp mkNot state
  
  POP -> case popStack (stack state) of
    Just (_, s') -> Continue state { stack = s', pc = pc state + 1 }
    Nothing -> Error "Stack underflow"
  
  _ | Just n <- dupDepth op -> case dupStack n (stack state) of
    Just s' -> Continue state { stack = s', pc = pc state + 1 }
    Nothing -> Error "Stack underflow on DUP"
  
  _ | Just n <- swapDepth op -> case swapStack n (stack state) of
    Just s' -> Continue state { stack = s', pc = pc state + 1 }
    Nothing -> Error "Stack underflow on SWAP"
  
  MLOAD -> case popStack (stack state) of
    Just (offset, s') ->
      let value = case offset of
            Concrete n -> mload n (memory state)
            _ -> Symbolic "mload" (symbolicId state)
          state' = state { symbolicId = symbolicId state + 1 }
      in case pushStack value s' of
        Just s'' -> Continue state' { stack = s'', pc = pc state' + 1 }
        Nothing -> Error "Stack overflow"
    Nothing -> Error "Stack underflow"
  
  MSTORE -> case popStackN 2 (stack state) of
    Just ([offset, value], s') ->
      let mem' = case offset of
            Concrete n -> mstore n value (memory state)
            _ -> memory state
      in Continue state { stack = s', memory = mem', pc = pc state + 1 }
    _ -> Error "Stack underflow"
  
  MSTORE8 -> case popStackN 2 (stack state) of
    Just ([offset, value], s') ->
      let mem' = case offset of
            Concrete n -> mstore8 n value (memory state)
            _ -> memory state
      in Continue state { stack = s', memory = mem', pc = pc state + 1 }
    _ -> Error "Stack underflow"
  
  SLOAD -> case popStack (stack state) of
    Just (key, s') ->
      let value = sload key (storage state)
      in case pushStack value s' of
        Just s'' -> Continue state { stack = s'', pc = pc state + 1 }
        Nothing -> Error "Stack overflow"
    Nothing -> Error "Stack underflow"
  
  SSTORE -> case popStackN 2 (stack state) of
    Just ([key, value], s') ->
      let stor' = sstore key value (storage state)
      in Continue state { stack = s', storage = stor', pc = pc state + 1 }
    _ -> Error "Stack underflow"
  
  JUMP -> case popStack (stack state) of
    Just (Concrete dest, s') ->
      Continue state { stack = s', pc = fromInteger dest }
    Just (_, _) -> Error "Symbolic jump destination not supported"
    Nothing -> Error "Stack underflow"
  
  JUMPI -> case popStackN 2 (stack state) of
    Just ([Concrete dest, cond], s') ->
      let trueBranch = addConstraint (CTrue cond) $
                       state { stack = s', pc = fromInteger dest }
          falseBranch = addConstraint (CFalse cond) $
                        state { stack = s', pc = pc state + 1 }
      in case cond of
        Concrete 0 -> Continue falseBranch
        Concrete _ -> Continue trueBranch
        _ -> Branch trueBranch falseBranch
    Just ([_, cond], s') -> Error "Symbolic jump destination not supported"
    _ -> Error "Stack underflow"
  
  JUMPDEST -> Continue state { pc = pc state + 1 }
  
  PC -> case pushStack (Concrete $ fromIntegral $ pc state) (stack state) of
    Just s' -> Continue state { stack = s', pc = pc state + 1 }
    Nothing -> Error "Stack overflow"
  
  CALLER -> case pushStack Caller (stack state) of
    Just s' -> Continue state { stack = s', pc = pc state + 1 }
    Nothing -> Error "Stack overflow"
  
  CALLVALUE -> case pushStack CallValue (stack state) of
    Just s' -> Continue state { stack = s', pc = pc state + 1 }
    Nothing -> Error "Stack overflow"
  
  CALLDATALOAD -> case popStack (stack state) of
    Just (offset, s') ->
      let value = case offset of
            Concrete n -> loadCalldata (calldata state) n
            _ -> CallDataLoad offset
      in case pushStack value s' of
        Just s'' -> Continue state { stack = s'', pc = pc state + 1 }
        Nothing -> Error "Stack overflow"
    Nothing -> Error "Stack underflow"
  
  CALLDATASIZE -> case pushStack CallDataSize (stack state) of
    Just s' -> Continue state { stack = s', pc = pc state + 1 }
    Nothing -> Error "Stack overflow"
  
  STOP -> Halt state { halted = True }
  
  RETURN -> case popStackN 2 (stack state) of
    Just ([offset, size], s') ->
      let retVal = case (offset, size) of
            (Concrete o, Concrete s) -> mload o (memory state)
            _ -> Symbolic "return" (symbolicId state)
      in Halt state { stack = s', returnData = Just retVal, halted = True }
    _ -> Error "Stack underflow"
  
  REVERT -> Halt state { reverted = True, halted = True }
  
  INVALID -> Halt state { reverted = True, halted = True }
  
  _ -> Error $ "Unsupported opcode: " ++ show op

binOp :: (SymExpr -> SymExpr -> SymExpr) -> ExecState -> StepResult
binOp f state = case popStackN 2 (stack state) of
  Just ([a, b], s') ->
    case pushStack (f a b) s' of
      Just s'' -> Continue state { stack = s'', pc = pc state + 1 }
      Nothing -> Error "Stack overflow"
  _ -> Error "Stack underflow"

unOp :: (SymExpr -> SymExpr) -> ExecState -> StepResult
unOp f state = case popStack (stack state) of
  Just (a, s') ->
    case pushStack (f a) s' of
      Just s'' -> Continue state { stack = s'', pc = pc state + 1 }
      Nothing -> Error "Stack overflow"
  Nothing -> Error "Stack underflow"

loadCalldata :: ByteString -> Integer -> SymExpr
loadCalldata cd offset
  | offset >= 0 && offset < fromIntegral (BS.length cd) =
      let bytes = BS.drop (fromInteger offset) cd
          value = foldl (\acc b -> acc * 256 + fromIntegral b) 0 (BS.unpack $ BS.take 32 bytes)
      in Concrete value
  | otherwise = Concrete 0

step :: ByteString -> ExecState -> StepResult
step code state
  | halted state = Halt state
  | pc state < 0 || pc state >= BS.length code = Halt state { halted = True }
  | otherwise =
      let byte = BS.index code (pc state)
          opcode = decodeOpcode byte
      in if isPush opcode
         then
           let n = pushBytes opcode
               value = readBytes code (pc state + 1) n
           in case pushStack (Concrete value) (stack state) of
                Just s' -> Continue state { stack = s', pc = pc state + 1 + n }
                Nothing -> Error "Stack overflow"
         else executeOpcode code opcode state

run :: ByteString -> ExecState -> Int -> [ExecState]
run code state maxSteps = go state 0
  where
    go st n
      | n >= maxSteps = [st]
      | halted st = [st]
      | otherwise = case step code st of
          Continue st' -> go st' (n + 1)
          Halt st' -> [st']
          Branch _ _ -> [st]
          Error msg -> [st { halted = True, reverted = True }]

explore :: ByteString -> ExecState -> Int -> [ExecState]
explore code initialState maxSteps = go [initialState] 0
  where
    go states n
      | n >= maxSteps = states
      | all halted states = states
      | otherwise =
          let states' = concatMap stepOne states
          in go states' (n + 1)
    
    stepOne st
      | halted st = [st]
      | otherwise = case step code st of
          Continue st' -> [st']
          Halt st' -> [st']
          Branch st1 st2 -> [st1, st2]
          Error msg -> [st { halted = True, reverted = True }]