package riscvemu

import (
	"fmt"
)

type ErrInvalidOpcode struct{}

func (e ErrInvalidOpcode) Error() string { return "invalid opcode" }

type ErrBreak struct{}

func (e ErrBreak) Error() string { return "EBREAK" }

type CPUState struct {
	Registers [32]uint32
	Pc        uint32
	w         EEI
}

func NewCPU(w EEI) CPUState {
	return CPUState{w: w}
}

func opcode_rd(opcode uint32) uint32  { return (opcode >> 7) & 0b11111 }
func opcode_rs1(opcode uint32) uint32 { return (opcode >> 15) & 0b11111 } // 2^5 is 32 for 32 registers
func opcode_rs2(opcode uint32) uint32 { return (opcode >> 20) & 0b11111 }

func (c *CPUState) ReadRegister(r uint32) uint32 {
	if r == 0 {
		return 0 // Read from the zero register always returns 0
	}

	return c.Registers[r]
}

func (c *CPUState) Step() error {
	opcode, err := c.w.Read32(c.Pc)
	if err != nil {
		return fmt.Errorf("ReadMemory: %w", err)
	}

	switch opcode & 0b1111111 {
	case 0b0110111:
		// LUI
		// LUI (load upper immediate) is used to build 32-bit constants and uses the U-type format. LUI
		// places the 32-bit U-immediate value into the destination register rd, filling in the lowest 12 bits
		// with zeros.
		c.Registers[opcode_rd(opcode)] = opcode & 0b11111111111111111111000000000000

	case 0b0010111:
		// AUIPC
		// AUIPC (add upper immediate to pc) is used to build pc-relative addresses and uses the U-type
		// format. AUIPC forms a 32-bit offset from the U-immediate, filling in the lowest 12 bits with zeros,
		// adds this offset to the address of the AUIPC instruction, then places the result in register rd.
		c.Registers[opcode_rd(opcode)] = c.Pc + (opcode & 0b11111111111111111111000000000000)

	case 0b01101111:
		// JAL
		// The jump and link (JAL) instruction uses the J-type format, where the J-immediate encodes a
		// signed offset in multiples of 2 bytes. The offset is sign-extended and added to the address of the
		// jump instruction to form the jump target address. Jumps can therefore target a ±1 MiB range.
		// JAL stores the address of the instruction that follows the JAL (pc+4) into register rd. The standard
		// software calling convention uses x1 as the return address register and x5 as an alternate link register.
		// Plain unconditional jumps (assembler pseudoinstruction J) are encoded as a JAL with rd=x0.

		var imm uint32 = (opcode >> 12) & 0b11111111111111111111
		if imm&0b10000000000000000000 != 0 { // 20th position
			imm |= 0b11111111111100000000000000000000 // sign extension: if MSB is set in imm, extend with 1's instead of 0's
		} // FIXME the encoding seems weirder than this??

		// 2^20 has ~1MiB unsigned range, but it's in multiples of two bytes
		offset := int32(imm) * 2

		c.Registers[opcode_rd(opcode)] = c.Pc + 4
		c.Pc = uint32(int32(c.Pc) + offset)
		return nil // Don't fallthrough, or else we increment Pc again

	case 0b1100111:
		// JALR
		funct3 := (opcode >> 12) & 0b111
		if funct3 != 0b000 {
			return ErrInvalidOpcode{}
		}

		// The indirect jump instruction JALR (jump and link register) uses the I-type encoding. The target
		// address is obtained by adding the sign-extended 12-bit I-immediate to the register rs1, then setting
		// the least-significant bit of the result to zero. The address of the instruction following the jump
		// (pc+4) is written to register rd. Register x0 can be used as the destination if the result is not
		// required.

		var imm uint32 = (opcode >> 20) & 0b111111111111
		if imm&0b100000000000 != 0 { // 12th position
			imm |= 0b11111111111111111111000000000000 // sign extension: if MSB is set in imm, extend with 1's instead of 0's
		}

		c.Registers[opcode_rd(opcode)] = c.Pc + 4
		c.Pc = (c.Registers[opcode_rs1(opcode)] + imm) & 0b11111111111111111111111111111110
		return nil // Don't fallthrough, or else we increment Pc again

	case 0b1100011:
		// BEQ/BNE/BLT/BGE/BLTU/BGEU
		// All branch instructions use the B-type instruction format. The 12-bit B-immediate encodes signed
		// offsets in multiples of 2 bytes. The offset is sign-extended and added to the address of the branch
		// instruction to give the target address. The conditional branch range is ±4 KiB.

		var imm uint32 = (opcode>>6)&0b11111 | (((opcode >> 25) & 0b11111) << 5) // FIXME the encoding seems weirder than this??
		if imm&0b100000000000 != 0 {                                             // 12th position
			imm |= 0b11111111111111111111000000000000 // sign extension: if MSB is set in imm, extend with 1's instead of 0's
		}

		var offset int32 = int32(imm) * 2

		// Branch instructions compare two registers. BEQ and BNE take the branch if registers rs1 and rs2
		// are equal or unequal respectively. BLT and BLTU take the branch if rs1 is less than rs2, using
		// signed and unsigned comparison respectively. BGE and BGEU take the branch if rs1 is greater
		// than or equal to rs2, using signed and unsigned comparison respectively. Note, BGT, BGTU,
		// BLE, and BLEU can be synthesized by reversing the operands to BLT, BLTU, BGE, and BGEU,
		// respectively.
		funct3 := (opcode >> 12) & 0b111
		switch funct3 {
		case 0b000:
			// BEQ
			if c.ReadRegister(opcode_rs1(opcode)) == c.ReadRegister(opcode_rs2(opcode)) {
				c.Pc = uint32(int32(c.Pc) + offset)
				return nil // Don't fallthrough to regular Pc advance
			}

		case 0b001:
			// BNE
			if c.ReadRegister(opcode_rs1(opcode)) != c.ReadRegister(opcode_rs2(opcode)) {
				c.Pc = uint32(int32(c.Pc) + offset)
				return nil // Don't fallthrough to regular Pc advance
			}

		case 0b100:
			// BLT
			if int32(c.ReadRegister(opcode_rs1(opcode))) < int32(c.ReadRegister(opcode_rs2(opcode))) {
				c.Pc = uint32(int32(c.Pc) + offset)
				return nil // Don't fallthrough to regular Pc advance
			}

		case 0b101:
			// BGE
			if int32(c.ReadRegister(opcode_rs1(opcode))) >= int32(c.ReadRegister(opcode_rs2(opcode))) {
				c.Pc = uint32(int32(c.Pc) + offset)
				return nil // Don't fallthrough to regular Pc advance
			}

		case 0b110:
			// BLTU
			if c.ReadRegister(opcode_rs1(opcode)) < c.ReadRegister(opcode_rs2(opcode)) {
				c.Pc = uint32(int32(c.Pc) + offset)
				return nil // Don't fallthrough to regular Pc advance
			}

		case 0b111:
			// BGEU
			if c.ReadRegister(opcode_rs1(opcode)) >= c.ReadRegister(opcode_rs2(opcode)) {
				c.Pc = uint32(int32(c.Pc) + offset)
				return nil // Don't fallthrough to regular Pc advance
			}

		default:
			return ErrInvalidOpcode{}
		}

	case 0b0000011:
		// LB/LH/LW/LBU/LHU
		// The effective address is obtained by adding register rs1
		// to the sign-extended 12-bit offset. Loads copy a value from memory to register rd.

		var imm uint32 = (opcode >> 20) & 0b111111111111
		if imm&0b100000000000 != 0 { // 12th position
			imm |= 0b11111111111111111111000000000000 // sign extension: if MSB is set in imm, extend with 1's instead of 0's
		}

		memAddr := imm + c.Registers[opcode_rs1(opcode)]

		// The LW instruction loads a 32-bit value from memory into rd. LH loads a 16-bit value from memory,
		// then sign-extends to 32-bits before storing in rd. LHU loads a 16-bit value from memory but then
		// zero extends to 32-bits before storing in rd. LB and LBU are defined analogously for 8-bit values.

		funct3 := (opcode >> 12) & 0b111
		switch funct3 {
		case 0b000:
			// LB
			memVal, err := c.w.ReadByte(memAddr)
			if err != nil {
				return err
			}

			storeVal := uint32(memVal)
			if storeVal&0b10000000 != 0 { // sign extend / 8th position
				storeVal |= 0b11111111111111111111111100000000
			}
			c.Registers[opcode_rd(opcode)] = storeVal

		case 0b001:
			// LH
			memVal, err := c.w.Read16(memAddr)
			if err != nil {
				return err
			}

			storeVal := uint32(memVal)
			if storeVal&0b1000000000000000 != 0 { // sign extend / 16th position
				storeVal |= 0b11111111111111110000000000000000
			}
			c.Registers[opcode_rd(opcode)] = storeVal

		case 0b010:
			// LW
			memVal, err := c.w.Read32(memAddr)
			if err != nil {
				return err
			}
			c.Registers[opcode_rd(opcode)] = memVal

		case 0b100:
			// LBU
			memVal, err := c.w.ReadByte(memAddr)
			if err != nil {
				return err
			}

			c.Registers[opcode_rd(opcode)] = uint32(memVal)

		case 0b101:
			// LHU
			memVal, err := c.w.Read16(memAddr)
			if err != nil {
				return err
			}

			c.Registers[opcode_rd(opcode)] = uint32(memVal)

		default:
			return ErrInvalidOpcode{}
		}
		panic("todo")

	case 0b0100011:
		// SB/SH/SW

		imm := ((opcode >> 7) & 0b11111) | (((opcode >> 25) & 0b1111111) << 5)
		memAddr := c.Registers[opcode_rs1(opcode)] + imm

		// The SW, SH, and SB instructions store 32-bit, 16-bit, and 8-bit values from the low bits of register
		// rs2 to memory.
		funct3 := (opcode >> 12) & 0b111
		switch funct3 {
		case 0b000:
			// SB
			err := c.w.WriteByte(memAddr, byte(c.Registers[opcode_rs2(opcode)]))
			if err != nil {
				return err
			}

		case 0b001:
			// SH
			err := c.w.Write16(memAddr, uint16(c.Registers[opcode_rs2(opcode)]))
			if err != nil {
				return err
			}
		case 0b010:
			// SW
			err := c.w.Write32(memAddr, c.Registers[opcode_rs2(opcode)])
			if err != nil {
				return err
			}
		default:
			return ErrInvalidOpcode{}
		}

	case 0b0010011:
		// ADDI/SLTI/SLTIU/XORI/ORI/ANDI/SLLI/SLRI/SRAI
		funct3 := (opcode >> 12) & 0b111

		var imm uint32 = (opcode >> 20) & 0b111111111111
		if imm&0b100000000000 != 0 { // 12th position
			imm |= 0b11111111111111111111000000000000 // sign extension: if MSB is set in imm, extend with 1's instead of 0's
		}

		switch funct3 {
		case 0b000:
			// ADDI/NOP
			// ADDI adds the sign-extended 12-bit immediate to register rs1. Arithmetic overflow is ignored and
			// the result is simply the low XLEN bits of the result. ADDI rd, rs1, 0 is used to implement the MV
			// rd, rs1 assembler pseudoinstruction.
			//
			// The NOP instruction does not change any architecturally visible state, except for advancing the
			// pc and incrementing any applicable performance counters. NOP is encoded as ADDI x0, x0, 0.
			c.Registers[opcode_rd(opcode)] = c.ReadRegister(opcode_rs1(opcode)) + imm

		case 0b010:
			// SLTI
			// SLTI (set less than immediate) places the value 1 in register rd if register rs1 is less than the
			// signextended immediate when both are treated as signed numbers, else 0 is written to rd. SLTIU is
			// similar but compares the values as unsigned numbers (i.e., the immediate is first sign-extended to
			// XLEN bits then treated as an unsigned number). Note, SLTIU rd, rs1, 1 sets rd to 1 if rs1 equals
			// zero, otherwise sets rd to 0 (assembler pseudoinstruction SEQZ rd, rs).
			if int32(c.ReadRegister(opcode_rs1(opcode))) < int32(imm) {
				c.Registers[opcode_rd(opcode)] = 1
			} else {
				c.Registers[opcode_rd(opcode)] = 0
			}

		case 0b011:
			// SLTIU
			if c.ReadRegister(opcode_rs1(opcode)) < imm {
				c.Registers[opcode_rd(opcode)] = 1
			} else {
				c.Registers[opcode_rd(opcode)] = 0
			}

		case 0b100:
			// XORI
			// ANDI, ORI, XORI are logical operations that perform bitwise AND, OR, and XOR on register rs1
			// and the sign-extended 12-bit immediate and place the result in rd. Note, XORI rd, rs1, -1 performs
			// a bitwise logical inversion of register rs1 (assembler pseudoinstruction NOT rd, rs).
			c.Registers[opcode_rd(opcode)] = c.ReadRegister(opcode_rs1(opcode)) ^ imm

		case 0b110:
			// ORI
			c.Registers[opcode_rd(opcode)] = c.ReadRegister(opcode_rs1(opcode)) | imm

		case 0b111:
			// ANDI
			c.Registers[opcode_rd(opcode)] = c.ReadRegister(opcode_rs1(opcode)) & imm

		case 0b001:
			// SLLI
			// Shifts by a constant are encoded as a specialization of the I-type format. The operand to be shifted
			// is in rs1, and the shift amount is encoded in the lower 5 bits of the I-immediate field. The right
			// shift type is encoded in bit 30. SLLI is a logical left shift (zeros are shifted into the lower bits);
			// SRLI is a logical right shift (zeros are shifted into the upper bits); and SRAI is an arithmetic right
			// shift (the original sign bit is copied into the vacated upper bits).

			shamt := (opcode >> 20) & 0b1111
			c.Registers[opcode_rd(opcode)] = c.ReadRegister(opcode_rs1(opcode)) << shamt

		case 0b101:
			// SRLI/SRAI share the same `funct3`
			shamt := (opcode >> 20) & 0b1111
			distinguish := (opcode >> 25) & 0b1111111
			if distinguish == 0b0000000 {
				// SRLI
				c.Registers[opcode_rd(opcode)] = c.ReadRegister(opcode_rs1(opcode)) >> shamt

			} else if distinguish == 0b0100000 {
				// SRAI
				// Go's shift operator implements arithmetic shifts if the left operand is a signed integer and logical shifts if it is an unsigned integer.
				c.Registers[opcode_rd(opcode)] = uint32(int32(c.ReadRegister(opcode_rs1(opcode))) >> shamt)

			} else {
				return ErrInvalidOpcode{}
			}

		default:
			return ErrInvalidOpcode{}
		}

	case 0b0110011:
		// ADD/SUB/SLL/SLT/SLTU/XOR/SRL/SRA/OR/AND
		// RV32I defines several arithmetic R-type operations. All operations read the rs1 and rs2 registers
		// as source operands and write the result into register rd. The funct7 and funct3 fields select the
		// type of operation.

		funct3 := (opcode >> 12) & 0b111
		funct7 := (opcode >> 25) & 0b1111111

		switch funct3 {
		case 0b000:
			// ADD/SUB
			// ADD performs the addition of rs1 and rs2. SUB performs the subtraction of rs2 from rs1. Overflows
			// are ignored and the low XLEN bits of results are written to the destination rd.

			if funct7 == 0b0000000 {
				c.Registers[opcode_rd(opcode)] = c.Registers[opcode_rs1(opcode)] + c.Registers[opcode_rs2(opcode)]

			} else if funct7 == 0b0100000 {
				c.Registers[opcode_rd(opcode)] = c.Registers[opcode_rs1(opcode)] - c.Registers[opcode_rs2(opcode)]

			} else {
				return ErrInvalidOpcode{}
			}

		case 0b111:
			// AND
			// AND, OR, and XOR perform bitwise logical operations.
			if funct7 != 0b0000000 {
				return ErrInvalidOpcode{}
			}

			c.Registers[opcode_rd(opcode)] = c.Registers[opcode_rs1(opcode)] & c.Registers[opcode_rs2(opcode)]

		case 0b110:
			// OR
			if funct7 != 0b0000000 {
				return ErrInvalidOpcode{}
			}

			c.Registers[opcode_rd(opcode)] = c.Registers[opcode_rs1(opcode)] | c.Registers[opcode_rs2(opcode)]

		case 0b100:
			// XOR
			if funct7 != 0b0000000 {
				return ErrInvalidOpcode{}
			}

			c.Registers[opcode_rd(opcode)] = c.Registers[opcode_rs1(opcode)] ^ c.Registers[opcode_rs2(opcode)]

		case 0b010:
			// SLT
			if funct7 != 0b0000000 {
				return ErrInvalidOpcode{}
			}

			// SLT and SLTU perform signed and unsigned compares respectively, writing 1 to rd if rs1 < rs2,
			// 0 otherwise. Note, SLTU rd, x0, rs2 sets rd to 1 if rs2 is not equal to zero, otherwise sets rd
			// to zero (assembler pseudoinstruction SNEZ rd, rs).
			if int32(c.Registers[opcode_rs1(opcode)]) < int32(c.Registers[opcode_rs2(opcode)]) {
				c.Registers[opcode_rd(opcode)] = 1
			} else {
				c.Registers[opcode_rd(opcode)] = 0
			}

		case 0b011:
			// SLTU
			if funct7 != 0b0000000 {
				return ErrInvalidOpcode{}
			}

			if c.Registers[opcode_rs1(opcode)] < c.Registers[opcode_rs2(opcode)] {
				c.Registers[opcode_rd(opcode)] = 1
			} else {
				c.Registers[opcode_rd(opcode)] = 0
			}

		case 0b001:
			// SLL
			if funct7 != 0b0000000 {
				return ErrInvalidOpcode{}
			}

			// SLL, SRL, and SRA perform logical left, logical right, and arithmetic right shifts on the value in
			// register rs1 by the shift amount held in the lower 5 bits of register rs2.
			c.Registers[opcode_rd(opcode)] = c.Registers[opcode_rs1(opcode)] << (c.Registers[opcode_rs2(opcode)] & 0b11111)

		case 0b101:
			// SRL/SRA
			if funct7 == 0b0000000 {
				// SRL
				c.Registers[opcode_rd(opcode)] = c.Registers[opcode_rs1(opcode)] >> (c.Registers[opcode_rs2(opcode)] & 0b11111)

			} else if funct7 == 0100000 {
				// SRA
				c.Registers[opcode_rd(opcode)] = uint32(int32(c.Registers[opcode_rs1(opcode)]) >> (c.Registers[opcode_rs2(opcode)] & 0b11111))

			} else {
				return ErrInvalidOpcode{}
			}

		default:
			return ErrInvalidOpcode{}
		}

	case 0b0001111:
		// FENCE/FENCE.TSO/PAUSE
		funct3 := (opcode >> 12) & 0b111
		if funct3 != 0b000 {
			return ErrInvalidOpcode{}
		}

		if opcode == 0b00000001000000000000000000001111 {
			// Zihintpause PAUSE extension
			// PAUSE is encoded as a FENCE instruction with pred=W, succ=0, fm=0, rd=x0, and rs1=x0.
			// PAUSE is encoded as a hint within the FENCE opcode because some implementations are ex-
			// pected to deliberately stall the PAUSE instruction until outstanding memory transactions have
			// completed. Because the successor set is null, however, PAUSE does not mandate any particular
			// memory ordering - hence, it truly is a HINT.
		}

		err := c.w.MemFence(opcode)
		if err != nil {
			return err
		}

	case 0b1110011:
		// ECALL/EBREAK
		if opcode == 0b00000000000000000000000001110011 {
			// ECALL
			err := c.w.Syscall()
			if err != nil {
				return err
			}

		} else if opcode == 0b00000000000100000000000001110011 {
			// EBREAK
			return ErrBreak{}

		} else {
			return ErrInvalidOpcode{}
		}

	default:
		return ErrInvalidOpcode{}
	}

	// Step program-counter forward
	c.Pc += 4

	return nil
}