"""Element lookup and attribute mutation helpers."""

from __future__ import annotations

from typing import Dict, List

from lxml import etree

from tokenmoulds.emitters.xml_constants import A_NS


def find_element(
    root: etree._Element,
    xpath: str,
    namespaces: Dict[str, str] | None = None,
) -> etree._Element | None:
    """Find a single element using XPath with namespace support."""
    ns = namespaces or {"a": A_NS}
    return root.find(xpath, ns)


def find_elements(
    root: etree._Element,
    xpath: str,
    namespaces: Dict[str, str] | None = None,
) -> List[etree._Element]:
    """Find all elements matching XPath."""
    ns = namespaces or {"a": A_NS}
    return root.findall(xpath, ns)


def set_element_attr(
    root: etree._Element,
    xpath: str,
    attr_name: str,
    attr_value: str,
    namespaces: Dict[str, str] | None = None,
) -> bool:
    """Set attribute on an element found by XPath."""
    element = find_element(root, xpath, namespaces)
    if element is not None:
        element.set(attr_name, attr_value)
        return True
    return False
