from typing import Any, Callable, Dict, Tuple, TYPE_CHECKING

if TYPE_CHECKING:
    from ....interpreter.program import INT_Program

from ...memory import Memory
from ....common.status import CONTINUE
from ....common.utils import b, b_inv, chi, compare, compare_bits_vectorized, z, z_inv, clamp_12, clamp_4, clamp_4_max0
from ....core.instruction_table import InstructionTable
from ....core.opcode import OpCode, OpReturn


class InstructionsWArgs2Reg1Imm(InstructionTable):
    def __init__(self, counter: int, program: "INT_Program", skip_index: int) -> None:
        self.counter = counter
        self.program = program
        self.skip_index = skip_index

    def get_props(self) -> list[int]:
        # Slice zeta once for better performance with large arrays
        # Most instructions need at most 6-8 bytes, so slice a reasonable window
        zeta_slice = self.program.zeta[self.counter + 1:self.counter + 8]
        
        # Cache the byte value and use bit operations for faster parsing
        byte_val = zeta_slice[0]  # equivalent to self.program.zeta[self.counter + 1]
        ra = clamp_12(byte_val & 0x0F)  # Lower 4 bits (equivalent to % 16)
        rb = clamp_12(byte_val >> 4)    # Upper 4 bits (equivalent to // 16)
        lx = clamp_4_max0(self.skip_index - 1)
        
        # Use sliced array for immediate value
        if lx > 0:
            imm_slice = zeta_slice[1:1+lx]  # equivalent to self.program.zeta[start:end]
            vx = chi(int.from_bytes(imm_slice, "little"), lx)
        else:
            vx = 0
        return [ra, rb, lx, vx]

    @classmethod
    def table(cls) -> Dict[int, OpCode]:
        return {
            120: OpCode(
                name="store_ind_u8", fn=cls.store_ind(8), gas=1, is_terminating=False
            ),
            121: OpCode(
                name="store_ind_u16", fn=cls.store_ind(16), gas=1, is_terminating=False
            ),
            122: OpCode(
                name="store_ind_u32", fn=cls.store_ind(32), gas=1, is_terminating=False
            ),
            123: OpCode(
                name="store_ind_u64", fn=cls.store_ind(64), gas=1, is_terminating=False
            ),
            124: OpCode(
                name="load_ind_u8", fn=cls.load_ind(8), gas=1, is_terminating=False
            ),
            125: OpCode(
                name="load_ind_i8",
                fn=cls.load_ind(8, True),
                gas=1,
                is_terminating=False,
            ),
            126: OpCode(
                name="load_ind_u16", fn=cls.load_ind(16), gas=1, is_terminating=False
            ),
            127: OpCode(
                name="load_ind_i16",
                fn=cls.load_ind(16, True),
                gas=1,
                is_terminating=False,
            ),
            128: OpCode(
                name="load_ind_u32", fn=cls.load_ind(32), gas=1, is_terminating=False
            ),
            129: OpCode(
                name="load_ind_i32",
                fn=cls.load_ind(32, True),
                gas=1,
                is_terminating=False,
            ),
            130: OpCode(
                name="load_ind_u64", fn=cls.load_ind(64), gas=1, is_terminating=False
            ),
            131: OpCode(
                name="add_imm_32", fn=cls.add_imm(32), gas=1, is_terminating=False
            ),
            132: OpCode(
                name="and_imm", fn=cls.op_imm("and"), gas=1, is_terminating=False
            ),
            133: OpCode(
                name="xor_imm", fn=cls.op_imm("xor"), gas=1, is_terminating=False
            ),
            134: OpCode(
                name="or_imm", fn=cls.op_imm("or"), gas=1, is_terminating=False
            ),
            135: OpCode(
                name="mul_imm_32", fn=cls.mul_imm(32), gas=1, is_terminating=False
            ),
            136: OpCode(
                name="set_lt_u_imm", fn=cls.set_lt_u_imm, gas=1, is_terminating=False
            ),
            137: OpCode(
                name="set_lt_s_imm", fn=cls.set_lt_s_imm, gas=1, is_terminating=False
            ),
            138: OpCode(
                name="shlo_l_imm_32", fn=cls.shlo_l_imm(32), gas=1, is_terminating=False
            ),
            139: OpCode(
                name="shlo_r_imm_32", fn=cls.shlo_r_imm(32), gas=1, is_terminating=False
            ),
            140: OpCode(
                name="shar_r_imm_32", fn=cls.shar_r_imm(32), gas=1, is_terminating=False
            ),
            141: OpCode(
                name="neg_add_imm_32",
                fn=cls.neg_add_imm(32),
                gas=1,
                is_terminating=False,
            ),
            142: OpCode(
                name="set_gt_u_imm", fn=cls.set_gt_u_imm, gas=1, is_terminating=False
            ),
            143: OpCode(
                name="set_gt_s_imm", fn=cls.set_gt_s_imm, gas=1, is_terminating=False
            ),
            144: OpCode(
                name="shlo_l_imm_alt_32",
                fn=cls.shlo_l_imm(32, True),
                gas=1,
                is_terminating=False,
            ),
            145: OpCode(
                name="shlo_r_imm_alt_32",
                fn=cls.shlo_r_imm(32, True),
                gas=1,
                is_terminating=False,
            ),
            146: OpCode(
                name="shar_r_imm_alt_32",
                fn=cls.shar_r_imm(32, True),
                gas=1,
                is_terminating=False,
            ),
            147: OpCode(
                name="cmov_iz_imm", fn=cls.cmov_iz_imm, gas=1, is_terminating=False
            ),
            148: OpCode(
                name="cmov_nz_imm", fn=cls.cmov_nz_imm, gas=1, is_terminating=False
            ),
            149: OpCode(
                name="add_imm_64", fn=cls.add_imm(64), gas=1, is_terminating=False
            ),
            150: OpCode(
                name="mul_imm_64", fn=cls.mul_imm(64), gas=1, is_terminating=False
            ),
            151: OpCode(
                name="shlo_l_imm_64", fn=cls.shlo_l_imm(64), gas=1, is_terminating=False
            ),
            152: OpCode(
                name="shlo_r_imm_64", fn=cls.shlo_r_imm(64), gas=1, is_terminating=False
            ),
            153: OpCode(
                name="shar_r_imm_64", fn=cls.shar_r_imm(64), gas=1, is_terminating=False
            ),
            154: OpCode(
                name="neg_add_imm_64",
                fn=cls.neg_add_imm(64),
                gas=1,
                is_terminating=False,
            ),
            155: OpCode(
                name="shlo_l_imm_alt_64",
                fn=cls.shlo_l_imm(64, True),
                gas=1,
                is_terminating=False,
            ),
            156: OpCode(
                name="shlo_r_imm_alt_64",
                fn=cls.shlo_r_imm(64, True),
                gas=1,
                is_terminating=False,
            ),
            157: OpCode(
                name="shar_r_imm_alt_64",
                fn=cls.shar_r_imm(64, True),
                gas=1,
                is_terminating=False,
            ),
            158: OpCode(
                name="rot_r_64_imm", fn=cls.rot_imm(64), gas=1, is_terminating=False
            ),
            159: OpCode(
                name="rot_r_64_imm_alt",
                fn=cls.rot_imm(64, True),
                gas=1,
                is_terminating=False,
            ),
            160: OpCode(
                name="rot_r_32_imm", fn=cls.rot_imm(32), gas=1, is_terminating=False
            ),
            161: OpCode(
                name="rot_r_32_imm_alt",
                fn=cls.rot_imm(32, True),
                gas=1,
                is_terminating=False,
            ),
        }

    @staticmethod
    def op_imm(op: str) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def op_imm_impl(self: "InstructionsWArgs2Reg1Imm", registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            wb_bits = b(registers[rb], 8)
            vx_bits = b(vx, 8)
            # Use vectorized comparison for massive performance boost
            result_bits = compare_bits_vectorized(wb_bits, vx_bits, op)
            registers[ra] = b_inv(result_bits)

            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return op_imm_impl

    def set_lt_u_imm(self, registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
        registers[ra] = int(registers[rb] < vx)
        return CONTINUE, self.counter + self.skip_index + 1, registers, memory

    def set_lt_s_imm(self, registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
        registers[ra] = int(z(registers[rb], 8) < z(vx, 8))
        return CONTINUE, self.counter + self.skip_index + 1, registers, memory

    @staticmethod
    def shlo_l_imm(bitsize: int, alt: bool = False) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def shlo_l_imm_impl(self: InstructionsWArgs2Reg1Imm, registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            a = int(registers[rb]) if not alt else vx
            b = vx if not alt else int(registers[rb])

            registers[ra] = chi(
                (a * 2 ** (b % bitsize)) % 2**bitsize, bitsize // 8
            )
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return shlo_l_imm_impl

    @staticmethod
    def shlo_r_imm(bitsize: int, alt: bool = False) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def shlo_r_imm_impl(self: InstructionsWArgs2Reg1Imm, registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            a = int(registers[rb]) if not alt else vx
            b = vx if not alt else int(registers[rb])

            registers[ra] = chi(
                ((a % 2 ** (bitsize)) // 2 ** (b % bitsize)) % 2 ** (bitsize),
                bitsize // 8,
            )
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return shlo_r_imm_impl

    @staticmethod
    def shar_r_imm(bitsize: int, alt: bool = False) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def shar_r_imm_impl(self: "InstructionsWArgs2Reg1Imm", registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            a = int(registers[rb]) if not alt else vx
            b = vx if not alt else int(registers[rb])

            registers[ra] = z_inv(
                z(a % 2 ** (bitsize), bitsize // 8) // 2 ** (b % bitsize), 8
            )
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return shar_r_imm_impl

    @staticmethod
    def neg_add_imm(bitsize: int) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def neg_add_imm_impl(self: "InstructionsWArgs2Reg1Imm", registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            value = (vx + 2**bitsize - int(registers[rb])) % 2**bitsize
            if bitsize < 64:
                value = chi(value, 4)
            registers[ra] = value
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return neg_add_imm_impl

    @staticmethod
    def add_imm(bitsize: int) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def add_imm_impl(self: "InstructionsWArgs2Reg1Imm", registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            value = (registers[rb] + vx) % 2 ** (bitsize)
            if bitsize < 64:
                value = chi(value, bitsize // 8)
            registers[ra] = value
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return add_imm_impl

    @staticmethod
    def mul_imm(bitsize: int) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def mul_imm_impl(self: "InstructionsWArgs2Reg1Imm", registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            value = (registers[rb] * vx) % 2 ** (bitsize)
            if bitsize < 64:
                value = chi(value, 4)
            registers[ra] = value
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return mul_imm_impl

    @staticmethod
    def rot_imm(bitsize: int, alt: bool = False) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def rot_imm_impl(self: "InstructionsWArgs2Reg1Imm", registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            a_val = int(registers[rb] if not alt else vx) % 2 ** (bitsize)
            b_val = int(vx if not alt else registers[rb]) % 2 ** (bitsize)

            a_bits = b(a_val, bitsize // 8)
            x = b_inv([a_bits[(i + b_val) % bitsize] for i in range(bitsize)])

            if bitsize < 64:
                x = chi(x, bitsize // 8)

            registers[ra] = x
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return rot_imm_impl

    def set_gt_u_imm(self, registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
        registers[ra] = int(registers[rb] > vx)
        return CONTINUE, self.counter + self.skip_index + 1, registers, memory

    def set_gt_s_imm(self, registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
        registers[ra] = int(z(registers[rb], 8) > z(vx, 8))
        return CONTINUE, self.counter + self.skip_index + 1, registers, memory

    def cmov_iz_imm(self, registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
        if registers[rb] == 0:
            registers[ra] = vx
        return CONTINUE, self.counter + self.skip_index + 1, registers, memory

    def cmov_nz_imm(self, registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
        if registers[rb] != 0:
            registers[ra] = vx
        return CONTINUE, self.counter + self.skip_index + 1, registers, memory

    @staticmethod
    def store_ind(bitsize: int) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def store_ind_impl(self: "InstructionsWArgs2Reg1Imm", registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            memory.write(
                registers[rb] + vx,
                int(registers[ra] % 2**bitsize).to_bytes(bitsize // 8, "little"),
            )
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return store_ind_impl

    @staticmethod
    def load_ind(bitsize: int, signed: bool = False) -> Callable[["InstructionsWArgs2Reg1Imm", list[int], Memory, int, int, int, int], OpReturn]:
        def load_ind_impl(self: "InstructionsWArgs2Reg1Imm", registers: list[int], memory: Memory, ra: int, rb: int, lx: int, vx: int) -> OpReturn:
            value = int.from_bytes(
                memory.read(registers[rb] + vx, bitsize // 8), "little"
            )
            if signed:
                value = z_inv(z(value, bitsize // 8), 8)

            registers[ra] = value
            return CONTINUE, self.counter + self.skip_index + 1, registers, memory

        return load_ind_impl
