import pytest
from reccmp.types import EntityType
from reccmp.compare.db import EntityDb
from reccmp.compare.asm.replacement import (
    create_name_lookup,
    NameReplacementProtocol,
)


@pytest.fixture(name="db")
def fixture_db() -> EntityDb:
    return EntityDb()


def create_lookup(
    db, addrs: dict[int, int] | None = None, is_orig: bool = True
) -> NameReplacementProtocol:
    if addrs is None:
        addrs = {}

    def bin_lookup(addr: int) -> int | None:
        return addrs.get(addr)

    if is_orig:
        return create_name_lookup(db.get_by_orig, bin_lookup, "orig_addr")

    return create_name_lookup(db.get_by_recomp, bin_lookup, "recomp_addr")


####


def test_name_replacement(db):
    """Should return a name for an entity that has one.
    Return None if there are no name attributes set or the entity does not exist."""
    with db.batch() as batch:
        batch.set_orig(100, name="Test")
        batch.set_orig(200, computed_name="Hello")
        batch.set_orig(300)  # No name

    lookup = create_lookup(db)
    entity_100 = lookup(100)
    entity_200 = lookup(200)

    assert entity_100 is not None
    # Using "in" here because the returned string may contain other information.
    # e.g. the entity type
    assert "Test" in entity_100

    assert entity_200 is not None
    assert "Hello" in entity_200

    assert lookup(300) is None


def test_name_hierarchy(db):
    """Use the "best" entity name. Currently there are only two.
    'computed_name' is preferred over just 'name'."""
    with db.batch() as batch:
        batch.set_orig(100, name="Test", computed_name="Hello")

    lookup = create_lookup(db)
    entity = lookup(100)

    assert entity is not None

    # Should prefer 'computed_name' over 'name'
    assert "Hello" in entity
    assert "Test" not in entity


@pytest.mark.parametrize("entity_type", (EntityType.DATA, EntityType.OFFSET))
def test_offset_name(db: EntityDb, entity_type: EntityType):
    """For some entities (i.e. variables) we will return a name if the search address
    is inside the address range of the entity. This is determined by the size attribute.
    """
    with db.batch() as batch:
        batch.set_orig(100, name="Hello", type=entity_type, size=10)

    lookup = create_lookup(db)

    assert lookup(100) is not None
    assert lookup(101) is not None

    # Outside the range = no name
    assert lookup(110) is None


def test_offset_name_non_variables(db):
    """Do not return an offset name for non-variable entities. (e.g. functions)."""
    with db.batch() as batch:
        batch.set_orig(100, name="Hello", type=EntityType.FUNCTION, size=10)
        batch.set_orig(200, name="Hello", size=10)  # No type

    lookup = create_lookup(db)

    assert lookup(100) is not None
    assert lookup(101) is None

    assert lookup(200) is not None
    assert lookup(201) is None


def test_offset_name_no_size(db):
    """An enity with no size attribute is considered to have size=0.
    Meaning: match only against the address value."""
    with db.batch() as batch:
        batch.set_orig(100, name="Hello", type=EntityType.DATA)

    lookup = create_lookup(db)

    assert lookup(100) is not None
    assert lookup(101) is None


def test_exact_restriction(db):
    """If exact=True, return a name only if the entity's address matches the search address.
    Otherwise we might return a name if the entity contains the search address."""
    with db.batch() as batch:
        batch.set_orig(100, name="Hello", type=EntityType.DATA, size=10)

    lookup = create_lookup(db)

    assert lookup(100, exact=True) is not None
    assert lookup(101, exact=True) is None

    # Proof that the exact parameter controls whether we get a name.
    assert lookup(101, exact=False) is not None


def test_indirect_function(db):
    """An instruction like `call dword ptr [0x1234]` means that we call the function
    whose address is at address 0x1234. This is an indirect lookup."""
    with db.batch() as batch:
        batch.set_orig(100, name="Hello", type=EntityType.FUNCTION)

    # Mock lookup so we will read 100 from address 200.
    lookup = create_lookup(db, {200: 100})

    # No entity at 200
    assert lookup(200) is None
    entity = lookup(200, indirect=True)
    assert entity is not None

    # Imitating ghidra asm display. Not every indirect lookup gets the arrow.
    assert "->" in entity


def test_indirect_function_variable(db):
    """If the indirect call instruction has the address of a variable in our database,
    prefer the variable name rather than reading the pointer."""
    with db.batch() as batch:
        batch.set_orig(100, name="Hello", type=EntityType.FUNCTION)
        batch.set_orig(200, name="Test", type=EntityType.DATA)

    # Mock lookup so we will read 100 from address 200.
    lookup = create_lookup(db, {200: 100})

    name = lookup(200, indirect=True)
    assert name is not None
    assert "Hello" not in name
    assert "Test" in name
    assert "->" not in name


def test_indirect_import(db):
    """If we are indirectly calling an imported funtion, we should see the import_name
    attribute used in the result. This will probably contain the DLL and function name.
    """
    with db.batch() as batch:
        batch.set_orig(100, name="Hello", type=EntityType.IMPORT)

    # No mock needed here because we will not need to read any data.
    lookup = create_lookup(db)

    # Should use arrow to suggest indirect call.
    name = lookup(100, indirect=True)
    assert name is not None
    assert "Hello" in name
    assert "->" in name

    # No arrow.
    name = lookup(100, indirect=False)
    assert name is not None
    assert "Hello" in name
    assert "->" not in name


def test_import_without_name(db):
    """If the import entity doesn't have a name, the lookup should return None."""
    with db.batch() as batch:
        batch.set_orig(100, type=EntityType.IMPORT)  # No name

    lookup = create_lookup(db)

    # Should not follow the pointer because it is an RVA, not a real virtual address.
    assert lookup(100, indirect=True) is None
    assert lookup(100, indirect=False) is None


def test_indirect_failed_lookup(db):
    """In the general case (i.e. we do not use the base entity to get the name)
    if there is no entity at the pointer location, return None."""
    with db.batch() as batch:
        batch.set_orig(200, name="Hello", type=EntityType.FUNCTION)

    # Mock lookup so we will read 100 from address 200.
    lookup = create_lookup(db, {200: 100})

    # There is an entity at 200 but we can't use it to generate a name.
    # There is no entity at 100 (indirect location)
    assert lookup(200, indirect=False) is not None
    assert lookup(200, indirect=True) is None
