"""
页面快照
支持 DOM 快照、截图、差异对比
"""
import json
import hashlib
import base64
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any, Set
from pathlib import Path
from datetime import datetime
from playwright.async_api import Page
from loguru import logger


@dataclass
class SnapshotConfig:
    """快照配置"""
    include_screenshot: bool = True
    include_dom: bool = True
    include_styles: bool = True
    include_scripts: bool = False
    screenshot_full_page: bool = False
    diff_ignore_attrs: List[str] = field(default_factory=lambda: ["data-testid", "aria-live"])
    diff_ignore_tags: List[str] = field(default_factory=lambda: ["script", "style", "noscript"])


@dataclass
class DOMSnapshot:
    """DOM 快照"""
    html: str
    structure: Dict[str, Any]  # 简化的 DOM 结构
    elements_count: int
    text_content: str
    styles: Dict[str, str] = field(default_factory=dict)
    scripts: List[str] = field(default_factory=list)
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


@dataclass
class SnapshotResult:
    """快照结果"""
    id: str
    url: str
    title: str
    dom: Optional[DOMSnapshot] = None
    screenshot: Optional[str] = None  # base64 编码
    screenshot_path: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())


@dataclass
class DiffResult:
    """差异对比结果"""
    has_changes: bool
    added_elements: List[str] = field(default_factory=list)
    removed_elements: List[str] = field(default_factory=list)
    modified_elements: List[Dict[str, Any]] = field(default_factory=list)
    text_changes: List[Dict[str, str]] = field(default_factory=list)
    style_changes: List[Dict[str, Any]] = field(default_factory=list)
    similarity_score: float = 1.0  # 0.0 - 1.0


class PageSnapshot:
    """
    页面快照工具
    支持创建、保存、加载和对比页面快照
    """

    def __init__(
        self,
        page: Page,
        config: Optional[SnapshotConfig] = None,
    ):
        self.page = page
        self.config = config or SnapshotConfig()
        self._snapshots: Dict[str, SnapshotResult] = {}

    # ==================== 创建快照 ====================

    async def capture(
        self,
        name: Optional[str] = None,
        save_screenshot: bool = True,
    ) -> SnapshotResult:
        """
        捕获页面快照

        Args:
            name: 快照名称
            save_screenshot: 是否保存截图

        Returns:
            SnapshotResult: 快照结果
        """
        snapshot_id = name or self._generate_id()
        url = self.page.url
        title = await self.page.title()

        result = SnapshotResult(
            id=snapshot_id,
            url=url,
            title=title,
            metadata={
                "viewport": self.page.viewport_size,
            },
        )

        # 捕获 DOM
        if self.config.include_dom:
            result.dom = await self._capture_dom()

        # 捕获截图
        if save_screenshot and self.config.include_screenshot:
            screenshot_data = await self.page.screenshot(
                full_page=self.config.screenshot_full_page
            )
            result.screenshot = base64.b64encode(screenshot_data).decode('utf-8')

        self._snapshots[snapshot_id] = result
        logger.debug(f"快照已捕获: {snapshot_id}")

        return result

    async def _capture_dom(self) -> DOMSnapshot:
        """捕获 DOM"""
        html = await self.page.content()

        # 获取简化的 DOM 结构
        structure = await self._get_simplified_structure()

        # 统计元素数量
        elements_count = await self.page.evaluate("""
            () => document.querySelectorAll('*').length
        """)

        # 获取文本内容
        text_content = await self.page.evaluate("""
            () => document.body ? document.body.innerText : ''
        """)

        # 获取样式
        styles = {}
        if self.config.include_styles:
            styles = await self._get_computed_styles()

        # 获取脚本
        scripts = []
        if self.config.include_scripts:
            scripts = await self.page.evaluate("""
                () => Array.from(document.scripts).map(s => s.src || s.textContent?.substring(0, 100))
            """)

        return DOMSnapshot(
            html=html,
            structure=structure,
            elements_count=elements_count,
            text_content=text_content,
            styles=styles,
            scripts=scripts,
        )

    async def _get_simplified_structure(self) -> Dict[str, Any]:
        """获取简化的 DOM 结构"""
        ignore_tags = self.config.diff_ignore_tags

        return await self.page.evaluate(f"""
            () => {{
                const ignoreTags = {json.dumps(ignore_tags)};
                function getElementInfo(el, depth = 0) {{
                    if (depth > 5 || !el.tagName) return null;
                    if (ignoreTags.includes(el.tagName.toLowerCase())) return null;

                    const info = {{
                        tag: el.tagName.toLowerCase(),
                        id: el.id || null,
                        classes: el.className ? el.className.split(' ').filter(c => c) : [],
                    }};

                    if (el.children.length > 0 && depth < 5) {{
                        info.children = Array.from(el.children)
                            .map(child => getElementInfo(child, depth + 1))
                            .filter(c => c !== null);
                    }}

                    return info;
                }}

                return getElementInfo(document.body);
            }}
        """)

    async def _get_computed_styles(self) -> Dict[str, str]:
        """获取计算样式"""
        return await self.page.evaluate("""
            () => {
                const styles = {};
                const body = document.body;
                if (!body) return styles;

                const computed = window.getComputedStyle(body);
                const important = [
                    'font-size', 'font-family', 'color', 'background-color',
                    'margin', 'padding', 'width', 'height'
                ];

                important.forEach(prop => {
                    styles[prop] = computed.getPropertyValue(prop);
                });

                return styles;
            }
        """)

    # ==================== 快照管理 ====================

    def get_snapshot(self, snapshot_id: str) -> Optional[SnapshotResult]:
        """获取快照"""
        return self._snapshots.get(snapshot_id)

    def list_snapshots(self) -> List[str]:
        """列出所有快照 ID"""
        return list(self._snapshots.keys())

    def delete_snapshot(self, snapshot_id: str) -> bool:
        """删除快照"""
        if snapshot_id in self._snapshots:
            del self._snapshots[snapshot_id]
            return True
        return False

    def clear_snapshots(self):
        """清除所有快照"""
        self._snapshots.clear()

    # ==================== 持久化 ====================

    async def save_snapshot(
        self,
        snapshot_id: str,
        file_path: str,
        include_screenshot: bool = True,
    ) -> bool:
        """
        保存快照到文件

        Args:
            snapshot_id: 快照 ID
            file_path: 保存路径
            include_screenshot: 是否包含截图

        Returns:
            bool: 是否保存成功
        """
        snapshot = self._snapshots.get(snapshot_id)
        if not snapshot:
            logger.warning(f"快照不存在: {snapshot_id}")
            return False

        try:
            Path(file_path).parent.mkdir(parents=True, exist_ok=True)

            data = {
                "id": snapshot.id,
                "url": snapshot.url,
                "title": snapshot.title,
                "dom": {
                    "html": snapshot.dom.html,
                    "structure": snapshot.dom.structure,
                    "elements_count": snapshot.dom.elements_count,
                    "text_content": snapshot.dom.text_content,
                    "styles": snapshot.dom.styles,
                } if snapshot.dom else None,
                "metadata": snapshot.metadata,
                "created_at": snapshot.created_at,
            }

            if include_screenshot and snapshot.screenshot:
                data["screenshot"] = snapshot.screenshot

            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)

            logger.info(f"快照已保存: {file_path}")
            return True

        except Exception as e:
            logger.error(f"保存快照失败: {e}")
            return False

    async def load_snapshot(self, file_path: str) -> Optional[SnapshotResult]:
        """
        从文件加载快照

        Args:
            file_path: 快照文件路径

        Returns:
            Optional[SnapshotResult]: 加载的快照
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)

            dom = None
            if data.get("dom"):
                dom = DOMSnapshot(
                    html=data["dom"]["html"],
                    structure=data["dom"]["structure"],
                    elements_count=data["dom"]["elements_count"],
                    text_content=data["dom"]["text_content"],
                    styles=data["dom"].get("styles", {}),
                )

            snapshot = SnapshotResult(
                id=data["id"],
                url=data["url"],
                title=data["title"],
                dom=dom,
                screenshot=data.get("screenshot"),
                metadata=data.get("metadata", {}),
                created_at=data["created_at"],
            )

            self._snapshots[snapshot.id] = snapshot
            return snapshot

        except Exception as e:
            logger.error(f"加载快照失败: {e}")
            return None

    # ==================== 差异对比 ====================

    async def compare(
        self,
        snapshot1_id: str,
        snapshot2_id: str,
    ) -> DiffResult:
        """
        对比两个快照

        Args:
            snapshot1_id: 第一个快照 ID
            snapshot2_id: 第二个快照 ID

        Returns:
            DiffResult: 差异结果
        """
        snap1 = self._snapshots.get(snapshot1_id)
        snap2 = self._snapshots.get(snapshot2_id)

        if not snap1 or not snap2:
            raise ValueError(f"快照不存在: {snapshot1_id}, {snapshot2_id}")

        return self._compare_snapshots(snap1, snap2)

    async def compare_with_current(
        self,
        snapshot_id: str,
    ) -> DiffResult:
        """
        将快照与当前页面对比

        Args:
            snapshot_id: 快照 ID

        Returns:
            DiffResult: 差异结果
        """
        stored = self._snapshots.get(snapshot_id)
        if not stored:
            raise ValueError(f"快照不存在: {snapshot_id}")

        # 捕获当前页面快照
        current = await self.capture("_temp_for_compare")

        result = self._compare_snapshots(stored, current)

        # 清理临时快照
        del self._snapshots["_temp_for_compare"]

        return result

    def _compare_snapshots(
        self,
        snap1: SnapshotResult,
        snap2: SnapshotResult,
    ) -> DiffResult:
        """对比两个快照"""
        result = DiffResult(has_changes=False)

        if not snap1.dom or not snap2.dom:
            return result

        # 对比 DOM 结构
        struct1 = snap1.dom.structure
        struct2 = snap2.dom.structure

        if struct1 and struct2:
            self._compare_structures(struct1, struct2, result)

        # 对比文本
        text1 = snap1.dom.text_content
        text2 = snap2.dom.text_content
        if text1 != text2:
            result.text_changes.append({
                "before": text1[:500],
                "after": text2[:500],
            })
            result.has_changes = True

        # 对比样式
        if snap1.dom.styles and snap2.dom.styles:
            style_diff = self._compare_styles(snap1.dom.styles, snap2.dom.styles)
            if style_diff:
                result.style_changes = style_diff
                result.has_changes = True

        # 计算相似度
        result.similarity_score = self._calculate_similarity(snap1, snap2, result)

        return result

    def _compare_structures(
        self,
        struct1: Dict[str, Any],
        struct2: Dict[str, Any],
        result: DiffResult,
        path: str = "body",
    ):
        """对比 DOM 结构"""
        # 对比标签
        if struct1.get("tag") != struct2.get("tag"):
            result.modified_elements.append({
                "path": path,
                "type": "tag_change",
                "before": struct1.get("tag"),
                "after": struct2.get("tag"),
            })
            result.has_changes = True
            return

        # 对比 ID
        if struct1.get("id") != struct2.get("id"):
            result.modified_elements.append({
                "path": path,
                "type": "id_change",
                "before": struct1.get("id"),
                "after": struct2.get("id"),
            })
            result.has_changes = True

        # 对比类名
        classes1 = set(struct1.get("classes", []))
        classes2 = set(struct2.get("classes", []))
        if classes1 != classes2:
            result.modified_elements.append({
                "path": path,
                "type": "class_change",
                "added": list(classes2 - classes1),
                "removed": list(classes1 - classes2),
            })
            result.has_changes = True

        # 对比子元素
        children1 = struct1.get("children", [])
        children2 = struct2.get("children", [])

        # 简单的数量对比
        if len(children1) != len(children2):
            result.modified_elements.append({
                "path": path,
                "type": "children_count_change",
                "before": len(children1),
                "after": len(children2),
            })
            result.has_changes = True

        # 递归对比子元素
        for i, (c1, c2) in enumerate(zip(children1, children2)):
            self._compare_structures(c1, c2, result, f"{path}/{c1.get('tag', '?')}[{i}]")

    def _compare_styles(
        self,
        styles1: Dict[str, str],
        styles2: Dict[str, str],
    ) -> List[Dict[str, Any]]:
        """对比样式"""
        changes = []
        all_keys = set(styles1.keys()) | set(styles2.keys())

        for key in all_keys:
            val1 = styles1.get(key, "")
            val2 = styles2.get(key, "")
            if val1 != val2:
                changes.append({
                    "property": key,
                    "before": val1,
                    "after": val2,
                })

        return changes

    def _calculate_similarity(
        self,
        snap1: SnapshotResult,
        snap2: SnapshotResult,
        diff: DiffResult,
    ) -> float:
        """计算相似度分数"""
        if not snap1.dom or not snap2.dom:
            return 0.0

        # 基于元素数量
        count_diff = abs(snap1.dom.elements_count - snap2.dom.elements_count)
        max_count = max(snap1.dom.elements_count, snap2.dom.elements_count)
        element_score = 1.0 - (count_diff / max_count) if max_count > 0 else 1.0

        # 基于修改数量
        change_count = (
            len(diff.added_elements) +
            len(diff.removed_elements) +
            len(diff.modified_elements) +
            len(diff.text_changes)
        )
        change_penalty = min(change_count * 0.05, 0.5)

        return max(0.0, min(1.0, element_score - change_penalty))

    # ==================== 辅助方法 ====================

    def _generate_id(self) -> str:
        """生成唯一 ID"""
        return f"snapshot_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{hashlib.md5(str(datetime.now().timestamp()).encode()).hexdigest()[:6]}"


def create_page_snapshot(
    page: Page,
    config: Optional[SnapshotConfig] = None,
) -> PageSnapshot:
    """创建页面快照实例"""
    return PageSnapshot(page, config)
