from __future__ import annotations

import logging
from enum import IntFlag

import ida_name
from ida_idaapi import ea_t
from typing_extensions import TYPE_CHECKING, Iterator, Optional, Tuple, Union

from .base import DatabaseEntity, InvalidEAError, check_db_open, decorate_all_methods

if TYPE_CHECKING:
    from .database import Database

logger = logging.getLogger(__name__)


class SetNameFlags(IntFlag):
    """Flags for set_name() function."""

    CHECK = ida_name.SN_CHECK  # Fail if the name contains invalid characters
    NOCHECK = ida_name.SN_NOCHECK  # Replace invalid characters silently
    PUBLIC = ida_name.SN_PUBLIC  # Make name public
    NON_PUBLIC = ida_name.SN_NON_PUBLIC  # Make name non-public
    WEAK = ida_name.SN_WEAK  # Make name weak
    NON_WEAK = ida_name.SN_NON_WEAK  # Make name non-weak
    AUTO = ida_name.SN_AUTO  # Make name autogenerated
    NON_AUTO = ida_name.SN_NON_AUTO  # Make name non-autogenerated
    NOLIST = ida_name.SN_NOLIST  # Exclude name from the list
    NOWARN = ida_name.SN_NOWARN  # Don't display a warning if failed
    LOCAL = ida_name.SN_LOCAL  # Create local name
    IDBENC = ida_name.SN_IDBENC  # Name is given in IDB encoding
    FORCE = ida_name.SN_FORCE  # Try variations if name exists
    NODUMMY = ida_name.SN_NODUMMY  # Don't prepend with '_' for dummy suffixes
    DELTAIL = ida_name.SN_DELTAIL  # Delete hindering item if tail byte


class DemangleFlags(IntFlag):
    """Flags for demangling operations."""

    # Memory model mask (0x7)
    PTRMSK = ida_name.MNG_PTRMSK  # Memory model mask
    DEFNEAR = ida_name.MNG_DEFNEAR  # Inhibit near, display everything else
    DEFNEARANY = ida_name.MNG_DEFNEARANY  # Inhibit near/__ptr64, display everything else
    DEFFAR = ida_name.MNG_DEFFAR  # Inhibit far, display everything else
    NOPTRTYP16 = ida_name.MNG_NOPTRTYP16  # Inhibit everything (disables vc7-extensions)
    DEFHUGE = ida_name.MNG_DEFHUGE  # Inhibit huge, display everything else
    DEFPTR64 = ida_name.MNG_DEFPTR64  # Inhibit __ptr64, display everything else
    DEFNONE = ida_name.MNG_DEFNONE  # Display everything
    NOPTRTYP = ida_name.MNG_NOPTRTYP  # Inhibit everything

    # Name component inhibition flags
    NODEFINIT = ida_name.MNG_NODEFINIT  # Inhibit everything except the main name
    NOUNDERSCORE = ida_name.MNG_NOUNDERSCORE  # Inhibit underscores in __ccall, __pascal...
    NOTYPE = ida_name.MNG_NOTYPE  # Inhibit callc&based
    NORETTYPE = ida_name.MNG_NORETTYPE  # Inhibit return type of functions
    NOBASEDT = ida_name.MNG_NOBASEDT  # Inhibit base types
    NOCALLC = ida_name.MNG_NOCALLC  # Inhibit __pascal/__ccall/etc
    NOPOSTFC = ida_name.MNG_NOPOSTFC  # Inhibit postfix const
    NOSCTYP = ida_name.MNG_NOSCTYP  # Inhibit public/private/protected
    NOTHROW = ida_name.MNG_NOTHROW  # Inhibit throw description
    NOSTVIR = ida_name.MNG_NOSTVIR  # Inhibit "static" & "virtual"
    NOECSU = ida_name.MNG_NOECSU  # Inhibit class/struct/union/enum
    NOCSVOL = ida_name.MNG_NOCSVOL  # Inhibit const/volatile/restrict
    NOCLOSUR = ida_name.MNG_NOCLOSUR  # Inhibit __closure for borland
    NOUNALG = ida_name.MNG_NOUNALG  # Inhibit __unaligned
    NOMANAGE = ida_name.MNG_NOMANAGE  # Inhibit __pin/__box/__gc for ms(.net)
    NOMODULE = ida_name.MNG_NOMODULE  # Inhibit module names (Swift)

    # Format flags
    SHORT_S = ida_name.MNG_SHORT_S  # signed (int) is displayed as s(int)
    SHORT_U = ida_name.MNG_SHORT_U  # unsigned (int) is displayed as u(int)
    ZPT_SPACE = ida_name.MNG_ZPT_SPACE  # Display space after comma in the arglist
    DROP_IMP = ida_name.MNG_DROP_IMP  # Inhibit __declspec(dllimport)

    # Name processing flags
    IGN_ANYWAY = ida_name.MNG_IGN_ANYWAY  # Ignore '_nn' at the end of name
    IGN_JMP = ida_name.MNG_IGN_JMP  # Ignore 'j_' at the beginning of name
    MOVE_JMP = ida_name.MNG_MOVE_JMP  # Move 'j_' prefix to the demangled name

    # Compiler and composite flags
    COMPILER_MSK = ida_name.MNG_COMPILER_MSK  # Compiler mask (0-autodetect)
    SHORT_FORM = ida_name.MNG_SHORT_FORM  # Composite flag for short form
    LONG_FORM = ida_name.MNG_LONG_FORM  # Composite flag for long form
    CALC_VALID = ida_name.MNG_CALC_VALID  # Valid flags for calculation mode


@decorate_all_methods(check_db_open)
class Names(DatabaseEntity):
    """
    Provides access to symbol and label management in the IDA database.

    Can be used to iterate over all names in the opened database.

    Args:
        database: Reference to the active IDA database.
    """

    def __init__(self, database: Database) -> None:
        super().__init__(database)

    def __iter__(self) -> Iterator[tuple[ea_t, str]]:
        return self.get_all()

    def __getitem__(self, index: int) -> Tuple[ea_t, str] | None:
        return self.get_at_index(index)

    def __len__(self) -> int:
        """
        Returns the total number of named elements in the database.

        Returns:
            The number of named elements.
        """
        return self.get_count()

    def get_count(self) -> int:
        """
        Retrieves the total number of named elements in the database.

        Returns:
            The number of named elements.
        """
        return ida_name.get_nlist_size()

    def get_at_index(self, index: int) -> Tuple[ea_t, str] | None:
        """
        Retrieves the named element at the specified index.

        Args:
            index: Index of the named element to retrieve.

        Returns:
            A tuple (effective address, name) at the given index.
            In case of error, returns None.
        """
        if index >= 0 and index < ida_name.get_nlist_size():
            return ida_name.get_nlist_ea(index), ida_name.get_nlist_name(index)
        raise IndexError(f'Name index {index} out of range [0, {self.get_count()})')

    def get_at(self, ea: ea_t) -> Optional[str]:
        """
        Retrieves the name at the specified address.

        Args:
            ea: The effective address.

        Returns:
            The name string if it exists, None otherwise.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        return ida_name.get_name(ea)

    def get_all(self) -> Iterator[Tuple[ea_t, str]]:
        """
        Returns an iterator over all named elements in the database.

        Returns:
            An iterator over (address, name) tuples.
        """
        index = 0
        while index < ida_name.get_nlist_size():
            yield ida_name.get_nlist_ea(index), ida_name.get_nlist_name(index)
            index += 1

    def set_name(
        self, ea: ea_t, name: str, flags: Union[int, SetNameFlags] = SetNameFlags.NOCHECK
    ) -> bool:
        """
        Set or delete name of an item at the specified address.

        Args:
            ea: Linear address.
            name: New name. Empty string to delete name.
            flags: Set name flags (SetNameFlags enum or raw int).

        Returns:
            True if successful, False otherwise.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        return ida_name.set_name(ea, name, flags)

    def force_name(
        self, ea: ea_t, name: str, flags: Union[int, SetNameFlags] = SetNameFlags.NOCHECK
    ) -> bool:
        """
        Force set a name, trying variations if the name already exists.

        Args:
            ea: Linear address.
            name: New name.
            flags: Set name flags (SetNameFlags enum or raw int).

        Returns:
            True if successful, False otherwise.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        return ida_name.force_name(ea, name, flags)

    def delete(self, ea: ea_t) -> bool:
        """
        Delete name at the specified address.

        Args:
            ea: Linear address.

        Returns:
            True if successful, False otherwise.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        return ida_name.del_global_name(ea)

    def is_valid_name(self, name: str) -> bool:
        """
        Check if a name is a valid user defined name.

        Args:
            name: Name to validate.

        Returns:
            True if valid, False otherwise.
        """
        return ida_name.is_uname(name)

    def is_public_name(self, ea: ea_t) -> bool:
        """
        Check if name at address is public.

        Args:
            ea: Linear address.

        Returns:
            True if public, False otherwise.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        return ida_name.is_public_name(ea)

    def make_name_public(self, ea: ea_t) -> None:
        """
        Make name at address public.

        Args:
            ea: Linear address.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        ida_name.make_name_public(ea)

    def make_name_non_public(self, ea: ea_t) -> None:
        """
        Make name at address non-public.

        Args:
            ea: Linear address.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        ida_name.make_name_non_public(ea)

    def is_weak_name(self, ea: ea_t) -> bool:
        """
        Check if name at address is weak.

        Args:
            ea: Linear address.

        Returns:
            True if weak, False otherwise.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        return ida_name.is_weak_name(ea)

    def make_name_weak(self, ea: ea_t) -> None:
        """
        Make name at address weak.

        Args:
            ea: Linear address.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        ida_name.make_name_weak(ea)

    def make_name_non_weak(self, ea: ea_t) -> None:
        """
        Make name at address non-weak.

        Args:
            ea: Linear address.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        ida_name.make_name_non_weak(ea)

    def get_demangled_name(
        self, ea: ea_t, inhibitor: Union[int, DemangleFlags] = 0, demform: int = 0
    ) -> Optional[str]:
        """
        Get demangled name at address.

        Args:
            ea: Linear address.
            inhibitor: Demangling inhibitor flags (DemangleFlags enum or raw int).
            demform: Demangling form flags.

        Returns:
            Demangled name or None if not available.

        Raises:
            InvalidEAError: If the effective address is invalid.
        """
        if not self.database.is_valid_ea(ea):
            raise InvalidEAError(ea)
        return ida_name.get_demangled_name(ea, inhibitor, demform)

    def demangle_name(self, name: str, disable_mask: Union[int, DemangleFlags] = 0) -> str:
        """
        Demangle a mangled name.

        Args:
            name: Mangled name to demangle.
            disable_mask: Bits to inhibit parts of demangled name (DemangleFlags enum or raw int).

        Returns:
            Demangled name or original name if demangling failed.
        """
        return ida_name.demangle_name(name, disable_mask)
