# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_navigator

from typing import Any, Dict, List, Optional, Set

from playwright.async_api import Page

from coreason_navigator.types import ElementBoundingBox
from coreason_navigator.utils.logger import logger

# Roles that are inherently interactive
INTERACTIVE_ROLES: Set[str] = {
    "button",
    "checkbox",
    "combobox",
    "link",
    "listbox",
    "menu",
    "menubar",
    "menuitem",
    "menuitemcheckbox",
    "menuitemradio",
    "radio",
    "radiogroup",
    "scrollbar",
    "searchbox",
    "slider",
    "spinbutton",
    "switch",
    "tab",
    "tablist",
    "textbox",
    "tree",
    "treegrid",
    "treeitem",
}

# Roles that are usually root containers and should be excluded even if focusable
EXCLUDED_ROLES: Set[str] = {
    "RootWebArea",
    "WebArea",
    "generic",  # unless focusable, handled by check
}


class AXTreeParser:
    """Parses the accessibility tree to find interactive elements with coordinates."""

    def __init__(self, page: Page):
        """
        Initializes the AXTreeParser.

        Args:
            page: The Playwright Page instance.
        """
        self.page = page

    async def get_interactive_elements(self) -> List[ElementBoundingBox]:
        """
        Fetches the full AX Tree via CDP, filters for interactive elements,
        resolves their bounding boxes, and returns a list of ElementBoundingBox.

        Returns:
            List[ElementBoundingBox]: A list of detected interactive elements.
        """
        try:
            # Connect to CDP
            session = await self.page.context.new_cdp_session(self.page)

            # Fetch the entire accessibility tree
            # This returns a flat list of nodes in 'nodes' key
            response = await session.send("Accessibility.getFullAXTree")
            nodes = response.get("nodes", [])

            # Create a map for quick lookup of children
            node_map = {node["nodeId"]: node for node in nodes}

            interactive_elements: List[ElementBoundingBox] = []

            for node in nodes:
                if self._is_interactive(node):
                    backend_id = node.get("backendDOMNodeId")
                    if not backend_id:
                        continue

                    bbox = await self._get_box_model(session, backend_id)
                    if bbox:
                        element_id = str(backend_id)  # Using backendID as a stable ID for this session

                        # Extract name (text)
                        text = self._get_node_name(node, node_map)
                        role = self._get_node_role(node)

                        interactive_elements.append(
                            ElementBoundingBox(
                                element_id=element_id,
                                x=bbox["x"],
                                y=bbox["y"],
                                width=bbox["width"],
                                height=bbox["height"],
                                tag_name=role,
                                text=text,
                                attributes={"role": role},
                            )
                        )

            return interactive_elements

        except Exception:
            logger.exception("Failed to fetch accessibility tree")
            return []

    def _is_interactive(self, node: Dict[str, Any]) -> bool:
        """
        Determines if a node is interactive based on role and properties.

        Args:
            node: The AX node dictionary.

        Returns:
            bool: True if interactive, False otherwise.
        """
        if node.get("ignored", False):
            return False

        role_obj = node.get("role", {})
        role_value = role_obj.get("value", "")

        # Explicitly exclude root roles to avoid noise
        if role_value in EXCLUDED_ROLES:
            # But wait, generic might be interactive if focusable?
            # RootWebArea is definitely not what we want to click on usually.
            if role_value in {"RootWebArea", "WebArea"}:
                return False

        # Check against allowlist
        if role_value in INTERACTIVE_ROLES:
            return True

        # Check if focusable (handles custom widgets, divs with tabindex)
        # Note: RootWebArea is often marked focusable, so we excluded it above.
        properties = node.get("properties", [])
        for prop in properties:
            if prop.get("name") == "focusable":
                # 'value' is a dict: {type: booleanOrUndefined, value: true}
                val = prop.get("value", {})
                if val.get("value") is True:
                    return True

        return False

    def _get_node_name(self, node: Dict[str, Any], node_map: Optional[Dict[str, Any]] = None) -> str:
        """
        Extracts the computed name/text of the node.

        Args:
            node: The AX node dictionary.
            node_map: Optional map of all nodes for child lookup.

        Returns:
            str: The name or text content of the node.
        """
        name_obj = node.get("name", {})
        name: str = name_obj.get("value", "")
        if name:
            return name

        # Fallback: check value property
        val_obj = node.get("value", {})
        val: Any = val_obj.get("value")
        if val and isinstance(val, str):
            return str(val)

        # Fallback for generic nodes: check children for StaticText
        if node_map and node.get("role", {}).get("value") == "generic":
            child_ids = node.get("childIds", [])
            child_texts = []
            for child_id in child_ids:
                child = node_map.get(child_id)
                if child and child.get("role", {}).get("value") == "StaticText":
                    child_name = child.get("name", {}).get("value", "")
                    if child_name:
                        child_texts.append(child_name)
            if child_texts:
                return " ".join(child_texts)

        return ""

    def _get_node_role(self, node: Dict[str, Any]) -> str:
        """
        Extracts the role of the node.

        Args:
            node: The AX node dictionary.

        Returns:
            str: The role value (e.g., 'button', 'link').
        """
        return str(node.get("role", {}).get("value", "unknown"))

    async def _get_box_model(self, session: Any, backend_node_id: int) -> Optional[Dict[str, float]]:
        """
        Fetches the quad for a given backend node ID via CDP.

        Args:
            session: The CDP session.
            backend_node_id: The backend DOM node ID.

        Returns:
            Optional[Dict[str, float]]: Dictionary with x, y, width, height, or None if hidden/detached.
        """
        try:
            # DOM.getBoxModel returns 'model' -> 'border', 'padding', 'content', 'margin' (all Quads)
            # Quad is [x1, y1, x2, y2, x3, y3, x4, y4]
            result = await session.send("DOM.getBoxModel", {"backendNodeId": backend_node_id})
            model = result.get("model", {})
            border_quad = model.get("border")  # Use border box

            if not border_quad or len(border_quad) < 8:
                return None

            # Convert quad to x, y, w, h
            xs = [border_quad[i] for i in range(0, 8, 2)]
            ys = [border_quad[i] for i in range(1, 8, 2)]

            x = min(xs)
            y = min(ys)
            width = max(xs) - x
            height = max(ys) - y

            if width <= 0 or height <= 0:
                return None

            return {"x": x, "y": y, "width": width, "height": height}

        except Exception:
            # Element might be hidden, detached, or display:none
            return None


async def get_interactive_elements(page: Page) -> List[ElementBoundingBox]:
    """
    Helper function to run the AXTreeParser.

    Args:
        page: The Playwright Page instance.

    Returns:
        List[ElementBoundingBox]: A list of detected interactive elements.
    """
    parser = AXTreeParser(page)
    return await parser.get_interactive_elements()
