import sys

class PhysicalMemory:
    def __init__(self, size_mb=2):
        self.memory = [0] * ((size_mb * 1024 * 1024) // 4)
    def write_word(self, paddr, value): self.memory[paddr >> 2] = value & 0xFFFFFFFF
    def read_word(self, paddr): return self.memory[paddr >> 2]

class TLBEntry:
    def __init__(self):
        self.vpn = None
        self.ppn = None
        self.valid = False
        self.last_access = 0

class PageTableSimulator:
    def __init__(self, phys_mem, tlb_size=4):
        self.phys_mem = phys_mem
        self.tlb = [TLBEntry() for _ in range(tlb_size)]
        self.timer = 0
        self.stats = {"tlb_hits": 0, "tlb_misses": 0, "page_faults": 0}

    def translate(self, dir_ptr, vaddr):
        self.timer += 1
        vpn = (vaddr >> 12)
        offset = vaddr & 0xFFF
        status = ""
        paddr = None

        # 1. TLB Lookup
        for entry in self.tlb:
            if entry.valid and entry.vpn == vpn:
                self.stats["tlb_hits"] += 1
                entry.last_access = self.timer
                paddr = (entry.ppn << 12) | offset
                status = "TLB Hit"
                print(f"0x{vaddr:08X}: {status} (PA: 0x{paddr:08X})")
                return paddr

        self.stats["tlb_misses"] += 1
        
        # 2. Page Walk
        l1_idx = (vaddr >> 22) & 0x3FF
        l2_idx = (vaddr >> 12) & 0x3FF

        pde = self.phys_mem.read_word(dir_ptr + (l1_idx * 4))
        if not (pde & 0x1):
            self.stats["page_faults"] += 1
            print(f"0x{vaddr:08X}: Page Fault")
            return "PAGE_FAULT"

        l2_base = (pde >> 12) << 12
        pte = self.phys_mem.read_word(l2_base + (l2_idx * 4))
        if not (pte & 0x1):
            self.stats["page_faults"] += 1
            print(f"0x{vaddr:08X}: Page Fault")
            return "PAGE_FAULT"

        ppn = pte >> 12
        paddr = (ppn << 12) | offset
        status = "TLB Miss"

        # 3. TLB Update (LRU)
        lru_entry = min(self.tlb, key=lambda x: x.last_access if x.valid else -1)
        lru_entry.vpn = vpn
        lru_entry.ppn = ppn
        lru_entry.valid = True
        lru_entry.last_access = self.timer

        print(f"0x{vaddr:08X}: {status} (PA: 0x{paddr:08X})")
        return paddr

def setup_memory(mem):
    # L1 Directory at 0x0000
    mem.write_word(0x0000, 0x00004001) # L1[0] -> L2 @ 0x4000
    mem.write_word(0x0004, 0x00005001) # L1[1] -> L2 @ 0x5000
    mem.write_word(0x0008, 0x00006001) # L1[2] -> L2 @ 0x6000
    
    # L2 Table A (0x4000)
    mem.write_word(0x4000 + (0x0*4), 0x0000A001) # VPN 0x00000 -> PPN 0x0A000
    mem.write_word(0x4000 + (0x1*4), 0x0000B001) # VPN 0x00001 -> PPN 0x0B000
    
    # L2 Table B (0x5000)
    mem.write_word(0x5000 + (0x2*4), 0x0000C001) # VPN 0x00402 -> PPN 0x0C000
    mem.write_word(0x5000 + (0x3*4), 0x0000D001) # VPN 0x00403 -> PPN 0x0D000
    
    # L2 Table C (0x6000)
    mem.write_word(0x6000 + (0x4*4), 0x0000E001) # VPN 0x00804 -> PPN 0x0E000
    mem.write_word(0x6000 + (0x5*4), 0x0000F001) # VPN 0x00805 -> PPN 0x0F000

def main():
    if len(sys.argv) < 2:
        print("Usage: python vm_sim.py <trace_file>")
        return

    mem = PhysicalMemory()
    setup_memory(mem)
    sim = PageTableSimulator(mem)
    
    with open(sys.argv[1], 'r') as f:
        for line in f:
            if line.strip():
                vaddr = int(line.strip(), 16)
                sim.translate(0x0000, vaddr)

    print("\n--- Final Statistics ---")
    print(f"TLB Hits: {sim.stats['tlb_hits']}")
    print(f"TLB Misses: {sim.stats['tlb_misses']}")
    print(f"Page Faults: {sim.stats['page_faults']}")

if __name__ == "__main__":
    main()