#!/usr/bin/env python3
"""
page_dom_extract.py — Text-based DOM 简化工具

用法:
    python3 page_dom_extract.py <url>          # CLI 模式
    from page_dom_extract import html_to_simplified
    result = html_to_simplified("https://example.com")

功能:
    把 HTML 转为 LLM 可读的简化文本格式，参考 alibaba/page-agent 思路
    - 优先用 curl 拉取（快速，静态页面）
    - 数字缩进表示层级 [N]<tag>text</tag>
    - 交互元素自动检测并标注 aria role
    - 中文支持（html.unescape 处理实体）

依赖:
    pip install lxml requests
"""

import sys
import html
import subprocess
import re
from urllib.parse import urljoin, urlparse

try:
    from lxml import etree
except ImportError:
    print("Error: lxml not installed. Run: pip install lxml", file=sys.stderr)
    sys.exit(1)

try:
    import requests
except ImportError:
    requests = None  # optional fallback


# 交互元素标签
INTERACTIVE_TAGS = frozenset({"a", "button", "input", "select", "textarea", "details", "summary", "label"})

# 有 onclick/onchange 事件的标签（简化检测）
EVENT_ATTRS = frozenset({"onclick", "onchange", "onsubmit", "onkeydown", "onkeyup", "onfocus", "onblur"})

# aria role 映射到优先级
ARIA_INTERACTIVE_ROLES = frozenset({
    "button", "link", "checkbox", "radio", "tab", "menuitem",
    "menuitemcheckbox", "menuitemradio", "option", "switch",
    "toggle", "treeitem", "slider", "spinbutton", "textbox",
    "combobox", "listbox", "grid", "tree", "dialog",
})


def is_interactive_element(el: etree._Element) -> tuple[bool, str | None]:
    """
    检测元素是否为交互元素。
    返回 (is_interactive, aria_role_or_none)
    """
    tag = el.tag if isinstance(el.tag, str) else ""
    tag = tag.lower()

    # 1. 交互标签检测
    if tag in INTERACTIVE_TAGS:
        # input 特殊处理：type=button/submit reset/image 才算交互
        if tag == "input":
            input_type = el.get("type", "text").lower()
            if input_type in ("submit", "button", "reset", "image", "checkbox", "radio"):
                return True, None
            return False, None
        return True, None

    # 2. 事件属性检测
    for attr in EVENT_ATTRS:
        if el.get(attr):
            role = el.get("role", "")
            return True, role if role else None

    # 3. contenteditable 检测
    if el.get("contenteditable") in ("true", "editable"):
        role = el.get("role", "")
        return True, role if role else "contenteditable"

    # 4. aria role 检测
    role = el.get("role", "")
    if role and role.lower() in ARIA_INTERACTIVE_ROLES:
        return True, role

    return False, None


def get_text_content(el: etree._Element) -> str:
    """提取元素的文本内容，保留结构化 whitespace"""
    text = etree.tostring(el, method="text", encoding="unicode", with_tail=False)
    # 合并连续空白，保留必要换行
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n\s*", "\n", text).strip()
    return text


def simplify_element(el: etree._Element, depth: int) -> list[str]:
    """
    递归简化单个元素，返回行列表。
    depth 为当前层级（从 0 开始）
    """
    lines = []
    tag = el.tag if isinstance(el.tag, str) else None

    # 跳过注释和特殊节点
    if tag is None or isinstance(el, etree._Comment) or isinstance(el, etree._ProcessingInstruction):
        if el.tail:
            tail = html.unescape(el.tail.strip())
            if tail:
                lines.append(f"[{depth}]{tail}")
        return lines

    tag = tag.lower()
    is_interactive, aria_role = is_interactive_element(el)
    text = get_text_content(el)
    text = html.unescape(text)

    # 构建属性字符串
    attrs = []
    if is_interactive:
        if tag == "a" and el.get("href"):
            attrs.append(f'href="{el.get("href")}"')
        elif tag == "input":
            for attr in ("name", "value", "placeholder", "type"):
                if el.get(attr):
                    attrs.append(f'{attr}="{el.get(attr)}"')
        elif tag == "button":
            if el.get("type"):
                attrs.append(f'type="{el.get("type")}"')
        elif tag in ("select", "textarea"):
            if el.get("name"):
                attrs.append(f'name="{el.get("name")}"')

    if aria_role:
        attrs.append(f'role="{aria_role}"')

    attr_str = " " + " ".join(attrs) if attrs else ""

    if is_interactive:
        # 交互元素：显示 tag + 内容
        if text:
            lines.append(f"[{depth}]<{tag}{attr_str}>{text}</{tag}>")
        else:
            # 无内容但有 href 的 link 也显示
            if tag == "a" and el.get("href"):
                lines.append(f"[{depth}]<{tag}{attr_str}>链接</{tag}>")
            elif tag in ("button",):
                lines.append(f"[{depth}]<{tag}{attr_str}>按钮</{tag}>")
            elif tag == "input":
                lines.append(f"[{depth}]<{tag}{attr_str}/>")
    else:
        # 非交互元素：只显示 tag（无内容时省略）
        if text:
            # 检查是否有子元素产生的内容
            child_texts = []
            for child in el:
                if isinstance(child, etree._Element):
                    child_texts.extend(simplify_element(child, depth + 1))
            # 如果直接文本内容比子元素多，用直接内容
            if text and not child_texts:
                lines.append(f"[{depth}]<{tag}>{text}</{tag}>")
            elif child_texts:
                lines.append(f"[{depth}]<{tag}>")
                lines.extend(child_texts)
                lines.append(f"[{depth}]</{tag}>")
        # 无内容时不显示（省略）

    # 处理 tail 文本（同级后续内容）
    if el.tail:
        tail = html.unescape(el.tail.strip())
        if tail:
            lines.append(f"[{depth}]{tail}")

    return lines


def html_to_simplified(url: str) -> str:
    """
    主函数：将 URL 的 HTML 转为简化文本格式。

    Args:
        url: 目标网页 URL

    Returns:
        简化后的文本格式字符串
    """
    html_content = fetch_html(url)
    return parse_and_simplify(html_content, url)


def fetch_html(url: str) -> str:
    """用 curl 拉取页面内容（快速，静态页面优先）"""
    try:
        result = subprocess.run(
            ["curl", "-s", "-L", "--max-time", "30", "-A",
             "Mozilla/5.0 (compatible; page_dom_extract/1.0)",
             url],
            capture_output=True,
            text=True,
            timeout=35,
        )
        if result.returncode == 0 and result.stdout:
            return result.stdout
    except (subprocess.SubprocessError, FileNotFoundError) as e:
        print(f"curl failed: {e}", file=sys.stderr)

    # curl 失败时尝试 requests（如果可用）
    if requests:
        try:
            resp = requests.get(url, timeout=30, headers={
                "User-Agent": "Mozilla/5.0 (compatible; page_dom_extract/1.0)"
            })
            if resp.status_code == 200:
                return resp.text
        except Exception as e:
            print(f"requests fallback failed: {e}", file=sys.stderr)

    raise ValueError(f"Failed to fetch: {url}")


def parse_and_simplify(html_content: str, base_url: str = "") -> str:
    """解析 HTML 并转为简化格式"""
    try:
        # 尝试解析为 HTML（容忍破碎 HTML）
        parser = etree.HTMLParser(recover=True, encoding="utf-8")
        tree = etree.fromstring(html_content.encode("utf-8") if isinstance(html_content, str) else html_content, parser)
    except Exception as e:
        print(f"Parse error: {e}", file=sys.stderr)
        # 降级：返回原文
        return html.unescape(html_content)

    lines = ["[0]<html>"]

    if tree is not None:
        # 处理 <head> 和 <body>
        head = tree.find("head")
        if head is not None:
            lines.append("[1]<head>")
            title = head.find("title")
            if title is not None and title.text:
                lines.append(f"[2]<title>{html.unescape(title.text.strip())}</title>")
            lines.append("[1]</head>")

        body = tree.find("body")
        if body is not None:
            lines.append("[1]<body>")
            for child in body:
                lines.extend(simplify_element(child, 2))
            lines.append("[1]</body>")
        else:
            # 没有 body，整个树遍历
            for child in tree:
                if isinstance(child, etree._Element):
                    lines.extend(simplify_element(child, 1))

    lines.append("[0]</html>")
    return "\n".join(lines)


def main():
    """CLI 入口"""
    if len(sys.argv) < 2:
        print("Usage: python3 page_dom_extract.py <url>", file=sys.stderr)
        sys.exit(1)

    url = sys.argv[1]

    try:
        result = html_to_simplified(url)
        print(result)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()