open Sequence

let example_program =
  [|
    Instr.PUSH (Instr.Int 5);
    Instr.PUSH (Instr.Int 3);
    Instr.ADD;
    Instr.DUP;
    Instr.PUSH (Instr.Int 2);
    Instr.MUL;
    Instr.ADD;
    Instr.HALT;
  |]

let arithmetic_program =
  [|
    Instr.PUSH (Instr.Int 10);
    Instr.PUSH (Instr.Int 7);
    Instr.SUB;
    Instr.PUSH (Instr.Int 2);
    Instr.MUL;
    Instr.HALT;
  |]

let conditional_program =
  [|
    Instr.PUSH (Instr.Int 5);
    Instr.PUSH (Instr.Int 3);
    Instr.LT;
    Instr.JMP_IF 2;
    Instr.PUSH (Instr.Int 100);
    Instr.JMP 1;
    Instr.PUSH (Instr.Int 200);
    Instr.HALT;
  |]

let comparison_program =
  [|
    Instr.PUSH (Instr.Int 10);
    Instr.PUSH (Instr.Int 5);
    Instr.GT;
    Instr.PUSH (Instr.Int 5);
    Instr.PUSH (Instr.Int 10);
    Instr.LE;
    Instr.AND;
    Instr.HALT;
  |]

let division_program =
  [|
    Instr.PUSH (Instr.Int 20);
    Instr.PUSH (Instr.Int 4);
    Instr.DIV;
    Instr.PUSH (Instr.Int 3);
    Instr.MOD;
    Instr.HALT;
  |]

let stack_ops_program =
  [|
    Instr.PUSH (Instr.Int 1);
    Instr.PUSH (Instr.Int 2);
    Instr.PUSH (Instr.Int 3);
    Instr.ROT;
    Instr.SWAP;
    Instr.OVER;
    Instr.ADD;
    Instr.HALT;
  |]

let logical_program =
  [|
    Instr.PUSH (Instr.Bool true);
    Instr.PUSH (Instr.Bool false);
    Instr.OR;
    Instr.PUSH (Instr.Bool true);
    Instr.AND;
    Instr.NOT;
    Instr.HALT;
  |]

let negation_program =
  [|
    Instr.PUSH (Instr.Int 42);
    Instr.NEG;
    Instr.PUSH (Instr.Int 10);
    Instr.ADD;
    Instr.HALT;
  |]

let closure_program =
  [|
    Instr.PUSH (Instr.Int 100);
    Instr.MK_CLOSURE 10;
    Instr.CALL;
    Instr.HALT;
    Instr.PUSH (Instr.Int 200);
    Instr.RET;
  |]

let exception_program =
  [|
    Instr.TRY 3;
    Instr.PUSH (Instr.Int 42);
    Instr.PUSH (Instr.Int 0);
    Instr.DIV;
    Instr.HALT;
    Instr.POP;
    Instr.PUSH (Instr.Int 999);
    Instr.HALT;
  |]

let complex_program =
  [|
    Instr.PUSH (Instr.Int 15);
    Instr.PUSH (Instr.Int 3);
    Instr.DIV;
    Instr.DUP;
    Instr.PUSH (Instr.Int 2);
    Instr.MUL;
    Instr.SWAP;
    Instr.ADD;
    Instr.PUSH (Instr.Int 10);
    Instr.GT;
    Instr.JMP_IF 2;
    Instr.PUSH (Instr.Int 0);
    Instr.JMP 1;
    Instr.PUSH (Instr.Int 1);
    Instr.HALT;
  |]

let run_test name program expected =
  Printf.printf "Test: %s\n" name;
  match Vm.execute_program program with
  | Vm.Success (Some (Instr.Int result)) ->
    if result = expected then
      Printf.printf "  ✓ Pass: %d (expected %d)\n\n" result expected
    else
      Printf.printf "  ✗ Fail: got %d, expected %d\n\n" result expected
  | Vm.Success (Some (Instr.Bool result)) ->
    Printf.printf "  Result: %b\n\n" result
  | Vm.Success None ->
    Printf.printf "  ✗ Fail: No result\n\n"
  | Vm.Failure e ->
    Printf.printf "  ✗ Error: %s\n\n" (State.show_error e)

let run_trace_test name program =
  Printf.printf "Trace test: %s\n" name;
  let result, trace = Vm.execute_with_trace program in
  (match result with
   | Vm.Success (Some (Instr.Int n)) ->
     Printf.printf "  Result: %d\n" n
   | Vm.Success (Some (Instr.Bool b)) ->
     Printf.printf "  Result: %b\n" b
   | Vm.Success None ->
     Printf.printf "  Result: None\n"
   | Vm.Failure e ->
     Printf.printf "  Error: %s\n" (State.show_error e));
  Printf.printf "  Trace entries: %d\n" (List.length trace);
  if List.length trace > 0 then
    (let last = List.hd (List.rev trace) in
     Printf.printf "  Final: pc=%d, stack_len=%d, cont_depth=%d\n"
       last.Trampoline.pc last.Trampoline.stack_len last.Trampoline.cont_depth);
  Printf.printf "\n"

let () =
  run_test "Basic arithmetic" example_program 24;
  run_test "Subtraction and multiplication" arithmetic_program 6;
  run_test "Conditional branching" conditional_program 200;
  run_test "Comparison operations" comparison_program 1;
  run_test "Division and modulo" division_program 2;
  run_test "Stack operations" stack_ops_program 5;
  run_test "Logical operations" logical_program 0;
  run_test "Negation" negation_program (-32);
  run_test "Complex expression" complex_program 1;
  
  let invalid_jump = [|
    Instr.PUSH (Instr.Int 1);
    Instr.JMP 1000;
    Instr.HALT;
  |] in
  let errors = Vm.validate invalid_jump in
  Printf.printf "Invalid jump validation: %d errors found\n" (List.length errors);
  List.iter (fun (pc, err) ->
    Printf.printf "  PC %d: %s\n" pc (State.show_error err))
    errors;
  Printf.printf "\n";
  
  let disasm = Vm.disassemble example_program in
  List.iter (fun (pc, instr) ->
    Printf.printf "  %d: %s\n" pc instr)
    disasm;
  Printf.printf "\n";
  
  run_trace_test "Traced execution" arithmetic_program;
  
  match Vm.execute_program exception_program with
  | Vm.Success (Some (Instr.Int n)) ->
    Printf.printf "Exception caught, result: %d\n" n
  | Vm.Success None ->
    Printf.printf "Exception caught, no result\n"
  | Vm.Failure e ->
    Printf.printf "Exception not caught: %s\n" (State.show_error e);
  Printf.printf "\n"