import re
import sys
import struct

class Y86_64Simulator:
    # Y86-64 Register IDs
    REG_RAX = 0
    REG_RCX = 1
    REG_RDX = 2
    REG_RBX = 3
    REG_RSP = 4
    REG_RBP = 5
    REG_RSI = 6
    REG_RDI = 7
    REG_R8 = 8
    REG_R9 = 9
    REG_R10 = 10
    REG_R11 = 11
    REG_R12 = 12
    REG_R13 = 13
    REG_R14 = 14
    REG_NONE = 15  # Used for instructions with no register operand

    # Status codes
    STAT_AOK = 1  # Normal operation
    STAT_HLT = 2  # Halt instruction encountered
    STAT_ADR = 3  # Invalid address
    STAT_INS = 4  # Invalid instruction
    
    # Condition codes
    CC_ZF = 0  # Zero flag
    CC_SF = 1  # Sign flag
    CC_OF = 2  # Overflow flag

    def __init__(self, mem_size=0x10000):
        # Initialize memory, registers, program counter and status
        self.memory = bytearray(mem_size)
        self.mem_size = mem_size
        self.registers = [0] * 15  # 15 registers (RAX to R14)
        self.pc = 0
        self.status = self.STAT_AOK
        
        # Condition code flags
        self.cc = [False] * 3  # ZF, SF, OF
        
        # Register name to ID mapping
        self.reg_name_to_id = {
            'rax': self.REG_RAX, '%rax': self.REG_RAX,
            'rcx': self.REG_RCX, '%rcx': self.REG_RCX,
            'rdx': self.REG_RDX, '%rdx': self.REG_RDX,
            'rbx': self.REG_RBX, '%rbx': self.REG_RBX,
            'rsp': self.REG_RSP, '%rsp': self.REG_RSP,
            'rbp': self.REG_RBP, '%rbp': self.REG_RBP,
            'rsi': self.REG_RSI, '%rsi': self.REG_RSI,
            'rdi': self.REG_RDI, '%rdi': self.REG_RDI,
            'r8': self.REG_R8, '%r8': self.REG_R8,
            'r9': self.REG_R9, '%r9': self.REG_R9,
            'r10': self.REG_R10, '%r10': self.REG_R10,
            'r11': self.REG_R11, '%r11': self.REG_R11,
            'r12': self.REG_R12, '%r12': self.REG_R12,
            'r13': self.REG_R13, '%r13': self.REG_R13,
            'r14': self.REG_R14, '%r14': self.REG_R14
        }
    
    def load_coe_file(self, filename):
        """
        Load a program from a Xilinx COE format file.
        COE file format typically has a header section followed by hex values.
        """
        try:
            with open(filename, 'r') as f:
                content = f.read()
            
            # Find memory initialization section
            memory_init = re.search(r'memory_initialization_vector\s*=\s*(.*?)(?:;|\Z)', content, re.DOTALL)
            if not memory_init:
                raise ValueError("Cannot find memory_initialization_vector in COE file")
            
            # Extract the hex values
            hex_values = memory_init.group(1).replace('\n', '').replace(' ', '').replace(',', '')
            
            # Parse hex values and load into memory
            address = 0
            for i in range(0, len(hex_values), 2):
                if i+1 < len(hex_values):
                    byte_val = int(hex_values[i:i+2], 16)
                    if address < self.mem_size:
                        self.memory[address] = byte_val
                        address += 1
                    else:
                        print(f"Warning: Address {address} exceeds memory size, truncating program")
                        break
                        
            print(f"Loaded {address} bytes from {filename}")
            return True
        except Exception as e:
            print(f"Error loading COE file: {e}")
            return False
    
    def read_byte(self, addr):
        """Read a byte from memory at address addr"""
        if 0 <= addr < self.mem_size:
            return self.memory[addr]
        else:
            self.status = self.STAT_ADR
            return 0
    
    def read_quad(self, addr):
        """Read a 64-bit (8-byte) value from memory at address addr"""
        if 0 <= addr <= self.mem_size - 8:
            return int.from_bytes(self.memory[addr:addr+8], byteorder='little')
        else:
            self.status = self.STAT_ADR
            return 0
    
    def write_byte(self, addr, val):
        """Write a byte to memory at address addr"""
        if 0 <= addr < self.mem_size:
            self.memory[addr] = val & 0xFF
        else:
            self.status = self.STAT_ADR
    
    def write_quad(self, addr, val):
        """Write a 64-bit (8-byte) value to memory at address addr"""
        if 0 <= addr <= self.mem_size - 8:
            val_bytes = val.to_bytes(8, byteorder='little')
            self.memory[addr:addr+8] = val_bytes
        else:
            self.status = self.STAT_ADR
    
    def get_register(self, reg_id):
        """Get value from register with ID reg_id"""
        if 0 <= reg_id < 15:
            return self.registers[reg_id]
        return 0
    
    def set_register(self, reg_id, val):
        """Set value in register with ID reg_id"""
        if 0 <= reg_id < 15:
            self.registers[reg_id] = val & 0xFFFFFFFFFFFFFFFF  # Ensure 64-bit value
    
    def set_cc(self, result):
        """Set condition codes based on result"""
        # Zero flag
        self.cc[self.CC_ZF] = (result == 0)
        
        # Sign flag (negative)
        self.cc[self.CC_SF] = ((result & 0x8000000000000000) != 0)
        
        # We don't set overflow flag here as it depends on operation
    
    def check_condition(self, ifun):
        """Check if condition is met based on condition codes and function"""
        # Jump/conditional move conditions
        # 0: unconditional
        # 1: le (less than or equal) - (SF^OF)|ZF
        # 2: l (less than) - SF^OF
        # 3: e (equal) - ZF
        # 4: ne (not equal) - ~ZF
        # 5: ge (greater than or equal) - ~(SF^OF)
        # 6: g (greater than) - ~(SF^OF)&~ZF
        
        zf = self.cc[self.CC_ZF]
        sf = self.cc[self.CC_SF]
        of = self.cc[self.CC_OF]
        
        if ifun == 0:    # unconditional
            return True
        elif ifun == 1:  # le
            return (sf != of) or zf
        elif ifun == 2:  # l
            return sf != of
        elif ifun == 3:  # e
            return zf
        elif ifun == 4:  # ne
            return not zf
        elif ifun == 5:  # ge
            return sf == of
        elif ifun == 6:  # g
            return (sf == of) and (not zf)
        else:
            return False
    
    def fetch_decode_execute(self):
        """Fetch, decode and execute a single instruction"""
        if self.status != self.STAT_AOK:
            return False
        
        # Fetch
        instr_byte = self.read_byte(self.pc)
        if self.status != self.STAT_AOK:
            return False
        
        # Decode instruction
        icode = (instr_byte >> 4) & 0xF
        ifun = instr_byte & 0xF
        
        # Increment PC (will be updated further based on instruction size)
        self.pc += 1
        
        # Initialize variables for the instruction
        valA = 0
        valB = 0
        valC = 0
        valE = 0
        valM = 0
        dstE = self.REG_NONE
        dstM = self.REG_NONE
        srcA = self.REG_NONE
        srcB = self.REG_NONE
        
        # Process based on instruction code
        if icode == 0:  # HALT
            self.status = self.STAT_HLT
            return False
            
        elif icode == 1:  # NOP
            pass  # No operation
            
        elif icode == 2:  # rrmovq or conditional move
            # Get register byte
            reg_byte = self.read_byte(self.pc)
            rA = (reg_byte >> 4) & 0xF
            rB = reg_byte & 0xF
            self.pc += 1
            
            # Check condition
            if self.check_condition(ifun):
                # Move value from rA to rB
                valA = self.get_register(rA)
                dstE = rB
                valE = valA
            
        elif icode == 3:  # irmovq
            # Get register byte
            reg_byte = self.read_byte(self.pc)
            rA = (reg_byte >> 4) & 0xF  # Should be 0xF (no register)
            rB = reg_byte & 0xF
            self.pc += 1
            
            # Get immediate value (8 bytes)
            valC = self.read_quad(self.pc)
            self.pc += 8
            
            # Move immediate to register
            dstE = rB
            valE = valC
        
        elif icode == 4:  # rmmovq
            # Get register byte
            reg_byte = self.read_byte(self.pc)
            rA = (reg_byte >> 4) & 0xF
            rB = reg_byte & 0xF
            self.pc += 1
            
            # Get displacement (8 bytes)
            valC = self.read_quad(self.pc)
            self.pc += 8
            
            # Calculate memory address
            valB = self.get_register(rB)
            valE = valB + valC
            
            # Move from register to memory
            valA = self.get_register(rA)
            self.write_quad(valE, valA)
        
        elif icode == 5:  # mrmovq
            # Get register byte
            reg_byte = self.read_byte(self.pc)
            rA = (reg_byte >> 4) & 0xF
            rB = reg_byte & 0xF
            self.pc += 1
            
            # Get displacement (8 bytes)
            valC = self.read_quad(self.pc)
            self.pc += 8
            
            # Calculate memory address
            valB = self.get_register(rB)
            valE = valB + valC
            
            # Move from memory to register
            valM = self.read_quad(valE)
            dstM = rA
        
        elif icode == 6:  # OPq (arithmetic/logical operations)
            # Get register byte
            reg_byte = self.read_byte(self.pc)
            rA = (reg_byte >> 4) & 0xF
            rB = reg_byte & 0xF
            self.pc += 1
            
            # Get operands
            valA = self.get_register(rA)
            valB = self.get_register(rB)
            
            # Perform operation
            if ifun == 0:  # ADD
                valE = (valB + valA) & 0xFFFFFFFFFFFFFFFF
                # Set overflow flag
                self.cc[self.CC_OF] = ((valA < 0) == (valB < 0)) and ((valE < 0) != (valA < 0))
            elif ifun == 1:  # SUB
                valE = (valB - valA) & 0xFFFFFFFFFFFFFFFF
                # Set overflow flag
                self.cc[self.CC_OF] = ((valA < 0) != (valB < 0)) and ((valE < 0) != (valB < 0))
            elif ifun == 2:  # AND
                valE = valB & valA
                self.cc[self.CC_OF] = False
            elif ifun == 3:  # XOR
                valE = valB ^ valA
                self.cc[self.CC_OF] = False
            else:
                self.status = self.STAT_INS
                return False
            
            # Set other condition codes
            self.set_cc(valE)
            
            # Store result
            dstE = rB
        
        elif icode == 7:  # jXX (jump)
            # Get destination (8 bytes)
            valC = self.read_quad(self.pc)
            self.pc += 8
            
            # Check condition
            if self.check_condition(ifun):
                # Jump
                self.pc = valC
        
        elif icode == 8:  # call
            # Get destination (8 bytes)
            valC = self.read_quad(self.pc)
            self.pc += 8
            
            # Push return address onto stack
            valB = self.get_register(self.REG_RSP)
            valE = valB - 8
            self.set_register(self.REG_RSP, valE)
            self.write_quad(valE, self.pc)
            
            # Jump to function
            self.pc = valC
        
        elif icode == 9:  # ret
            # Pop return address from stack
            valA = self.get_register(self.REG_RSP)
            valB = valA
            valE = valB + 8
            
            valM = self.read_quad(valA)
            
            # Update stack pointer
            self.set_register(self.REG_RSP, valE)
            
            # Jump to return address
            self.pc = valM
        
        elif icode == 10:  # pushq
            # Get register byte
            reg_byte = self.read_byte(self.pc)
            rA = (reg_byte >> 4) & 0xF
            rB = reg_byte & 0xF  # Should be 0xF (no register)
            self.pc += 1
            
            # Get value to push
            valA = self.get_register(rA)
            
            # Decrement stack pointer
            valB = self.get_register(self.REG_RSP)
            valE = valB - 8
            
            # Push value onto stack
            self.set_register(self.REG_RSP, valE)
            self.write_quad(valE, valA)
        
        elif icode == 11:  # popq
            # Get register byte
            reg_byte = self.read_byte(self.pc)
            rA = (reg_byte >> 4) & 0xF
            rB = reg_byte & 0xF  # Should be 0xF (no register)
            self.pc += 1
            
            # Pop value from stack
            valA = self.get_register(self.REG_RSP)
            valB = valA
            valE = valB + 8
            
            valM = self.read_quad(valA)
            
            # Update registers
            dstM = rA
            self.set_register(self.REG_RSP, valE)
        
        else:
            self.status = self.STAT_INS
            return False
        
        # Write results to registers
        if dstE != self.REG_NONE:
            self.set_register(dstE, valE)
        
        if dstM != self.REG_NONE:
            self.set_register(dstM, valM)
        
        return self.status == self.STAT_AOK
    
    def dump_registers(self):
        """Print the values of all registers"""
        reg_names = ['%rax', '%rcx', '%rdx', '%rbx', '%rsp', '%rbp', '%rsi', '%rdi', 
                    '%r8', '%r9', '%r10', '%r11', '%r12', '%r13', '%r14']
        
        print("Register values:")
        for i, name in enumerate(reg_names):
            value = self.registers[i]
            print(f"{name} = 0x{value:016x} ({value})")
        
        print(f"PC = 0x{self.pc:x}")
        print(f"ZF = {1 if self.cc[self.CC_ZF] else 0}, SF = {1 if self.cc[self.CC_SF] else 0}, OF = {1 if self.cc[self.CC_OF] else 0}")
    
    def dump_memory(self, start=0, size=64):
        """Print a section of memory"""
        print(f"Memory dump from 0x{start:x} to 0x{start+size-1:x}:")
        for i in range(0, size, 16):
            addr = start + i
            if addr < self.mem_size:
                line = f"0x{addr:04x}: "
                for j in range(16):
                    if addr + j < self.mem_size:
                        line += f"{self.memory[addr+j]:02x} "
                    else:
                        line += "   "
                    if j == 7:
                        line += " "
                
                # ASCII representation
                line += " |"
                for j in range(16):
                    if addr + j < self.mem_size:
                        char = self.memory[addr+j]
                        if 32 <= char <= 126:  # Printable ASCII
                            line += chr(char)
                        else:
                            line += "."
                    else:
                        line += " "
                line += "|"
                print(line)
    
    def run(self, max_instructions=1000000):
        """Run the program until halt or error, or until max_instructions is reached"""
        instructions_executed = 0
        
        while self.status == self.STAT_AOK and instructions_executed < max_instructions:
            if not self.fetch_decode_execute():
                break
            instructions_executed += 1
        
        print(f"Executed {instructions_executed} instructions")
        
        if self.status == self.STAT_HLT:
            print("Program halted normally")
        elif self.status == self.STAT_ADR:
            print("Error: Invalid memory address")
        elif self.status == self.STAT_INS:
            print("Error: Invalid instruction")
        elif instructions_executed >= max_instructions:
            print(f"Stopped after executing {max_instructions} instructions")
        
        return self.status

def main():
    if len(sys.argv) < 2:
        print("Usage: python y86_64_simulator.py <coe_file>")
        return
    
    simulator = Y86_64Simulator()
    if simulator.load_coe_file(sys.argv[1]):
        simulator.run()
        simulator.dump_registers()
        simulator.dump_memory(0, 128)  # Dump first 128 bytes of memory

if __name__ == "__main__":
    main()
