import re
import xml.etree.ElementTree as ET
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple


@dataclass
class SpatialNode:
    """A single node in the Spatial Graph, representing a UI element and its geometry."""

    bounds: Tuple[int, int, int, int]  # (x1, y1, x2, y2)
    node_id: str = ""
    class_name: str = ""
    text: str = ""
    content_desc: str = ""
    resource_id: str = ""
    clickable: bool = False
    scrollable: bool = False

    # Spatial Properties
    children: List["SpatialNode"] = field(default_factory=list)
    parent: Optional["SpatialNode"] = None

    @property
    def x1(self) -> int:
        return self.bounds[0]

    @property
    def y1(self) -> int:
        return self.bounds[1]

    @property
    def x2(self) -> int:
        return self.bounds[2]

    @property
    def y2(self) -> int:
        return self.bounds[3]

    @property
    def width(self) -> int:
        return self.x2 - self.x1

    @property
    def height(self) -> int:
        return self.y2 - self.y1

    @property
    def center_x(self) -> int:
        return self.x1 + (self.width // 2)

    @property
    def center_y(self) -> int:
        return self.y1 + (self.height // 2)

    @property
    def area(self) -> int:
        return self.width * self.height

    def contains(self, other: "SpatialNode") -> bool:
        """Returns True if this node completely encompasses the other node geometrically."""
        return self.x1 <= other.x1 and self.y1 <= other.y1 and self.x2 >= other.x2 and self.y2 >= other.y2

    def intersects(self, other: "SpatialNode") -> bool:
        """Returns True if this node's bounding box overlaps with the other's bounding box."""
        if self.x1 >= other.x2 or other.x1 >= self.x2:
            return False
        if self.y1 >= other.y2 or other.y1 >= self.y2:
            return False
        return True

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.node_id,
            "class": self.class_name,
            "text": self.text,
            "content_desc": self.content_desc,
            "resource_id": self.resource_id,
            "bounds": self.bounds,
            "clickable": self.clickable,
            "scrollable": self.scrollable,
            "center": (self.center_x, self.center_y),
        }


class SpatialParser:
    """
    Parses Android UI XML into a structured 2D Spatial Tree.
    Calculates parent-child relationships structurally, not just based on XML nesting.
    """

    def __init__(self):
        self._node_counter = 0

    def parse(self, xml_string: str) -> Optional[SpatialNode]:
        """Parses the raw XML dump into a Spatial Graph."""
        try:
            clean_xml = re.sub(r"<\?xml.*?\?>", "", xml_string).strip()
            if not clean_xml:
                return None
            root_elem = ET.fromstring(clean_xml)

            # 1. First Pass: Create flat list of spatial nodes
            all_nodes = []
            self._flatten_xml(root_elem, all_nodes)

            if not all_nodes:
                return None

            # 2. Second Pass: Reconstruct tree based on strict spatial containment
            # Sort nodes by area descending (largest first)
            all_nodes.sort(key=lambda n: n.area, reverse=True)

            root_node = all_nodes[0]

            for i in range(1, len(all_nodes)):
                child = all_nodes[i]
                # Find the smallest node that contains this child
                # Since we sorted by area descending, we search backwards to find the tightest fit
                parent_found = False
                for j in range(i - 1, -1, -1):
                    potential_parent = all_nodes[j]
                    if potential_parent.contains(child):
                        potential_parent.children.append(child)
                        child.parent = potential_parent
                        parent_found = True
                        break

                # Fallback to root if no parent found (floating node)
                if not parent_found and child != root_node:
                    root_node.children.append(child)
                    child.parent = root_node

            return root_node

        except ET.ParseError:
            return None

    def _flatten_xml(self, element: ET.Element, nodes_list: List[SpatialNode]):
        """Recursively traverses the XML and creates a flat list of SpatialNodes."""
        attrib = element.attrib

        bounds_str = attrib.get("bounds", "")
        match = re.match(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]", bounds_str)

        if match:
            left, top, right, bottom = map(int, match.groups())

            # Filter zero-area nodes early
            if right > left and bottom > top:
                self._node_counter += 1
                text_val = attrib.get("text", "").strip()
                hint_val = attrib.get("hint", "").strip()
                if not text_val and hint_val:
                    text_val = hint_val

                node = SpatialNode(
                    node_id=f"n_{self._node_counter}",
                    class_name=attrib.get("class", ""),
                    text=text_val,
                    content_desc=attrib.get("content-desc", "").strip(),
                    resource_id=attrib.get("resource-id", "").strip(),
                    bounds=(left, top, right, bottom),
                    clickable=attrib.get("clickable", "false") == "true",
                    scrollable=attrib.get("scrollable", "false") == "true",
                )
                nodes_list.append(node)

        for child in element:
            self._flatten_xml(child, nodes_list)

    def get_all_nodes(self, root: SpatialNode) -> List[SpatialNode]:
        """Flattens the Spatial Tree into a list for easy filtering."""
        result = [root]
        for child in root.children:
            result.extend(self.get_all_nodes(child))
        return result

    def get_clickable_nodes(self, root: SpatialNode) -> List[SpatialNode]:
        """Returns all nodes that are clickable or have strong semantic meaning."""
        all_nodes = self.get_all_nodes(root)
        clickables = []

        for n in all_nodes:
            has_semantic = bool(n.text or n.content_desc)
            semantic_res = n.resource_id and any(
                x in n.resource_id.lower() for x in ["button", "tab", "icon", "action", "menu", "imageview"]
            )

            if n.clickable or n.scrollable or semantic_res or (has_semantic and n.area < 500000 and n.area > 0):
                # Filter out pure massive containers (like whole screen) if they aren't explicitly clickable
                if not n.clickable and not n.scrollable and n.area > 2000000:
                    continue
                # Also exclude if it's just a ViewGroup with a description but no action
                if not n.clickable and n.class_name == "android.view.ViewGroup":
                    continue
                clickables.append(n)

        return clickables
