open Instr

type result =
  | Continue of { state : State.t; cont : Cont.t }
  | Halted of value option
  | Error of State.error

let step state cont : result =
  if state.pc >= Array.length state.code then
    Halted None
  else
    let instr = state.code.(state.pc) in
    let next_pc = state.pc + 1 in
    let next_state = State.with_pc state next_pc in
    
    match instr with
    | PUSH v ->
      let new_state = State.push next_state v in
      Continue { state = new_state; cont }
    
    | POP ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (_, new_state) ->
         Continue { state = new_state; cont })
    
    | DUP ->
      (match State.peek next_state with
       | State.Error e -> Error e
       | State.Ok v ->
         let new_state = State.push next_state v in
         Continue { state = new_state; cont })
    
    | SWAP ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([a; b], new_state) ->
         let new_state = State.push (State.push new_state a) b in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "SWAP"))
    
    | OVER ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([a; b], new_state) ->
         let new_state = State.push (State.push new_state b) a in
         let new_state = State.push new_state b in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "OVER"))
    
    | ROT ->
      (match State.pop_n next_state 3 with
       | State.Error e -> Error e
       | State.Ok ([a; b; c], new_state) ->
         let new_state = State.push (State.push (State.push new_state a) c) b in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "ROT"))
    
    | ADD ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (a + b)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "ADD"))
    
    | SUB ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (a - b)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "SUB"))
    
    | MUL ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (a * b)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "MUL"))
    
    | DIV ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int 0], _) -> Error State.DivisionByZero
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (a / b)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "DIV"))
    
    | MOD ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int 0], _) -> Error State.DivisionByZero
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (a mod b)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "MOD"))
    
    | EQ ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (if a = b then 1 else 0)) in
         Continue { state = new_state; cont }
       | State.Ok ([Bool a; Bool b], new_state) ->
         let new_state = State.push new_state (Bool (a = b)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "EQ"))
    
    | NE ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (if a <> b then 1 else 0)) in
         Continue { state = new_state; cont }
       | State.Ok ([Bool a; Bool b], new_state) ->
         let new_state = State.push new_state (Bool (a <> b)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "NE"))
    
    | LT ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (if a < b then 1 else 0)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "LT"))
    
    | LE ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (if a <= b then 1 else 0)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "LE"))
    
    | GT ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (if a > b then 1 else 0)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "GT"))
    
    | GE ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([Int a; Int b], new_state) ->
         let new_state = State.push new_state (Int (if a >= b then 1 else 0)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "GE"))
    
    | AND ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([a; b], new_state) ->
         (match (bool_of_value a, bool_of_value b) with
          | (Some a', Some b') ->
            let new_state = State.push new_state (Bool (a' && b')) in
            Continue { state = new_state; cont }
          | _ -> Error (State.TypeMismatch "AND")))
    
    | OR ->
      (match State.pop_n next_state 2 with
       | State.Error e -> Error e
       | State.Ok ([a; b], new_state) ->
         (match (bool_of_value a, bool_of_value b) with
          | (Some a', Some b') ->
            let new_state = State.push new_state (Bool (a' || b')) in
            Continue { state = new_state; cont }
          | _ -> Error (State.TypeMismatch "OR")))
    
    | NOT ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (v, new_state) ->
         (match bool_of_value v with
          | Some b ->
            let new_state = State.push new_state (Bool (not b)) in
            Continue { state = new_state; cont }
          | None -> Error (State.TypeMismatch "NOT")))
    
    | NEG ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (Int n, new_state) ->
         let new_state = State.push new_state (Int (-n)) in
         Continue { state = new_state; cont }
       | State.Ok _ -> Error (State.TypeMismatch "NEG"))
    
    | JMP offset ->
      let new_pc = next_state.pc + offset - 1 in
      if new_pc < 0 || new_pc >= Array.length state.code then
        Error (State.InvalidJump new_pc)
      else
        let new_state = State.with_pc next_state new_pc in
        Continue { state = new_state; cont }
    
    | JMP_IF offset ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (v, new_state) ->
         (match bool_of_value v with
          | Some true ->
            let new_pc = new_state.pc + offset - 1 in
            if new_pc < 0 || new_pc >= Array.length state.code then
              Error (State.InvalidJump new_pc)
            else
              let new_state = State.with_pc new_state new_pc in
              Continue { state = new_state; cont }
          | Some false ->
            Continue { state = new_state; cont }
          | None -> Error (State.TypeMismatch "JMP_IF")))
    
    | JMP_IF_NOT offset ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (v, new_state) ->
         (match bool_of_value v with
          | Some false ->
            let new_pc = new_state.pc + offset - 1 in
            if new_pc < 0 || new_pc >= Array.length state.code then
              Error (State.InvalidJump new_pc)
            else
              let new_state = State.with_pc new_state new_pc in
              Continue { state = new_state; cont }
          | Some true ->
            Continue { state = new_state; cont }
          | None -> Error (State.TypeMismatch "JMP_IF_NOT")))
    
    | CALL ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (func, new_state) ->
         (match func with
          | Closure { code; env } ->
            let new_state = State.push_env new_state in
            let new_state = State.set_env new_state env in
            let new_pc = code in
            if new_pc < 0 || new_pc >= Array.length state.code then
              Error (State.InvalidJump new_pc)
            else
              let new_state = State.with_pc new_state new_pc in
              let new_cont = Cont.return_to cont in
              Continue { state = new_state; cont = new_cont }
          | _ -> Error State.InvalidClosure))
    
    | RET ->
      (match State.pop_env next_state with
       | State.Error _ ->
         (match cont with
          | Return_to k ->
            Continue { state = next_state; cont = k }
          | Halt ->
            Halted (match State.peek next_state with State.Ok v -> Some v | _ -> None)
          | _ ->
            Halted (match State.peek next_state with State.Ok v -> Some v | _ -> None))
       | State.Ok (new_state, _) ->
         (match cont with
          | Return_to k ->
            Continue { state = new_state; cont = k }
          | Halt ->
            Halted (match State.peek new_state with State.Ok v -> Some v | _ -> None)
          | _ ->
            Halted (match State.peek new_state with State.Ok v -> Some v | _ -> None)))
    
    | MK_CLOSURE code ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (env_val, new_state) ->
         let env = match env_val with
           | Closure { env; _ } -> env
           | _ -> State.get_env new_state
         in
         let closure = Closure { code; env } in
         let new_state = State.push new_state closure in
         Continue { state = new_state; cont })
    
    | LOAD idx ->
      if idx < 0 || idx >= List.length next_state.env then
        Error (State.InvalidIndex idx)
      else
        let v = List.nth next_state.env idx in
        let new_state = State.push next_state v in
        Continue { state = new_state; cont }
    
    | STORE idx ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (v, new_state) ->
         let env = new_state.env in
         let env' = if idx < List.length env then
             let rec update i = function
               | [] -> []
               | x :: xs -> if i = idx then v :: xs else x :: update (i + 1) xs
             in update 0 env
           else
             let padding = List.init (idx - List.length env) (fun _ -> Int 0) in
             env @ padding @ [v]
         in
         let new_state = State.set_env new_state env' in
         Continue { state = new_state; cont })
    
    | PUSH_ENV ->
      let new_state = State.push_env next_state in
      Continue { state = new_state; cont }
    
    | POP_ENV ->
      (match State.pop_env next_state with
       | State.Error e -> Error e
       | State.Ok (new_state, _) ->
         Continue { state = new_state; cont })
    
    | TRY offset ->
      let handler_pc = next_state.pc + offset - 1 in
      if handler_pc < 0 || handler_pc >= Array.length state.code then
        Error (State.InvalidJump handler_pc)
      else
        let handler_cont = Cont.eval_to (state.code.(handler_pc)) Cont.halt in
        let new_cont = Cont.exception_to handler_cont cont in
        Continue { state = next_state; cont = new_cont }
    
    | RAISE ->
      (match State.pop next_state with
       | State.Error e -> Error e
       | State.Ok (exn_val, new_state) ->
         let new_state = State.with_exception new_state exn_val in
         let rec find_handler = function
           | Exception_to { handler; cont } ->
             let new_state = State.clear_exception new_state in
             let new_state = State.push new_state exn_val in
             Continue { state = new_state; cont = handler }
           | Return_to k -> find_handler k
           | Apply_to { cont; _ } -> find_handler cont
           | Eval_to { cont; _ } -> find_handler cont
           | Pop_to { cont; _ } -> find_handler cont
           | Halt -> Error State.NoExceptionHandler
         in
         find_handler cont)
    
    | RERAISE ->
      (match next_state.exception_value with
       | None -> Error State.NoExceptionHandler
       | Some exn_val ->
         let rec find_handler = function
           | Exception_to { handler; cont } ->
             let new_state = State.clear_exception new_state in
             let new_state = State.push new_state exn_val in
             Continue { state = new_state; cont = handler }
           | Return_to k -> find_handler k
           | Apply_to { cont; _ } -> find_handler cont
           | Eval_to { cont; _ } -> find_handler cont
           | Pop_to { cont; _ } -> find_handler cont
           | Halt -> Error State.NoExceptionHandler
         in
         find_handler cont)
    
    | HALT ->
      Halted (match State.peek next_state with State.Ok v -> Some v | _ -> None)