"""Type stubs for tsrkit_asm module - Complete API"""

from typing import Optional, List, Union
from enum import Enum

class Reg(Enum):
    rax = 0
    rcx = 1
    rdx = 2
    rbx = 3
    rsp = 4
    rbp = 5
    rsi = 6
    rdi = 7
    r8 = 8
    r9 = 9
    r10 = 10
    r11 = 11
    r12 = 12
    r13 = 13
    r14 = 14
    r15 = 15

class ImmKind(Enum):
    I8 = 0
    I16 = 1
    I32 = 2
    I64 = 3

class LoadKind(Enum):
    U8 = 0
    U16 = 1
    U32 = 2
    U64 = 3
    I8 = 4
    I16 = 5
    I32 = 6

class RegSize(Enum):
    R32 = 0
    R64 = 1

class Size(Enum):
    U8 = 0
    U16 = 1
    U32 = 2
    U64 = 3

class Condition(Enum):
    Overflow = 0
    NotOverflow = 1
    Below = 2
    AboveOrEqual = 3
    Equal = 4
    NotEqual = 5
    BelowOrEqual = 6
    Above = 7
    Sign = 8
    NotSign = 9
    Parity = 10
    NotParity = 11
    Less = 12
    GreaterOrEqual = 13
    LessOrEqual = 14
    Greater = 15

class RegIndex(Enum):
    rax = 0
    rcx = 1
    rdx = 2
    rbx = 3
    rbp = 5
    rsi = 6
    rdi = 7
    r8 = 8
    r9 = 9
    r10 = 10
    r11 = 11
    r12 = 12
    r13 = 13
    r14 = 14
    r15 = 15

class SegReg(Enum):
    fs = 0
    gs = 1

class Scale(Enum):
    x1 = 0
    x2 = 1
    x4 = 2
    x8 = 3

class MemOpBaseOffset:
    def __init__(self, seg: Optional[SegReg], size: RegSize, base: Reg, offset: int) -> None: ...

class MemOpBaseIndexScaleOffset:
    def __init__(self, seg: Optional[SegReg], size: RegSize, base: Reg, index: RegIndex, scale: Scale, offset: int) -> None: ...

class MemOpIndexScaleOffset:
    def __init__(self, seg: Optional[SegReg], size: RegSize, index: RegIndex, scale: Scale, offset: int) -> None: ...

class MemOpOffset:
    def __init__(self, seg: Optional[SegReg], size: RegSize, offset: int) -> None: ...

class MemOpRipRelative:
    def __init__(self, seg: Optional[SegReg], offset: int) -> None: ...

MemOp = Union[MemOpBaseOffset, MemOpBaseIndexScaleOffset, MemOpIndexScaleOffset, MemOpOffset, MemOpRipRelative]

class RegMemReg:
    def __init__(self, reg: Reg) -> None: ...

class RegMemMem:
    def __init__(self, mem: MemOp) -> None: ...

RegMem = Union[RegMemReg, RegMemMem]

class OperandsRegMemReg:
    def __init__(self, size: Size, reg_mem: RegMem, reg: Reg) -> None: ...

class OperandsRegRegMem:
    def __init__(self, size: Size, reg: Reg, reg_mem: RegMem) -> None: ...

class OperandsRegMemImm:
    def __init__(self, reg_mem: RegMem, imm: ImmKind) -> None: ...

Operands = Union[OperandsRegMemReg, OperandsRegRegMem, OperandsRegMemImm]


class PyAssembler:
    """Python wrapper for the PolkaVM assembler with complete x86-64 instruction set support"""
    
    def __init__(self) -> None: ...
    
    # ========== ASSEMBLER MANAGEMENT ==========
    
    # Label management
    def forward_declare_label(self) -> int: ...
    def create_label(self) -> int: ...
    def define_label(self, raw: int) -> None: ...
    def get_label_origin_offset(self, label: int) -> Optional[int]: ...
    
    # Origin and address management
    def origin(self) -> int: ...
    def set_origin(self, o: int) -> None: ...
    def current_address(self) -> int: ...
    def len(self) -> int: ...
    def is_empty(self) -> bool: ...
    def spare_capacity(self) -> int: ...
    
    # Memory management
    def reserve_code(self, length: int) -> None: ...
    def reserve_labels(self, length: int) -> None: ...
    def resize(self, size: int, fill_with: int) -> None: ...
    
    # Raw bytes
    def push_raw(self, b: bytes) -> None: ...
    def push_with_label(self, label: int, _opcode: str) -> None: ...
    
    # Finalization
    def finalize(self) -> bytes: ...
    def clear(self) -> None: ...
    def code_mut(self) -> bytes: ...
    
    # ========== BASIC INSTRUCTIONS ==========
    
    # No-operation variants
    def ud2(self) -> None: ...
    def endbr64(self) -> None: ...
    def syscall(self) -> None: ...
    def push_reg(self, reg: Reg) -> None: ...  # Alias for push
    def push(self, reg: Reg) -> None: ...
    def push_imm(self, val: int) -> None: ...
    def pop_reg(self, reg: Reg) -> None: ...  # Alias for pop
    def pop(self, reg: Reg) -> None: ...
    def nop(self) -> None: ...
    def nop2(self) -> None: ...
    def nop3(self) -> None: ...
    def nop4(self) -> None: ...
    def nop5(self) -> None: ...
    def nop6(self) -> None: ...
    def nop7(self) -> None: ...
    def nop8(self) -> None: ...
    def nop9(self) -> None: ...
    def nop10(self) -> None: ...
    def nop11(self) -> None: ...
    
    # Move instructions
    def mov_imm64(self, reg: Reg, val: int) -> None: ...
    def mov_imm(self, mem: RegMem, imm: ImmKind) -> None: ...
    def mov_(self, size: RegSize, a: Reg, b: Reg) -> None: ...  # Alias for mov
    def mov(self, size: RegSize, a: Reg, b: Reg) -> None: ...
    
    # Move with sign/zero extension
    def movsx_8_to_64(self, reg_size: RegSize, dst: Reg, src: Reg) -> None: ...
    def movsx_16_to_64(self, reg_size: RegSize, dst: Reg, src: Reg) -> None: ...
    def movzx_16_to_64(self, reg_size: RegSize, dst: Reg, src: Reg) -> None: ...
    def movsxd_32_to_64(self, dst: Reg, src: Reg) -> None: ...
    
    # Return
    def ret(self) -> None: ...
    
    # Compare and test
    def cmp(self, val: Operands) -> None: ...
    def test(self, operands: Operands) -> None: ...
    
    # Memory operations
    def load(self, kind: LoadKind, reg: Reg, mem: MemOp) -> None: ...
    def store(self, size: Size, mem: MemOp, reg: Reg) -> None: ...
    
    # Cache operations
    def clflushopt(self, seg: Optional[int], reg_size: int, base: Reg, offset: int) -> None: ...

    # CPU state
    def cpuid(self) -> None: ...
    def rdtscp(self) -> None: ...
    def rdpmc(self) -> None: ...
    def cdq(self) -> None: ...
    def cqo(self) -> None: ...
    
    # Memory barriers
    def mfence(self) -> None: ...
    def lfence(self) -> None: ...

    # ========== ARITHMETIC OPERATIONS ==========
    def add(self, operands: Operands) -> None: ...
    def sub(self, operands: Operands) -> None: ...
    def inc(self, size: Size, reg_mem: RegMem) -> None: ...
    def dec(self, size: Size, reg_mem: RegMem) -> None: ...
    def neg(self, size: Size, reg_mem: RegMem) -> None: ...
    
    # Multiply operations
    def imul(self, reg_size: RegSize, dst: Reg, src: RegMem) -> None: ...
    def imul_imm(self, reg_size: RegSize, dst: Reg, src: RegMem, imm: int) -> None: ...
    def imul_dx_ax(self, reg_size: RegSize, src: RegMem) -> None: ...
    def mul(self, reg_size: RegSize, src: RegMem) -> None: ...
    def mul_dx_ax(self, reg_size: RegSize, src: RegMem) -> None: ...
    
    # Divide operations
    def div(self, reg_size: RegSize, src: RegMem) -> None: ...
    def idiv(self, reg_size: RegSize, src: RegMem) -> None: ...

    # ========== LOGICAL OPERATIONS ==========
    def and_(self, operands: Operands) -> None: ...
    def or_(self, operands: Operands) -> None: ...
    def xor(self, operands: Operands) -> None: ...
    def not_(self, size: Size, reg_mem: RegMem) -> None: ...

    # ========== SHIFT AND ROTATE OPERATIONS ==========
    def sar_cl(self, reg_size: RegSize, reg_mem: RegMem) -> None: ...
    def sar_imm(self, reg_size: RegSize, reg_mem: RegMem, imm: int) -> None: ...
    def shl_cl(self, reg_size: RegSize, reg_mem: RegMem) -> None: ...
    def shl_imm(self, reg_size: RegSize, reg_mem: RegMem, imm: int) -> None: ...
    def shr_cl(self, reg_size: RegSize, reg_mem: RegMem) -> None: ...
    def shr_imm(self, reg_size: RegSize, reg_mem: RegMem, imm: int) -> None: ...
    def rol_cl(self, reg_size: RegSize, reg_mem: RegMem) -> None: ...
    def ror_cl(self, reg_size: RegSize, reg_mem: RegMem) -> None: ...
    def ror_imm(self, reg_size: RegSize, reg_mem: RegMem, imm: int) -> None: ...
    def rcr_imm(self, reg_size: RegSize, reg_mem: RegMem, imm: int) -> None: ...

    # ========== BIT MANIPULATION ==========
    def bts(self, reg_size: RegSize, reg_mem: RegMem, bit: int) -> None: ...
    def popcnt(self, reg_size: RegSize, dst: Reg, src: RegMem) -> None: ...
    def lzcnt(self, reg_size: RegSize, dst: Reg, src: RegMem) -> None: ...
    def tzcnt(self, reg_size: RegSize, dst: Reg, src: RegMem) -> None: ...
    def bswap(self, reg_size: RegSize, reg: Reg) -> None: ...

    # ========== CONDITIONAL OPERATIONS ==========
    def cmov(self, condition: Condition, reg_size: RegSize, dst: Reg, src: RegMem) -> None: ...
    def setcc(self, condition: Condition, reg_mem: RegMem) -> None: ...

    # ========== LEA OPERATIONS ==========
    def lea(self, reg_size: RegSize, dst: Reg, mem: MemOp) -> None: ...
    def lea_rip_label(self, dst: Reg, label: int) -> None: ...

    # ========== EXCHANGE ==========
    def xchg_mem(self, reg_size: RegSize, reg: Reg, mem: MemOp) -> None: ...

    # ========== CONTROL FLOW ==========
    
    # Jump instructions
    def jmp_rel32(self, offset: int) -> None: ...
    def jmp_rel8(self, offset: int) -> None: ...
    def jmp_label32(self, label: int) -> None: ...
    def jmp_label8(self, label: int) -> None: ...
    def jmp(self, reg_mem: RegMem) -> None: ...
    
    # Call instructions
    def call_rel32(self, offset: int) -> None: ...
    def call_label32(self, label: int) -> None: ...
    def call(self, reg_mem: RegMem) -> None: ...
    
    # Conditional jump instructions
    def jcc_rel8(self, condition: Condition, offset: int) -> None: ...
    def jcc_rel32(self, condition: Condition, offset: int) -> None: ...
    def jcc_label8(self, condition: Condition, label: int) -> None: ...
    def jcc_label32(self, condition: Condition, label: int) -> None: ...
    def jcc_label32_default(self, condition: Condition, label: int, default_offset: int) -> None: ...

# ========== CONSTANTS ==========

AllRegs: List[Reg]