import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import re
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
from enum import Enum

class CPUMode(Enum):
    SIMPLE = "simple"
    X86 = "x86"
    ARM = "arm"

@dataclass
class CPUState:
    """Complete CPU state representation"""
    registers: Dict[str, int]
    memory: bytearray
    ip: int
    sp: int
    flags: Dict[str, bool]
    halted: bool
    mode: CPUMode
    
    def __init__(self, mem_size=65536, mode=CPUMode.SIMPLE):
        self.registers = {f"R{i}": 0 for i in range(8)}
        self.memory = bytearray(mem_size)
        self.ip = 0
        self.sp = mem_size - 4  # Stack starts at top
        self.flags = {"Z": False, "N": False, "C": False, "O": False}
        self.halted = False
        self.mode = mode
        self.bp = mem_size - 4  # Base pointer

class Breakpoint:
    def __init__(self, address=None, label=None, condition=None):
        self.address = address
        self.label = label
        self.condition = condition
        self.enabled = True

class Instruction:
    def __init__(self, opcode, operands, line_num, label=None, comment=None):
        self.opcode = opcode
        self.operands = operands
        self.line_num = line_num
        self.label = label
        self.comment = comment
        self.address = 0

class AssemblyDebugger:
    def __init__(self, mode=CPUMode.SIMPLE):
        self.cpu = CPUState(mode=mode)
        self.instructions: List[Instruction] = []
        self.labels: Dict[str, int] = {}
        self.breakpoints: List[Breakpoint] = []
        self.execution_history: List[str] = []
        self.source_lines: List[str] = []
        self.pipeline_enabled = False
        self.cache_enabled = False
        
    def reset(self):
        """Reset CPU to initial state"""
        mem_size = len(self.cpu.memory)
        mode = self.cpu.mode
        self.cpu = CPUState(mem_size, mode)
        self.execution_history.clear()
        
    def parse_assembly(self, code: str) -> Tuple[bool, str]:
        """Parse assembly code into instructions"""
        self.instructions.clear()
        self.labels.clear()
        self.source_lines = code.split('\n')
        
        address = 0
        for line_num, line in enumerate(self.source_lines, 1):
            line = line.strip()
            
            # Skip empty lines and comments
            if not line or line.startswith(';'):
                continue
                
            # Extract comment
            comment = None
            if ';' in line:
                line, comment = line.split(';', 1)
                line = line.strip()
                comment = comment.strip()
            
            # Check for label
            label = None
            if ':' in line:
                label, line = line.split(':', 1)
                label = label.strip()
                self.labels[label] = address
                line = line.strip()
                
            if not line:
                continue
                
            # Parse instruction
            parts = line.split(None, 1)
            if not parts:
                continue
                
            opcode = parts[0].upper()
            operands = []
            
            if len(parts) > 1:
                operands = [op.strip() for op in parts[1].split(',')]
            
            instr = Instruction(opcode, operands, line_num, label, comment)
            instr.address = address
            self.instructions.append(instr)
            
            # Estimate instruction size (simplified)
            address += 4
            
        return True, "Assembly parsed successfully"
    
    def parse_operand(self, operand: str) -> Tuple[str, int]:
        """Parse operand and return (type, value)"""
        operand = operand.strip()
        
        # Register
        if operand.upper() in self.cpu.registers or operand.upper() in ['SP', 'BP', 'IP']:
            return ('register', operand.upper())
        
        # Memory reference [addr] or [register]
        if operand.startswith('[') and operand.endswith(']'):
            inner = operand[1:-1].strip()
            
            # [register+offset] or [register-offset]
            if '+' in inner or '-' in inner:
                return ('memory_offset', inner)
            
            # [register]
            if inner.upper() in self.cpu.registers or inner.upper() in ['SP', 'BP']:
                return ('memory_reg', inner.upper())
            
            # [address]
            try:
                addr = self.parse_number(inner)
                return ('memory', addr)
            except:
                return ('error', f"Invalid memory reference: {operand}")
        
        # Immediate value
        try:
            value = self.parse_number(operand)
            return ('immediate', value)
        except:
            pass
        
        # Label
        if operand in self.labels:
            return ('label', self.labels[operand])
        
        return ('error', f"Invalid operand: {operand}")
    
    def parse_number(self, s: str) -> int:
        """Parse number in various formats"""
        s = s.strip()
        if s.startswith('0x') or s.startswith('0X'):
            return int(s, 16)
        elif s.startswith('0b') or s.startswith('0B'):
            return int(s, 2)
        else:
            return int(s)
    
    def get_operand_value(self, operand: str) -> int:
        """Get the value of an operand"""
        op_type, value = self.parse_operand(operand)
        
        if op_type == 'register':
            if value == 'SP':
                return self.cpu.sp
            elif value == 'BP':
                return self.cpu.bp
            elif value == 'IP':
                return self.cpu.ip
            return self.cpu.registers[value]
        
        elif op_type == 'immediate' or op_type == 'label':
            return value
        
        elif op_type == 'memory':
            return self.read_memory(value, 4)
        
        elif op_type == 'memory_reg':
            addr = self.get_operand_value(value)
            return self.read_memory(addr, 4)
        
        elif op_type == 'memory_offset':
            # Parse BP+4 or SP-8 etc
            if '+' in value:
                reg, offset = value.split('+')
                addr = self.get_operand_value(reg.strip()) + self.parse_number(offset.strip())
            else:
                reg, offset = value.split('-')
                addr = self.get_operand_value(reg.strip()) - self.parse_number(offset.strip())
            return self.read_memory(addr, 4)
        
        raise ValueError(f"Cannot get value of operand: {operand}")
    
    def set_operand_value(self, operand: str, value: int):
        """Set the value of an operand"""
        op_type, op_value = self.parse_operand(operand)
        
        # Ensure value is in 32-bit range
        value = value & 0xFFFFFFFF
        
        if op_type == 'register':
            if op_value == 'SP':
                self.cpu.sp = value
            elif op_value == 'BP':
                self.cpu.bp = value
            elif op_value == 'IP':
                self.cpu.ip = value
            else:
                self.cpu.registers[op_value] = value
        
        elif op_type == 'memory':
            self.write_memory(op_value, value, 4)
        
        elif op_type == 'memory_reg':
            addr = self.get_operand_value(op_value)
            self.write_memory(addr, value, 4)
        
        elif op_type == 'memory_offset':
            if '+' in op_value:
                reg, offset = op_value.split('+')
                addr = self.get_operand_value(reg.strip()) + self.parse_number(offset.strip())
            else:
                reg, offset = op_value.split('-')
                addr = self.get_operand_value(reg.strip()) - self.parse_number(offset.strip())
            self.write_memory(addr, value, 4)
        
        else:
            raise ValueError(f"Cannot set value of operand: {operand}")
    
    def read_memory(self, address: int, size: int) -> int:
        """Read from memory (little-endian)"""
        if address < 0 or address + size > len(self.cpu.memory):
            raise ValueError(f"Memory access violation at 0x{address:08X}")
        
        value = 0
        for i in range(size):
            value |= self.cpu.memory[address + i] << (i * 8)
        return value
    
    def write_memory(self, address: int, value: int, size: int):
        """Write to memory (little-endian)"""
        if address < 0 or address + size > len(self.cpu.memory):
            raise ValueError(f"Memory access violation at 0x{address:08X}")
        
        for i in range(size):
            self.cpu.memory[address + i] = (value >> (i * 8)) & 0xFF
    
    def update_flags(self, result: int, original_bits=32):
        """Update CPU flags based on result"""
        mask = (1 << original_bits) - 1
        result = result & mask
        
        self.cpu.flags['Z'] = (result == 0)
        self.cpu.flags['N'] = (result & (1 << (original_bits - 1))) != 0
        # Carry handled by individual operations
    
    def execute_instruction(self, instr: Instruction) -> Tuple[bool, str]:
        """Execute a single instruction and return detailed explanation"""
        explanation = []
        
        try:
            # Save state before execution
            old_registers = self.cpu.registers.copy()
            old_flags = self.cpu.flags.copy()
            old_sp = self.cpu.sp
            old_ip = self.cpu.ip
            
            explanation.append(f"▶ EXECUTING: {instr.opcode} {', '.join(instr.operands)}")
            if instr.comment:
                explanation.append(f"  Comment: {instr.comment}")
            
            # Execute based on opcode
            if instr.opcode == 'MOV':
                src_val = self.get_operand_value(instr.operands[1])
                explanation.append(f"  Moving value {src_val} (0x{src_val:08X}) to {instr.operands[0]}")
                self.set_operand_value(instr.operands[0], src_val)
            
            elif instr.opcode == 'ADD':
                dst_val = self.get_operand_value(instr.operands[0])
                src_val = self.get_operand_value(instr.operands[1])
                result = dst_val + src_val
                explanation.append(f"  {dst_val} + {src_val} = {result & 0xFFFFFFFF}")
                self.cpu.flags['C'] = result > 0xFFFFFFFF
                self.set_operand_value(instr.operands[0], result)
                self.update_flags(result)
            
            elif instr.opcode == 'SUB':
                dst_val = self.get_operand_value(instr.operands[0])
                src_val = self.get_operand_value(instr.operands[1])
                result = dst_val - src_val
                explanation.append(f"  {dst_val} - {src_val} = {result & 0xFFFFFFFF}")
                self.cpu.flags['C'] = result < 0
                self.set_operand_value(instr.operands[0], result)
                self.update_flags(result)
            
            elif instr.opcode == 'MUL':
                dst_val = self.get_operand_value(instr.operands[0])
                src_val = self.get_operand_value(instr.operands[1])
                result = dst_val * src_val
                explanation.append(f"  {dst_val} × {src_val} = {result & 0xFFFFFFFF}")
                self.cpu.flags['C'] = result > 0xFFFFFFFF
                self.set_operand_value(instr.operands[0], result)
                self.update_flags(result)
            
            elif instr.opcode == 'DIV':
                dst_val = self.get_operand_value(instr.operands[0])
                src_val = self.get_operand_value(instr.operands[1])
                if src_val == 0:
                    return False, "❌ DIVISION BY ZERO"
                result = dst_val // src_val
                explanation.append(f"  {dst_val} ÷ {src_val} = {result}")
                self.set_operand_value(instr.operands[0], result)
                self.update_flags(result)
            
            elif instr.opcode == 'CMP':
                val1 = self.get_operand_value(instr.operands[0])
                val2 = self.get_operand_value(instr.operands[1])
                result = val1 - val2
                explanation.append(f"  Comparing {val1} with {val2}: difference = {result}")
                self.cpu.flags['Z'] = (val1 == val2)
                self.cpu.flags['N'] = (result < 0)
                self.cpu.flags['C'] = (val1 < val2)
                explanation.append(f"  Flags set: Z={int(self.cpu.flags['Z'])}, N={int(self.cpu.flags['N'])}, C={int(self.cpu.flags['C'])}")
            
            elif instr.opcode == 'JMP':
                target = self.labels.get(instr.operands[0])
                if target is None:
                    return False, f"❌ Unknown label: {instr.operands[0]}"
                explanation.append(f"  Unconditional jump to {instr.operands[0]} (address 0x{target:04X})")
                self.cpu.ip = target
                return True, '\n'.join(explanation)
            
            elif instr.opcode == 'JE' or instr.opcode == 'JZ':
                target = self.labels.get(instr.operands[0])
                if target is None:
                    return False, f"❌ Unknown label: {instr.operands[0]}"
                if self.cpu.flags['Z']:
                    explanation.append(f"  Zero flag is SET → Taking jump to {instr.operands[0]}")
                    self.cpu.ip = target
                    return True, '\n'.join(explanation)
                else:
                    explanation.append(f"  Zero flag is CLEAR → Not jumping, continuing to next instruction")
            
            elif instr.opcode == 'JNE' or instr.opcode == 'JNZ':
                target = self.labels.get(instr.operands[0])
                if target is None:
                    return False, f"❌ Unknown label: {instr.operands[0]}"
                if not self.cpu.flags['Z']:
                    explanation.append(f"  Zero flag is CLEAR → Taking jump to {instr.operands[0]}")
                    self.cpu.ip = target
                    return True, '\n'.join(explanation)
                else:
                    explanation.append(f"  Zero flag is SET → Not jumping, continuing to next instruction")
            
            elif instr.opcode == 'JG':
                target = self.labels.get(instr.operands[0])
                if target is None:
                    return False, f"❌ Unknown label: {instr.operands[0]}"
                if not self.cpu.flags['Z'] and not self.cpu.flags['N']:
                    explanation.append(f"  Greater than condition met → Taking jump")
                    self.cpu.ip = target
                    return True, '\n'.join(explanation)
                else:
                    explanation.append(f"  Greater than condition not met → Not jumping")
            
            elif instr.opcode == 'JL':
                target = self.labels.get(instr.operands[0])
                if target is None:
                    return False, f"❌ Unknown label: {instr.operands[0]}"
                if self.cpu.flags['N']:
                    explanation.append(f"  Less than condition met → Taking jump")
                    self.cpu.ip = target
                    return True, '\n'.join(explanation)
                else:
                    explanation.append(f"  Less than condition not met → Not jumping")
            
            elif instr.opcode == 'PUSH':
                value = self.get_operand_value(instr.operands[0])
                self.cpu.sp -= 4
                explanation.append(f"  Pushing value {value} (0x{value:08X}) onto stack")
                explanation.append(f"  Stack pointer: 0x{old_sp:04X} → 0x{self.cpu.sp:04X}")
                self.write_memory(self.cpu.sp, value, 4)
            
            elif instr.opcode == 'POP':
                if self.cpu.sp >= len(self.cpu.memory) - 4:
                    return False, "❌ STACK UNDERFLOW"
                value = self.read_memory(self.cpu.sp, 4)
                explanation.append(f"  Popping value {value} (0x{value:08X}) from stack")
                self.set_operand_value(instr.operands[0], value)
                old_sp = self.cpu.sp
                self.cpu.sp += 4
                explanation.append(f"  Stack pointer: 0x{old_sp:04X} → 0x{self.cpu.sp:04X}")
            
            elif instr.opcode == 'CALL':
                target = self.labels.get(instr.operands[0])
                if target is None:
                    return False, f"❌ Unknown label: {instr.operands[0]}"
                # Push return address
                return_addr = self.cpu.ip + 4
                self.cpu.sp -= 4
                self.write_memory(self.cpu.sp, return_addr, 4)
                explanation.append(f"  Calling function {instr.operands[0]}")
                explanation.append(f"  Return address 0x{return_addr:04X} pushed to stack")
                explanation.append(f"  Jumping to 0x{target:04X}")
                self.cpu.ip = target
                return True, '\n'.join(explanation)
            
            elif instr.opcode == 'RET':
                if self.cpu.sp >= len(self.cpu.memory) - 4:
                    return False, "❌ STACK UNDERFLOW on RET"
                return_addr = self.read_memory(self.cpu.sp, 4)
                self.cpu.sp += 4
                explanation.append(f"  Returning to address 0x{return_addr:04X}")
                self.cpu.ip = return_addr
                return True, '\n'.join(explanation)
            
            elif instr.opcode == 'LOAD':
                addr_val = self.get_operand_value(instr.operands[1])
                value = self.read_memory(addr_val, 4)
                explanation.append(f"  Loading value {value} (0x{value:08X}) from memory[0x{addr_val:04X}]")
                self.set_operand_value(instr.operands[0], value)
            
            elif instr.opcode == 'STORE':
                addr_val = self.get_operand_value(instr.operands[0])
                value = self.get_operand_value(instr.operands[1])
                explanation.append(f"  Storing value {value} (0x{value:08X}) to memory[0x{addr_val:04X}]")
                self.write_memory(addr_val, value, 4)
            
            elif instr.opcode == 'HLT':
                explanation.append("  Halting CPU")
                self.cpu.halted = True
                return True, '\n'.join(explanation)
            
            else:
                return False, f"❌ Unknown instruction: {instr.opcode}"
            
            # Advance IP if not already modified by jump/call/ret
            if self.cpu.ip == old_ip:
                self.cpu.ip += 4
            
            # Show what changed
            changes = []
            for reg, val in self.cpu.registers.items():
                if val != old_registers[reg]:
                    changes.append(f"  {reg}: 0x{old_registers[reg]:08X} → 0x{val:08X}")
            
            if self.cpu.sp != old_sp:
                changes.append(f"  SP: 0x{old_sp:04X} → 0x{self.cpu.sp:04X}")
            
            if self.cpu.ip != old_ip and instr.opcode not in ['JMP', 'JE', 'JNE', 'JG', 'JL', 'CALL', 'RET']:
                changes.append(f"  IP: 0x{old_ip:04X} → 0x{self.cpu.ip:04X}")
            
            flag_changes = []
            for flag, val in self.cpu.flags.items():
                if val != old_flags[flag]:
                    flag_changes.append(f"{flag}={int(val)}")
            
            if flag_changes:
                changes.append(f"  FLAGS: {' '.join(flag_changes)}")
            
            if changes:
                explanation.append("\n🔄 CHANGES:")
                explanation.extend(changes)
            else:
                explanation.append("\n  (No register or flag changes)")
            
            return True, '\n'.join(explanation)
            
        except Exception as e:
            return False, f"❌ ERROR: {str(e)}"
    
    def find_instruction_at_ip(self) -> Optional[Instruction]:
        """Find instruction at current IP"""
        for instr in self.instructions:
            if instr.address == self.cpu.ip:
                return instr
        return None
    
    def check_breakpoints(self) -> Optional[Breakpoint]:
        """Check if any breakpoint is hit"""
        for bp in self.breakpoints:
            if not bp.enabled:
                continue
            
            if bp.address is not None and bp.address == self.cpu.ip:
                return bp
            
            if bp.condition:
                # Evaluate condition (simplified)
                try:
                    # Replace register names with values
                    cond = bp.condition
                    for reg in self.cpu.registers:
                        cond = cond.replace(reg, str(self.cpu.registers[reg]))
                    if eval(cond):
                        return bp
                except:
                    pass
        
        return None
    
    def get_state_display(self) -> str:
        """Get formatted CPU state"""
        lines = []
        lines.append("REGISTERS:")
        
        # Show registers in rows of 4
        for i in range(0, 8, 4):
            reg_line = "  "
            for j in range(4):
                if i + j < 8:
                    reg = f"R{i+j}"
                    val = self.cpu.registers[reg]
                    reg_line += f"{reg}=0x{val:08X}  "
            lines.append(reg_line)
        
        lines.append(f"  SP=0x{self.cpu.sp:04X}  BP=0x{self.cpu.bp:04X}  IP=0x{self.cpu.ip:04X}")
        
        flags_str = ' '.join([f"{k}={int(v)}" for k, v in self.cpu.flags.items()])
        lines.append(f"  FLAGS: {flags_str}")
        
        # Show stack
        lines.append("\nSTACK (top 16 bytes):")
        for addr in range(self.cpu.sp, min(self.cpu.sp + 16, len(self.cpu.memory)), 4):
            value = self.read_memory(addr, 4)
            ascii_repr = ''.join([chr(b) if 32 <= b < 127 else '.' 
                                 for b in self.cpu.memory[addr:addr+4]])
            lines.append(f"  0x{addr:04X}: {value:08X}  {ascii_repr}")
        
        return '\n'.join(lines)
    
    def compile_c_to_asm(self, c_code: str) -> str:
        """Simple C to Assembly compiler (educational)"""
        output = []
        output.append("C → ASSEMBLY COMPILATION")
        output.append("=" * 50)
        output.append("")
        
        # Very simplified C parser for educational purposes
        # Handle simple function definitions
        
        # Example: int add(int a, int b) { return a + b; }
        func_pattern = r'(\w+)\s+(\w+)\s*\(([^)]*)\)\s*\{([^}]*)\}'
        matches = re.findall(func_pattern, c_code, re.DOTALL)
        
        if not matches:
            output.append("❌ No functions found or unsupported C syntax")
            return '\n'.join(output)
        
        for return_type, func_name, params, body in matches:
            output.append(f"FUNCTION: {func_name}")
            output.append("")
            
            # Parse parameters
            param_list = [p.strip() for p in params.split(',') if p.strip()]
            param_names = []
            for p in param_list:
                parts = p.split()
                if len(parts) >= 2:
                    param_names.append(parts[-1])
            
            # Stack frame explanation
            output.append("STACK FRAME LAYOUT:")
            output.append("  [BP+0]  ← saved BP")
            output.append("  [BP-4]  ← return address")
            
            offset = -8
            var_offsets = {}
            for param in param_names:
                output.append(f"  [BP{offset}]  ← {param}")
                var_offsets[param] = offset
                offset -= 4
            
            output.append("")
            output.append("GENERATED ASSEMBLY:")
            output.append("")
            
            # Function prologue
            asm = []
            asm.append(f"{func_name}:")
            asm.append("    PUSH BP")
            asm.append("    MOV BP, SP")
            
            # Parse body for return statement
            return_match = re.search(r'return\s+([^;]+);', body)
            if return_match:
                expr = return_match.group(1).strip()
                
                # Simple expression parsing
                if '+' in expr:
                    parts = expr.split('+')
                    var1 = parts[0].strip()
                    var2 = parts[1].strip()
                    
                    asm.append(f"    ; Calculate {var1} + {var2}")
                    asm.append(f"    MOV R0, [BP{var_offsets.get(var1, -8)}]")
                    asm.append(f"    ADD R0, [BP{var_offsets.get(var2, -12)}]")
                    asm.append("    ; Result in R0")
                
                elif '-' in expr:
                    parts = expr.split('-')
                    var1 = parts[0].strip()
                    var2 = parts[1].strip()
                    
                    asm.append(f"    ; Calculate {var1} - {var2}")
                    asm.append(f"    MOV R0, [BP{var_offsets.get(var1, -8)}]")
                    asm.append(f"    SUB R0, [BP{var_offsets.get(var2, -12)}]")
                
                elif expr in var_offsets:
                    asm.append(f"    ; Return {expr}")
                    asm.append(f"    MOV R0, [BP{var_offsets[expr]}]")
                
                else:
                    # Try to parse as number
                    try:
                        val = int(expr)
                        asm.append(f"    ; Return constant {val}")
                        asm.append(f"    MOV R0, {val}")
                    except:
                        asm.append(f"    ; Return expression: {expr}")
                        asm.append(f"    MOV R0, 0  ; SIMPLIFIED")
            
            # Function epilogue
            asm.append("    POP BP")
            asm.append("    RET")
            
            output.extend(asm)
            output.append("")
        
        output.append("=" * 50)
        output.append("✓ Compilation complete. Copy assembly code to execute.")
        
        return '\n'.join(output)


class DebuggerGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("🧠 Assembly Debugger & CPU Simulator")
        self.root.geometry("1400x900")
        
        self.debugger = AssemblyDebugger()
        self.running = False
        self.stepping = False
        
        self.setup_ui()
        self.show_startup_message()
    
    def setup_ui(self):
        # Main container with paned window
        main_paned = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        main_paned.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # Left panel - Code editor
        left_frame = ttk.Frame(main_paned)
        main_paned.add(left_frame, weight=1)
        
        # Code input
        ttk.Label(left_frame, text="Assembly Code:", font=('Courier', 10, 'bold')).pack(anchor='w')
        
        self.code_text = scrolledtext.ScrolledText(left_frame, width=50, height=25, 
                                                    font=('Courier', 10))
        self.code_text.pack(fill=tk.BOTH, expand=True, pady=5)
        
        # Sample code
        sample = """; Simple addition example