.section .text
.global _start

.macro INC_TESTS
  la s10, test_count
  lw s11, 0(s10)
  addi s11, s11, 1
  sw s11, 0(s10)
.endm

.macro ASSERT_EQ reg, expected, msg
  li s10, \expected
  mv s11, \reg
  beq s10, s11, .L_success_\@
  # Failure
  la a0, .L_msg_\@
  mv a1, s11
  li a2, \expected
  jal report_failure
  j .L_end_\@
.L_success_\@:
  # Success
  la s10, success_count
  lw s11, 0(s10)
  addi s11, s11, 1
  sw s11, 0(s10)
.L_end_\@:
  .section .rodata
.L_msg_\@:  .string "\msg"
  .section .text
.endm

.macro ASSERT_EQ_REG reg, expected_reg, msg
  mv s10, \expected_reg
  mv s11, \reg
  beq s10, s11, .L_success_reg_\@
  # Failure
  la a0, .L_msg_reg_\@
  mv a1, s11
  mv a2, s10
  jal report_failure
  j .L_end_reg_\@
.L_success_reg_\@:
  # Success
  la s10, success_count
  lw s11, 0(s10)
  addi s11, s11, 1
  sw s11, 0(s10)
.L_end_reg_\@:
  .section .rodata
.L_msg_reg_\@:  .string "\msg"
  .section .text
.endm

.macro TEST_VAL_RI inst, val1, imm, expected, name
  INC_TESTS
  li t1, \val1
  \inst t0, t1, \imm
  ASSERT_EQ t0, \expected, \name
.endm

.macro TEST_VAL_RR inst, val1, val2, expected, name
  INC_TESTS
  li t1, \val1
  li t2, \val2
  \inst t0, t1, t2
  ASSERT_EQ t0, \expected, \name
.endm

_start:
  # ==========================================
  # RV32I Immediate Arithmetic
  # ==========================================
  TEST_VAL_RI addi, 10, 20, 30, addi_pos
  TEST_VAL_RI addi, 10, -20, -10, addi_neg
  TEST_VAL_RI addi, 0x7fffffff, 1, 0x80000000, addi_overflow
  
  TEST_VAL_RI slti, 10, 20, 1, slti_lt
  TEST_VAL_RI slti, 20, 10, 0, slti_gt
  TEST_VAL_RI slti, -10, 10, 1, slti_neg_lt
  
  TEST_VAL_RI sltiu, 10, 20, 1, sltiu_lt
  TEST_VAL_RI sltiu, -10, 10, 0, sltiu_neg_gt
  
  TEST_VAL_RI xori, 0xf0f0f0f0, 0x0ff, 0xf0f0f00f, xori_val
  TEST_VAL_RI ori, 0xf0f0f0f0, 0x0ff, 0xf0f0f0ff, ori_val
  TEST_VAL_RI andi, 0xf0f0f0f0, 0x0ff, 0x000000f0, andi_val
  
  TEST_VAL_RI slli, 1, 5, 32, slli_val
  TEST_VAL_RI slli, 0x80000000, 1, 0, slli_overflow
  
  TEST_VAL_RI srli, 32, 5, 1, srli_val
  TEST_VAL_RI srli, 0x80000000, 1, 0x40000000, srli_sign
  
  TEST_VAL_RI srai, 32, 5, 1, srai_pos
  TEST_VAL_RI srai, 0x80000000, 1, 0xc0000000, srai_neg

  # ==========================================
  # RV32I Register-Register Arithmetic
  # ==========================================
  TEST_VAL_RR add, 10, 20, 30, add_pos
  TEST_VAL_RR add, 0x7fffffff, 1, 0x80000000, add_overflow
  
  TEST_VAL_RR sub, 20, 10, 10, sub_pos
  TEST_VAL_RR sub, 10, 20, -10, sub_neg
  
  TEST_VAL_RR sll, 1, 5, 32, sll_val
  TEST_VAL_RR sll, 1, 37, 32, sll_wrap
  
  TEST_VAL_RR slt, 10, 20, 1, slt_lt
  TEST_VAL_RR slt, -10, 10, 1, slt_neg
  
  TEST_VAL_RR sltu, 10, 20, 1, sltu_lt
  TEST_VAL_RR sltu, -10, 10, 0, sltu_neg
  
  TEST_VAL_RR xor, 0xf0f0f0f0, 0x0f0f0f0f, 0xffffffff, xor_val
  
  TEST_VAL_RR srl, 32, 5, 1, srl_val
  TEST_VAL_RR srl, 32, 37, 1, srl_wrap
  
  TEST_VAL_RR sra, 0x80000000, 1, 0xc0000000, sra_neg
  TEST_VAL_RR sra, 0x80000000, 33, 0xc0000000, sra_wrap
  
  TEST_VAL_RR or, 0xf0f0f0f0, 0x0f0f0f0f, 0xffffffff, or_val
  TEST_VAL_RR and, 0xf0f0f0f0, 0x0f0f0f0f, 0, and_val

  # ==========================================
  # RV32I LUI and AUIPC
  # ==========================================
  INC_TESTS
  lui t0, 0x12345
  ASSERT_EQ t0, 0x12345000, lui_val
  
  INC_TESTS
  1: auipc t0, 0x10
  la t1, 1b
  sub t0, t0, t1
  ASSERT_EQ t0, 0x10000, auipc_val

  # ==========================================
  # RV32I Jumps
  # ==========================================
  INC_TESTS
  jal t0, 2f
  li t1, 0
  j 3f
2:
  li t1, 1
3:
  ASSERT_EQ t1, 1, jal_branch
  INC_TESTS
  la t2, 2b - 8
  ASSERT_EQ_REG t0, t2, jal_link

  INC_TESTS
  la t1, 4f
  jalr t0, t1, 0
  li t2, 0
  j 5f
4:
  li t2, 1
5:
  ASSERT_EQ t2, 1, jalr_branch
  INC_TESTS
  la t3, 4b - 8
  ASSERT_EQ_REG t0, t3, jalr_link

  INC_TESTS
  la t1, 6f - 4
  jalr t0, t1, 4
  li t2, 0
  j 7f
6:
  li t2, 1
7:
  ASSERT_EQ t2, 1, jalr_offset

  # ==========================================
  # RV32I Branches
  # ==========================================
  # BEQ
  INC_TESTS
  li t1, 10
  li t2, 10
  beq t1, t2, 1f
  li t0, 0
  j 2f
1:
  li t0, 1
2:
  ASSERT_EQ t0, 1, beq_taken

  INC_TESTS
  li t1, 10
  li t2, 20
  beq t1, t2, 1f
  li t0, 1
  j 2f
1:
  li t0, 0
2:
  ASSERT_EQ t0, 1, beq_not_taken

  # BNE
  INC_TESTS
  li t1, 10
  li t2, 20
  bne t1, t2, 1f
  li t0, 0
  j 2f
1:
  li t0, 1
2:
  ASSERT_EQ t0, 1, bne_taken

  INC_TESTS
  li t1, 10
  li t2, 10
  bne t1, t2, 1f
  li t0, 1
  j 2f
1:
  li t0, 0
2:
  ASSERT_EQ t0, 1, bne_not_taken

  # BLT
  INC_TESTS
  li t1, -10
  li t2, 10
  blt t1, t2, 1f
  li t0, 0
  j 2f
1:
  li t0, 1
2:
  ASSERT_EQ t0, 1, blt_taken

  INC_TESTS
  li t1, 10
  li t2, -10
  blt t1, t2, 1f
  li t0, 1
  j 2f
1:
  li t0, 0
2:
  ASSERT_EQ t0, 1, blt_not_taken

  # BGE
  INC_TESTS
  li t1, 10
  li t2, -10
  bge t1, t2, 1f
  li t0, 0
  j 2f
1:
  li t0, 1
2:
  ASSERT_EQ t0, 1, bge_taken

  INC_TESTS
  li t1, -10
  li t2, 10
  bge t1, t2, 1f
  li t0, 1
  j 2f
1:
  li t0, 0
2:
  ASSERT_EQ t0, 1, bge_not_taken

  # BLTU
  INC_TESTS
  li t1, 10
  li t2, -10
  bltu t1, t2, 1f
  li t0, 0
  j 2f
1:
  li t0, 1
2:
  ASSERT_EQ t0, 1, bltu_taken

  INC_TESTS
  li t1, -10
  li t2, 10
  bltu t1, t2, 1f
  li t0, 1
  j 2f
1:
  li t0, 0
2:
  ASSERT_EQ t0, 1, bltu_not_taken

  # BGEU
  INC_TESTS
  li t1, -10
  li t2, 10
  bgeu t1, t2, 1f
  li t0, 0
  j 2f
1:
  li t0, 1
2:
  ASSERT_EQ t0, 1, bgeu_taken

  INC_TESTS
  li t1, 10
  li t2, -10
  bgeu t1, t2, 1f
  li t0, 1
  j 2f
1:
  li t0, 0
2:
  ASSERT_EQ t0, 1, bgeu_not_taken

  # ==========================================
  # RV32I Loads and Stores
  # ==========================================
  # Store word and load word/half/byte
  la t1, data_buffer
  li t2, 0x12345678
  sw t2, 0(t1)
  
  INC_TESTS
  lw t0, 0(t1)
  ASSERT_EQ t0, 0x12345678, load_word
  
  INC_TESTS
  lh t0, 0(t1)
  ASSERT_EQ t0, 0x5678, load_half
  
  INC_TESTS
  lhu t0, 0(t1)
  ASSERT_EQ t0, 0x5678, load_half_u
  
  INC_TESTS
  lb t0, 0(t1)
  ASSERT_EQ t0, 0x78, load_byte
  
  INC_TESTS
  lbu t0, 0(t1)
  ASSERT_EQ t0, 0x78, load_byte_u

  # Store half and byte with sign extension verification
  li t2, 0x8000
  sh t2, 0(t1)
  INC_TESTS
  lh t0, 0(t1)
  ASSERT_EQ t0, -32768, load_half_neg
  
  INC_TESTS
  lhu t0, 0(t1)
  ASSERT_EQ t0, 32768, load_half_neg_u

  li t2, 0x80
  sb t2, 0(t1)
  INC_TESTS
  lb t0, 0(t1)
  ASSERT_EQ t0, -128, load_byte_neg
  
  INC_TESTS
  lbu t0, 0(t1)
  ASSERT_EQ t0, 128, load_byte_neg_u

  # ==========================================
  # RV32M Extensions
  # ==========================================
  TEST_VAL_RR mul, 10, 20, 200, mul_pos
  TEST_VAL_RR mul, -10, 20, -200, mul_neg
  TEST_VAL_RR mul, 0x10000, 0x10000, 0, mul_overflow
  
  TEST_VAL_RR mulh, 0x10000, 0x10000, 1, mulh_pos
  TEST_VAL_RR mulh, -10, 20, -1, mulh_neg
  
  TEST_VAL_RR mulhsu, -10, 20, -1, mulhsu_neg
  TEST_VAL_RR mulhu, 0x80000000, 2, 1, mulhu_val
  
  TEST_VAL_RR div, 20, 6, 3, div_pos
  TEST_VAL_RR div, -20, 6, -3, div_neg
  TEST_VAL_RR div, 20, 0, -1, div_by_zero
  TEST_VAL_RR div, -2147483648, -1, -2147483648, div_overflow
  
  TEST_VAL_RR divu, 20, 6, 3, divu_pos
  TEST_VAL_RR divu, -20, 6, 715827879, divu_neg
  TEST_VAL_RR divu, 20, 0, 4294967295, divu_by_zero
  
  TEST_VAL_RR rem, 20, 6, 2, rem_pos
  TEST_VAL_RR rem, -20, 6, -2, rem_neg
  TEST_VAL_RR rem, 20, 0, 20, rem_by_zero
  TEST_VAL_RR rem, -2147483648, -1, 0, rem_overflow
  
  TEST_VAL_RR remu, 20, 6, 2, remu_pos
  TEST_VAL_RR remu, -20, 6, 2, remu_neg
  TEST_VAL_RR remu, 20, 0, 20, remu_by_zero

  # ==========================================
  # RV32A Extensions
  # ==========================================
  # Load-Reserved / Store-Conditional
  la t1, atomic_var
  li t2, 42
  lr.w t0, (t1)
  sc.w t3, t2, (t1)
  INC_TESTS
  ASSERT_EQ t3, 0, sc_success
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, 42, sc_val
  
  # SC without LR (should fail)
  li t2, 99
  sc.w t3, t2, (t1)
  INC_TESTS
  ASSERT_EQ t3, 1, sc_fail_no_lr

  # AMOSWAP.W
  li t2, 10
  sw t2, 0(t1)
  li t3, 20
  amoswap.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, 10, amoswap_old
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, 20, amoswap_new

  # AMOADD.W
  li t2, 10
  sw t2, 0(t1)
  li t3, 20
  amoadd.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, 10, amoadd_old
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, 30, amoadd_new

  # AMOXOR.W
  li t2, 0x12345678
  sw t2, 0(t1)
  li t3, 0x0f0f0f0f
  amoxor.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, 0x12345678, amoxor_old
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, 0x1d3b5977, amoxor_new

  # AMOAND.W
  li t2, 0x12345678
  sw t2, 0(t1)
  li t3, 0x0f0f0f0f
  amoand.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, 0x12345678, amoand_old
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, 0x02040608, amoand_new

  # AMOOR.W
  li t2, 0x12345678
  sw t2, 0(t1)
  li t3, 0x0f0f0f0f
  amoor.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, 0x12345678, amoor_old
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, 0x1f3f5f7f, amoor_new

  # AMOMIN.W
  li t2, -10
  sw t2, 0(t1)
  li t3, 10
  amomin.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, -10, amomin_old1
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, -10, amomin_new1

  li t2, 10
  sw t2, 0(t1)
  li t3, -10
  amomin.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, 10, amomin_old2
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, -10, amomin_new2

  # AMOMAX.W
  li t2, -10
  sw t2, 0(t1)
  li t3, 10
  amomax.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, -10, amomax_old
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, 10, amomax_new

  # AMOMINU.W
  li t2, -10  # 0xfffffff6
  sw t2, 0(t1)
  li t3, 10
  amominu.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, -10, amominu_old
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, 10, amominu_new

  # AMOMAXU.W
  li t2, -10  # 0xfffffff6
  sw t2, 0(t1)
  li t3, 10
  amomaxu.w t0, t3, (t1)
  INC_TESTS
  ASSERT_EQ t0, -10, amomaxu_old
  INC_TESTS
  lw t4, 0(t1)
  ASSERT_EQ t4, -10, amomaxu_new

  # ==========================================
  # RV32I Fence
  # ==========================================
  fence

  # ==========================================
  # Report Results
  # ==========================================
  # Print Summary
  la a0, str_summary_prefix
  jal print_string
  
  # Print total test count
  la t0, test_count
  lw a0, 0(t0)
  jal print_decimal
  
  la a0, str_summary_tests
  jal print_string
  
  # Print passed count
  la t0, success_count
  lw a0, 0(t0)
  jal print_decimal
  
  la a0, str_summary_passed
  jal print_string
  
  # Print failed count
  la t0, failure_count
  lw a0, 0(t0)
  jal print_decimal
  
  la a0, str_summary_failed
  jal print_string

  la t0, failure_count
  lw t1, 0(t0)
  bnez t1, report_error
  
  la a0, str_all_passed
  jal print_string
  li a0, 0
  li a7, 93
  ecall

report_error:
  la a0, str_some_failed
  jal print_string
  li a0, 1
  li a7, 93
  ecall

report_failure:
  addi sp, sp, -32
  sw ra, 28(sp)
  sw s0, 24(sp)
  sw s1, 20(sp)
  sw s2, 16(sp)
  
  mv s0, a0
  mv s1, a1
  mv s2, a2
  
  la a0, str_fail_prefix
  jal print_string
  
  mv a0, s0
  jal print_string
  
  la a0, str_expected
  jal print_string
  
  mv a0, s2
  jal print_hex
  
  la a0, str_got
  jal print_string
  
  mv a0, s1
  jal print_hex
  
  li a0, 10
  jal print_char
  
  la t0, failure_count
  lw t1, 0(t0)
  addi t1, t1, 1
  sw t1, 0(t0)
  
  lw ra, 28(sp)
  lw s0, 24(sp)
  lw s1, 20(sp)
  lw s2, 16(sp)
  addi sp, sp, 32
  ret

print_string:
  addi sp, sp, -16
  sw ra, 12(sp)
  sw s0, 8(sp)
  mv s0, a0
print_string_loop:
  lbu a0, 0(s0)
  beqz a0, print_string_done
  jal print_char
  addi s0, s0, 1
  j print_string_loop
print_string_done:
  lw ra, 12(sp)
  lw s0, 8(sp)
  addi sp, sp, 16
  ret

print_char:
  addi sp, sp, -16
  sb a0, 0(sp)
  li a7, 64
  li a0, 1
  mv a1, sp
  li a2, 1
  ecall
  addi sp, sp, 16
  ret

print_hex:
  addi sp, sp, -32
  sw ra, 28(sp)
  sw s0, 24(sp)
  sw s1, 20(sp)
  sw s2, 16(sp)
  
  mv s0, a0
  li s1, 8
  la s2, hex_chars
  
print_hex_loop:
  srli t0, s0, 28
  andi t0, t0, 0xF
  add t1, s2, t0
  lbu a0, 0(t1)
  jal print_char
  slli s0, s0, 4
  addi s1, s1, -1
  bnez s1, print_hex_loop
  
  lw ra, 28(sp)
  lw s0, 24(sp)
  lw s1, 20(sp)
  lw s2, 16(sp)
  addi sp, sp, 32
  ret

print_decimal:
  addi sp, sp, -32
  sw ra, 28(sp)
  sw s0, 24(sp)
  sw s1, 20(sp)
  
  mv s0, a0
  
  # Handle 0 case
  bnez s0, 1f
  li a0, 48 # '0'
  jal print_char
  j print_decimal_done
1:
  addi s1, sp, 8
  li t2, 10
2:
  divu t0, s0, t2
  remu t1, s0, t2
  addi t1, t1, 48 # convert to ASCII
  sb t1, 0(s1)
  addi s1, s1, 1
  mv s0, t0
  bnez s0, 2b
  
  # Now print digits in reverse order
3:
  addi s1, s1, -1
  lbu a0, 0(s1)
  jal print_char
  addi t0, sp, 8
  bne s1, t0, 3b

print_decimal_done:
  lw ra, 28(sp)
  lw s0, 24(sp)
  lw s1, 20(sp)
  addi sp, sp, 32
  ret

.section .data
.align 2
test_count:    .word 0
success_count: .word 0
failure_count: .word 0

.align 4
data_buffer:   .zero 16
atomic_var:    .word 0

.section .rodata
hex_chars:          .string "0123456789abcdef"
str_fail_prefix:    .string "  FAIL: "
str_expected:       .string ", expected 0x"
str_got:            .string ", got 0x"
str_summary_prefix: .string "Executed "
str_summary_tests:  .string " tests: "
str_summary_passed: .string " passed, "
str_summary_failed: .string " failed.\n"
str_all_passed:     .string "All tests PASSED!\n"
str_some_failed:    .string "Some tests FAILED.\n"