# -*- coding: utf-8 -*-
"""
 Copyright 2025 NXP

 NXP Confidential and Proprietary. This software is owned or controlled by NXP and may only be
 used strictly in accordance with the applicable license terms. By expressly
 accepting such terms or by downloading, installing, activating and/or otherwise
 using the software, you are agreeing that you have read, and that you agree to
 comply with and are bound by, such license terms. If you do not agree to be
 bound by the applicable license terms, then you may not retain, install,
 activate or otherwise use the software.

"""

class ZENVRegisters:
    """
    RISC-V processor register definitions
    """
    
    # Base addresses from TCL file
    ABSTR_REGISTERS_BASE = 0x0
    DEBUG_REGISTERS_BASE = 0x10000
    CSR_REGISTERS_BASE = 0x20000
    
    # Abstract Registers (accessed via Abstract commands)
    # Formula: abstr_registers_base + (reg & 0xFFFF)
    ABSTRACT_REGISTERS = {
        "tselect": ABSTR_REGISTERS_BASE + (0x07a0 & 0xFFFF),
        "tdata1": ABSTR_REGISTERS_BASE + (0x07a1 & 0xFFFF),
        "mcontrol": ABSTR_REGISTERS_BASE + (0x07a1 & 0xFFFF),
        "tdata2": ABSTR_REGISTERS_BASE + (0x07a2 & 0xFFFF),
        "tinfo": ABSTR_REGISTERS_BASE + (0x07a4 & 0xFFFF),
        "dcsr": ABSTR_REGISTERS_BASE + (0x07b0 & 0xFFFF),
        "dpc": ABSTR_REGISTERS_BASE + (0x07b1 & 0xFFFF),
        "gpr0": ABSTR_REGISTERS_BASE + (0x1000 & 0xFFFF),
        "gpr1": ABSTR_REGISTERS_BASE + (0x1001 & 0xFFFF),
        "gpr2": ABSTR_REGISTERS_BASE + (0x1002 & 0xFFFF),
        "gpr3": ABSTR_REGISTERS_BASE + (0x1003 & 0xFFFF),
        "gpr4": ABSTR_REGISTERS_BASE + (0x1004 & 0xFFFF),
        "gpr5": ABSTR_REGISTERS_BASE + (0x1005 & 0xFFFF),
        "gpr6": ABSTR_REGISTERS_BASE + (0x1006 & 0xFFFF),
        "gpr7": ABSTR_REGISTERS_BASE + (0x1007 & 0xFFFF),
        "gpr8": ABSTR_REGISTERS_BASE + (0x1008 & 0xFFFF),
        "gpr9": ABSTR_REGISTERS_BASE + (0x1009 & 0xFFFF),
        "gpr10": ABSTR_REGISTERS_BASE + (0x100A & 0xFFFF),
        "gpr11": ABSTR_REGISTERS_BASE + (0x100B & 0xFFFF),
        "gpr12": ABSTR_REGISTERS_BASE + (0x100C & 0xFFFF),
        "gpr13": ABSTR_REGISTERS_BASE + (0x100D & 0xFFFF),
        "gpr14": ABSTR_REGISTERS_BASE + (0x100E & 0xFFFF),
        "gpr15": ABSTR_REGISTERS_BASE + (0x100F & 0xFFFF),
        # Just for ZV0035, not applicable for ZV0020
        "gpr16": ABSTR_REGISTERS_BASE + (0x1010 & 0xFFFF),
        "gpr17": ABSTR_REGISTERS_BASE + (0x1011 & 0xFFFF),
        "gpr18": ABSTR_REGISTERS_BASE + (0x1012 & 0xFFFF),
        "gpr19": ABSTR_REGISTERS_BASE + (0x1013 & 0xFFFF),
        "gpr20": ABSTR_REGISTERS_BASE + (0x1014 & 0xFFFF),
        "gpr21": ABSTR_REGISTERS_BASE + (0x1015 & 0xFFFF),
        "gpr22": ABSTR_REGISTERS_BASE + (0x1016 & 0xFFFF),
        "gpr23": ABSTR_REGISTERS_BASE + (0x1017 & 0xFFFF),
        "gpr24": ABSTR_REGISTERS_BASE + (0x1018 & 0xFFFF),
        "gpr25": ABSTR_REGISTERS_BASE + (0x1019 & 0xFFFF),
        "gpr26": ABSTR_REGISTERS_BASE + (0x101A & 0xFFFF),
        "gpr27": ABSTR_REGISTERS_BASE + (0x101B & 0xFFFF),
        "gpr28": ABSTR_REGISTERS_BASE + (0x101C & 0xFFFF),
        "gpr29": ABSTR_REGISTERS_BASE + (0x101D & 0xFFFF),
        "gpr30": ABSTR_REGISTERS_BASE + (0x101E & 0xFFFF),
        "gpr31": ABSTR_REGISTERS_BASE + (0x101F & 0xFFFF),
    }
    
    # Debug Registers
    # Formula: debug_registers_base + (reg & 0x3FF)
    DEBUG_REGISTERS = {
        "data0": DEBUG_REGISTERS_BASE + (0x004 & 0x3FF),
        "dmcontrol": DEBUG_REGISTERS_BASE + (0x010 & 0x3FF),
        "dmstatus": DEBUG_REGISTERS_BASE + (0x011 & 0x3FF),
        "abstractcs": DEBUG_REGISTERS_BASE + (0x016 & 0x3FF),
        "command": DEBUG_REGISTERS_BASE + (0x017 & 0x3FF),
        "abstractauto": DEBUG_REGISTERS_BASE + (0x018 & 0x3FF),
        "progbuf0": DEBUG_REGISTERS_BASE + (0x020 & 0x3FF),
        "progbuf1": DEBUG_REGISTERS_BASE + (0x021 & 0x3FF),
        "progbuf2": DEBUG_REGISTERS_BASE + (0x022 & 0x3FF),
        "dmcs2": DEBUG_REGISTERS_BASE + (0x032 & 0x3FF),
        "haltsum0": DEBUG_REGISTERS_BASE + (0x040 & 0x3FF),
        "devarch": DEBUG_REGISTERS_BASE + (0x3ef & 0x3FF),
        "devtype": DEBUG_REGISTERS_BASE + (0x3f0 & 0x3FF),
        "devid2": DEBUG_REGISTERS_BASE + (0x3f1 & 0x3FF),
        "devid1": DEBUG_REGISTERS_BASE + (0x3f2 & 0x3FF),
        "devid": DEBUG_REGISTERS_BASE + (0x3f3 & 0x3FF),
        "pidr4": DEBUG_REGISTERS_BASE + (0x3f4 & 0x3FF),
        "pidr5": DEBUG_REGISTERS_BASE + (0x3f5 & 0x3FF),
        "pidr6": DEBUG_REGISTERS_BASE + (0x3f6 & 0x3FF),
        "pidr7": DEBUG_REGISTERS_BASE + (0x3f7 & 0x3FF),
        "pidr0": DEBUG_REGISTERS_BASE + (0x3f8 & 0x3FF),
        "pidr1": DEBUG_REGISTERS_BASE + (0x3f9 & 0x3FF),
        "pidr2": DEBUG_REGISTERS_BASE + (0x3fa & 0x3FF),
        "pidr3": DEBUG_REGISTERS_BASE + (0x3fb & 0x3FF),
        "cidr0": DEBUG_REGISTERS_BASE + (0x3fc & 0x3FF),
        "cidr1": DEBUG_REGISTERS_BASE + (0x3fd & 0x3FF),
        "cidr2": DEBUG_REGISTERS_BASE + (0x3fe & 0x3FF),
        "cidr3": DEBUG_REGISTERS_BASE + (0x3ff & 0x3FF),
    }
    
    # CSR Registers
    # Formula: csr_registers_base + (reg & 0xFFF)
    CSR_REGISTERS = {
        "misa": CSR_REGISTERS_BASE + (0x301 & 0xFFF),
        "mvendorid": CSR_REGISTERS_BASE + (0xF11 & 0xFFF),
        "marchid": CSR_REGISTERS_BASE + (0xF12 & 0xFFF),
        "mimpid": CSR_REGISTERS_BASE + (0xF13 & 0xFFF),
        "mhartid": CSR_REGISTERS_BASE + (0xF14 & 0xFFF),
        "mstatus": CSR_REGISTERS_BASE + (0x300 & 0xFFF),
        "sstatus": CSR_REGISTERS_BASE + (0x100 & 0xFFF),
        "mtvec": CSR_REGISTERS_BASE + (0x305 & 0xFFF),
        "stvec": CSR_REGISTERS_BASE + (0x105 & 0xFFF),
        "medeleg": CSR_REGISTERS_BASE + (0x302 & 0xFFF),
        "mideleg": CSR_REGISTERS_BASE + (0x303 & 0xFFF),
        "mip": CSR_REGISTERS_BASE + (0x344 & 0xFFF),
        "sip": CSR_REGISTERS_BASE + (0x144 & 0xFFF),
        "mie": CSR_REGISTERS_BASE + (0x304 & 0xFFF),
        "sie": CSR_REGISTERS_BASE + (0x104 & 0xFFF),
        "mscratch": CSR_REGISTERS_BASE + (0x340 & 0xFFF),
        "sscratch": CSR_REGISTERS_BASE + (0x140 & 0xFFF),
        "mepc": CSR_REGISTERS_BASE + (0x341 & 0xFFF),
        "sepc": CSR_REGISTERS_BASE + (0x141 & 0xFFF),
        "mcause": CSR_REGISTERS_BASE + (0x342 & 0xFFF),
        "scause": CSR_REGISTERS_BASE + (0x142 & 0xFFF),
        "mtval": CSR_REGISTERS_BASE + (0x343 & 0xFFF),
        "stval": CSR_REGISTERS_BASE + (0x143 & 0xFFF),
        "mcycle": CSR_REGISTERS_BASE + (0xB00 & 0xFFF),
        "cycle": CSR_REGISTERS_BASE + (0xC00 & 0xFFF),
        "mcycleh": CSR_REGISTERS_BASE + (0xB80 & 0xFFF),
        "cycleh": CSR_REGISTERS_BASE + (0xC80 & 0xFFF),
        "minstret": CSR_REGISTERS_BASE + (0xB02 & 0xFFF),
        "instret": CSR_REGISTERS_BASE + (0xC02 & 0xFFF),
        "minstreth": CSR_REGISTERS_BASE + (0xB82 & 0xFFF),
        "instreth": CSR_REGISTERS_BASE + (0xC82 & 0xFFF),
        "mhpmcounter3": CSR_REGISTERS_BASE + (0xB03 & 0xFFF),
        "mhpmcounter31": CSR_REGISTERS_BASE + (0xB1F & 0xFFF),
        "hpmcounter3": CSR_REGISTERS_BASE + (0xC03 & 0xFFF),
        "hpmcounter31": CSR_REGISTERS_BASE + (0xC1F & 0xFFF),
        "mhpmcounter3h": CSR_REGISTERS_BASE + (0xB83 & 0xFFF),
        "mhpmcounter31h": CSR_REGISTERS_BASE + (0xB9F & 0xFFF),
        "hpmcounter3h": CSR_REGISTERS_BASE + (0xC83 & 0xFFF),
        "hpmcounter31h": CSR_REGISTERS_BASE + (0xC9F & 0xFFF),
        "mhpmevent3": CSR_REGISTERS_BASE + (0x323 & 0xFFF),
        "mhpmevent31": CSR_REGISTERS_BASE + (0x33F & 0xFFF),
        "mcounteren": CSR_REGISTERS_BASE + (0x306 & 0xFFF),
        "scounteren": CSR_REGISTERS_BASE + (0x106 & 0xFFF),
        "mcountinhibit": CSR_REGISTERS_BASE + (0x320 & 0xFFF),
        "bucsr": CSR_REGISTERS_BASE + (0x7C0 & 0xFFF),
        "tmcfg0": CSR_REGISTERS_BASE + (0xFC0 & 0xFFF),
        "l1cfg0": CSR_REGISTERS_BASE + (0xCC1 & 0xFFF),
        "l1cfg1": CSR_REGISTERS_BASE + (0xCC2 & 0xFFF),
        "dmemcfg0": CSR_REGISTERS_BASE + (0xFD0 & 0xFFF),
        "imemcfg0": CSR_REGISTERS_BASE + (0xFD4 & 0xFFF),
        "ucsstate": CSR_REGISTERS_BASE + (0x820 & 0xFFF),
        "ucscorr": CSR_REGISTERS_BASE + (0x823 & 0xFFF),
        "upcsstate": CSR_REGISTERS_BASE + (0x821 & 0xFFF),
        "uacsstate": CSR_REGISTERS_BASE + (0x822 & 0xFFF),
        # Thread management
        "ten": CSR_REGISTERS_BASE + (0xBC0 & 0xFFF),
        "tensr": CSR_REGISTERS_BASE + (0xFC1 & 0xFFF),
        "inia0": CSR_REGISTERS_BASE + (0xBE0 & 0xFFF),
        "inia1": CSR_REGISTERS_BASE + (0xBE1 & 0xFFF),
        # Just for ZV0035, not applicable for ZV0020
        "satp": CSR_REGISTERS_BASE + (0x180 & 0xFFF),
        "pmpcfg0": CSR_REGISTERS_BASE + (0x3A0 & 0xFFF),
        "pmpcfg1": CSR_REGISTERS_BASE + (0x3A1 & 0xFFF),
        "pmpcfg2": CSR_REGISTERS_BASE + (0x3A2 & 0xFFF),
        "pmpcfg3": CSR_REGISTERS_BASE + (0x3A3 & 0xFFF),
        "pmpaddr0": CSR_REGISTERS_BASE + (0x3B0 & 0xFFF),
        "pmpaddr15": CSR_REGISTERS_BASE + (0x3BF & 0xFFF),
        # Just for ZV2337
        "mncause": CSR_REGISTERS_BASE + (0x742 & 0xFFF),
        "mnepc": CSR_REGISTERS_BASE + (0x741 & 0xFFF),
        "mnscratch": CSR_REGISTERS_BASE + (0x740 & 0xFFF),
        "mnstatus": CSR_REGISTERS_BASE + (0x744 & 0xFFF),
    }
    
    # Generate PMP address registers (pmpaddr1-pmpaddr14)
    for n in range(1, 15):
        CSR_REGISTERS["pmpaddr{}".format(n)] = CSR_REGISTERS_BASE + ((0x3B0 + n) & 0xFFF)
    
    # Generate performance counter registers (mhpmcounter4-mhpmcounter30)
    for n in range(4, 31):
        CSR_REGISTERS["mhpmcounter{}".format(n)] = CSR_REGISTERS_BASE + ((0xB00 + n) & 0xFFF)
        CSR_REGISTERS["hpmcounter{}".format(n)] = CSR_REGISTERS_BASE + ((0xC00 + n) & 0xFFF)
        CSR_REGISTERS["mhpmcounter{}h".format(n)] = CSR_REGISTERS_BASE + ((0xB80 + n) & 0xFFF)
        CSR_REGISTERS["hpmcounter{}h".format(n)] = CSR_REGISTERS_BASE + ((0xC80 + n) & 0xFFF)
        CSR_REGISTERS["mhpmevent{}".format(n)] = CSR_REGISTERS_BASE + ((0x320 + n) & 0xFFF)
    
    # Register Aliases
    ALIASES = {}
    
    def __init__(self):
        """Initialize the ZENV register manager"""
        self.all_registers = {}
        self.all_registers.update(self.ABSTRACT_REGISTERS)
        self.all_registers.update(self.DEBUG_REGISTERS)
        self.all_registers.update(self.CSR_REGISTERS)
    
    def get_address_by_name(self, name):
        """Get register address by name"""
        name_lower = name.lower()
        
        # Check aliases first
        if name_lower in self.ALIASES:
            name_lower = self.ALIASES[name_lower]
        
        if name_lower in self.all_registers:
            return self.all_registers[name_lower]
        return None
    
    def get_abstract_reg(self, name):
        """Get abstract register address by name"""
        name_lower = name.lower()
        if name_lower in self.ABSTRACT_REGISTERS:
            return self.ABSTRACT_REGISTERS[name_lower]
        return None
    
    def get_debug_reg(self, name):
        """Get debug register address by name"""
        name_lower = name.lower()
        if name_lower in self.DEBUG_REGISTERS:
            return self.DEBUG_REGISTERS[name_lower]
        return None
    
    def get_csr_reg(self, name):
        """Get CSR register address by name"""
        name_lower = name.lower()
        if name_lower in self.CSR_REGISTERS:
            return self.CSR_REGISTERS[name_lower]
        return None
    
    def is_abstract_reg(self, name):
        """Check if register is an abstract register"""
        return name.lower() in self.ABSTRACT_REGISTERS
    
    def is_debug_reg(self, name):
        """Check if register is a debug register"""
        return name.lower() in self.DEBUG_REGISTERS
    
    def is_csr_reg(self, name):
        """Check if register is a CSR register"""
        return name.lower() in self.CSR_REGISTERS
    
    def is_gpr_reg(self, name):
        """Check if register is a general purpose register (GPR)"""
        name_lower = name.lower()
        return name_lower.startswith('gpr') and name_lower in self.ABSTRACT_REGISTERS
    
    def is_machine_mode_reg(self, name):
        """Check if register is a machine mode CSR register"""
        name_lower = name.lower()
        return (name_lower in self.CSR_REGISTERS and 
                (name_lower.startswith('m') or name_lower.startswith('cycle')))
    
    def is_supervisor_mode_reg(self, name):
        """Check if register is a supervisor mode CSR register"""
        name_lower = name.lower()
        return (name_lower in self.CSR_REGISTERS and 
                name_lower.startswith('s'))
    
    def is_user_mode_reg(self, name):
        """Check if register is a user mode CSR register"""
        name_lower = name.lower()
        return (name_lower in self.CSR_REGISTERS and 
                name_lower.startswith('u'))