type error =
  | StackUnderflow
  | InvalidIndex of int
  | TypeMismatch of string
  | DivisionByZero
  | InvalidJump of int
  | InvalidClosure
  | NoExceptionHandler

type 'a result = Ok of 'a | Error of error

type t = {
  pc : int;
  stack : Instr.value list;
  code : Instr.t array;
  env : Instr.value list;
  env_stack : Instr.value list list;
  exception_value : Instr.value option;
}

let create code = {
  pc = 0;
  stack = [];
  code;
  env = [];
  env_stack = [];
  exception_value = None;
}

let with_pc state pc = { state with pc }
let with_stack state stack = { state with stack }
let with_env state env = { state with env }
let with_exception state exn = { state with exception_value = Some exn }
let clear_exception state = { state with exception_value = None }

let push state v = { state with stack = v :: state.stack }

let pop state = match state.stack with
  | [] -> Error StackUnderflow
  | v :: rest -> Ok (v, { state with stack = rest })

let pop_n state n =
  if n < 0 then Error (InvalidIndex n)
  else if n = 0 then Ok ([], state)
  else
    let rec aux acc n = function
      | [] when n > 0 -> Error StackUnderflow
      | rest when n = 0 -> Ok (List.rev acc, { state with stack = rest })
      | v :: rest -> aux (v :: acc) (n - 1) rest
    in
    aux [] n state.stack

let peek state = match state.stack with
  | [] -> Error StackUnderflow
  | v :: _ -> Ok v

let push_env state =
  { state with env_stack = state.env :: state.env_stack }

let pop_env state = match state.env_stack with
  | [] -> Error StackUnderflow
  | env :: rest -> Ok ({ state with env = env; env_stack = rest })

let get_env state = state.env

let set_env state env = { state with env }

let show state =
  Printf.sprintf "State(pc=%d, stack_len=%d, env_len=%d, env_stack_len=%d)"
    state.pc (List.length state.stack) (List.length state.env) (List.length state.env_stack)

let show_error = function
  | StackUnderflow -> "StackUnderflow"
  | InvalidIndex i -> Printf.sprintf "InvalidIndex(%d)" i
  | TypeMismatch s -> Printf.sprintf "TypeMismatch(%s)" s
  | DivisionByZero -> "DivisionByZero"
  | InvalidJump i -> Printf.sprintf "InvalidJump(%d)" i
  | InvalidClosure -> "InvalidClosure"
  | NoExceptionHandler -> "NoExceptionHandler"