# mypy: ignore-errors
import copy
import json
import re
from typing import Dict, List, Set, Tuple, Union

from lxml import etree

# Type aliases for better readability
DOMElement = etree._Element
DOMTree = etree._Element
NodeID = str
IDMapping = Dict[str, int]
TreeReprResult = Tuple[str, IDMapping]


def get_descendants(
    node: DOMElement, max_depth: int, current_depth: int = 0
) -> List[DOMElement]:
    """Get descendants of a node up to max_depth.

    Args:
        node: The DOM element to get descendants for
        max_depth: Maximum depth to traverse
        current_depth: Current depth in the traversal (used recursively)

    Returns:
        List of descendant DOM elements
    """
    if current_depth > max_depth:
        return []

    descendants = []
    for child in node:
        descendants.append(child)
        descendants.extend(get_descendants(child, max_depth, current_depth + 1))

    return descendants


def simplify_attributes(
    node: DOMElement, max_value_length: int = 5, max_length: int = 20
) -> None:
    """Get attribute representation for a node.

    Extracts relevant attributes from a DOM node and creates a 'meta' attribute
    containing a simplified representation of the node's attributes.

    Args:
        node: The DOM element to process
        max_value_length: Maximum number of words to include from each attribute value
        max_length: Maximum number of words to include in the meta attribute

    Returns:
        None - modifies the node in place
    """
    # Get attribute values in order of importance
    attr_values_set: Set[str] = set()
    attr_values = ""

    # List of attributes to extract, in order of importance
    important_attrs = [
        "role",
        "aria_role",
        "type",
        "alt",
        "aria_description",
        "aria_label",
        "label",
        "title",
        "name",
        "text_value",
        "value",
        "placeholder",
        "input_checked",
        "input_value",
        "option_selected",
        "class",
    ]

    # Extract values from attributes
    for attr in important_attrs:
        if attr in node.attrib and node.attrib[attr] is not None:
            value = node.attrib[attr].lower()

            # Skip less meaningful values
            less_meaningful = [
                "hidden",
                "none",
                "presentation",
                "null",
                "undefined",
            ]

            if value in less_meaningful or value.startswith("http"):
                continue

            # For normal values, split into words and limit length
            value_words = value.split()
            filtered_words = [w for w in value_words if len(w) < 15][:max_value_length]
            value = " ".join(filtered_words)

            # Add to result if not already included
            if value and value not in attr_values_set:
                attr_values_set.add(value)
                attr_values += value + " "

    # Save the backend_node_id before clearing attributes
    uid = node.attrib.get("backend_node_id", "")

    # Clear all attributes
    node.attrib.clear()

    # Restore the node ID as 'id' attribute
    if uid:
        node.attrib["id"] = uid

    # Add meta attribute with limited length
    if attr_values:
        node.attrib["meta"] = " ".join(attr_values.split()[:max_length])


def prune_tree(
    dom_tree: DOMTree,
    candidate_set: Set[str],
    max_depth: int = 5,
    max_children: int = 50,
    max_sibling: int = 3,
) -> DOMTree:
    """Prune DOM tree to keep only relevant nodes.

    Reduces the size of a DOM tree by keeping only nodes that are:
    1. Candidates (in candidate_set)
    2. Ancestors of candidates
    3. Descendants of candidates (up to max_depth)
    4. Siblings of candidates (within max_sibling range)

    Args:
        dom_tree: The DOM tree to prune
        candidate_set: Set of node IDs to keep
        max_depth: Maximum depth of descendants to keep
        max_children: Maximum number of children to keep
        max_sibling: Maximum number of siblings to keep on each side

    Returns:
        Pruned DOM tree
    """
    # Process candidate IDs to handle both string and JSON string formats
    processed_candidate_set: Set[str] = set()
    for candidate_id in candidate_set:
        # Handle JSON string candidates
        try:
            parsed_id = json.loads(candidate_id)
            if isinstance(parsed_id, str):
                processed_candidate_set.add(parsed_id)
                continue
        except (json.JSONDecodeError, TypeError):
            pass

        # Handle regular string candidates
        processed_candidate_set.add(str(candidate_id))

    # Set of node IDs to keep
    nodes_to_keep: Set[str] = set()

    # Find all candidate nodes in the tree
    for candidate_id in processed_candidate_set:
        try:
            # Find the candidate node - use original dom_tree to match reference implementation
            candidate_node = dom_tree.xpath(f'//*[@backend_node_id="{candidate_id}"]')[
                0
            ]

            # Add the candidate node ID to the keep set
            nodes_to_keep.add(candidate_node.attrib["backend_node_id"])

            # Add all ancestors
            nodes_to_keep.update(
                [
                    x.attrib.get("backend_node_id", "")
                    for x in candidate_node.xpath("ancestor::*")
                ]
            )

            # Add descendants with max depth
            nodes_to_keep.update(
                [
                    x.attrib.get("backend_node_id", "")
                    for x in get_descendants(candidate_node, max_depth)
                ][:max_children]
            )

            # Add siblings within range
            parent = candidate_node.getparent()
            if parent is not None:
                siblings = [x for x in parent.getchildren() if x.tag != "text"]
                try:
                    idx_in_sibling = siblings.index(candidate_node)
                    nodes_to_keep.update(
                        [
                            x.attrib.get("backend_node_id", "")
                            for x in siblings[
                                max(0, idx_in_sibling - max_sibling) : idx_in_sibling
                                + max_sibling
                                + 1
                            ]
                        ]
                    )
                except ValueError:
                    # If candidate_node is not in siblings, skip sibling processing
                    pass
        except IndexError:
            # Skip if candidate node not found
            continue

    # Clone the tree
    new_tree = copy.deepcopy(dom_tree)

    # Remove nodes not in nodes_to_keep (process in reverse to avoid modifying during iteration)
    for node in new_tree.xpath("//*")[::-1]:
        if node.tag != "text":
            is_keep = node.attrib.get("backend_node_id", "") in nodes_to_keep
            is_candidate = (
                node.attrib.get("backend_node_id", "") in processed_candidate_set
            )
        else:
            parent = node.getparent()
            if parent is not None:
                is_keep = parent.attrib.get("backend_node_id", "") in nodes_to_keep
                is_candidate = (
                    parent.attrib.get("backend_node_id", "") in processed_candidate_set
                )
            else:
                is_keep = False
                is_candidate = False

        if not is_keep and node.getparent() is not None:
            node.getparent().remove(node)
        else:
            # Remove backend_node_id for all non-candidate nodes, following original Mind2Web approach
            if not is_candidate or node.tag == "text":
                node.attrib.pop("backend_node_id", None)

            # Simplify nodes with no attributes and only one child
            if (
                len(node.attrib) == 0
                and not any([x.tag == "text" for x in node.getchildren()])
                and node.getparent() is not None
                and node.tag != "text"
                and len(node.getchildren()) <= 1
            ):
                # Insert all children into parent
                for child in node.getchildren():
                    node.addprevious(child)
                node.getparent().remove(node)

    return new_tree


def get_tree_repr(
    tree: Union[str, DOMTree],
    max_value_length: int = 5,
    max_length: int = 20,
    id_mapping: IDMapping = None,
    keep_html_brackets: bool = False,
) -> TreeReprResult:
    """Get string representation of DOM tree.

    Converts a DOM tree to a string representation, with options to control
    the format and length of the output.

    Args:
        tree: The DOM tree to convert (either as a string or Element)
        max_value_length: Maximum number of words to include from each attribute value
        max_length: Maximum number of words to include in meta attributes
        id_mapping: Dictionary mapping backend_node_ids to sequential integers
        keep_html_brackets: Whether to keep HTML brackets in the output

    Returns:
        Tuple of (tree_representation_string, updated_id_mapping)
    """
    # Initialize id_mapping if None
    if id_mapping is None:
        id_mapping = {}
    # Parse string to tree if needed
    if isinstance(tree, str):
        try:
            tree = etree.fromstring(tree)
        except etree.XMLSyntaxError as e:
            # Handle malformed XML
            raise ValueError(f"Invalid XML: {e}")
    else:
        # Make a deep copy to avoid modifying the original
        tree = copy.deepcopy(tree)

    # Process all nodes in the tree
    for node in tree.xpath("//*"):
        if node.tag != "text":
            # Process regular nodes
            if "backend_node_id" in node.attrib:
                # Map backend_node_id to a sequential integer
                node_id = node.attrib["backend_node_id"]
                if node_id not in id_mapping:
                    id_mapping[node_id] = len(id_mapping)
                node.attrib["backend_node_id"] = str(id_mapping[node_id])

            # Get attribute representation
            simplify_attributes(node, max_value_length, max_length)
        else:
            # Process text nodes
            if node.text:
                # Limit text length
                node.text = " ".join(node.text.split()[:max_length])

    # Convert tree to string
    try:
        tree_repr = etree.tostring(tree, encoding="unicode")
    except Exception as e:
        raise ValueError(f"Failed to convert tree to string: {e}")

    # Clean up the representation
    tree_repr = tree_repr.replace('"', " ")
    tree_repr = (
        tree_repr.replace("meta= ", "").replace("id= ", "id=").replace(" >", ">")
    )

    # Replace text tags with their content
    tree_repr = re.sub(r"<text>(.*?)</text>", r"\1", tree_repr)

    # Optionally replace HTML brackets with parentheses
    if not keep_html_brackets:
        tree_repr = tree_repr.replace("/>", "$/$>")
        tree_repr = re.sub(r"</(.+?)>", r")", tree_repr)
        tree_repr = re.sub(r"<(.+?)>", r"(\1", tree_repr)
        tree_repr = tree_repr.replace("$/$", ")")

    # Unescape HTML entities
    html_escape_table = [
        ("&quot;", '"'),
        ("&amp;", "&"),
        ("&lt;", "<"),
        ("&gt;", ">"),
        ("&nbsp;", " "),
        ("&ndash;", "-"),
        ("&rsquo;", "'"),
        ("&lsquo;", "'"),
        ("&ldquo;", '"'),
        ("&rdquo;", '"'),
        ("&#39;", "'"),
        ("&#40;", "("),
        ("&#41;", ")"),
    ]
    for k, v in html_escape_table:
        tree_repr = tree_repr.replace(k, v)

    # Normalize whitespace
    tree_repr = re.sub(r"\s+", " ", tree_repr).strip()

    return tree_repr, id_mapping
