"""
基础数据提取器
从网页中提取结构化数据
"""
import re
import json
from typing import Optional, List, Dict, Any
from playwright.async_api import Page
from loguru import logger

from .types import ExtractionRule, ExtractionResult


class DataExtractor:
    """
    数据提取器
    从网页中提取结构化数据
    """

    def __init__(self, page: Page):
        self.page = page
        self._transforms = {
            "int": lambda x: int(x) if x else 0,
            "float": lambda x: float(x) if x else 0.0,
            "str": lambda x: str(x).strip() if x else "",
            "bool": lambda x: bool(x) if x else False,
            "lower": lambda x: str(x).lower() if x else "",
            "upper": lambda x: str(x).upper() if x else "",
            "strip": lambda x: str(x).strip() if x else "",
            "price": lambda x: float(re.sub(r'[^\d.]', '', str(x))) if x else 0.0,
            "number": lambda x: int(re.sub(r'[^\d]', '', str(x))) if x else 0,
        }

    def add_transform(self, name: str, func):
        """添加转换函数"""
        self._transforms[name] = func

    async def extract_text(self, selector: str) -> Optional[str]:
        """提取文本"""
        try:
            element = self.page.locator(selector)
            text = await element.text_content()
            return text.strip() if text else None
        except:
            return None

    async def extract_texts(self, selector: str) -> List[str]:
        """提取多个文本 - 批量优化版"""
        try:
            texts = await self.page.locator(selector).evaluate_all(
                "elements => elements.map(el => el.textContent?.trim() || '').filter(t => t)"
            )
            return texts
        except:
            return []

    async def extract_attribute(
        self,
        selector: str,
        attribute: str,
    ) -> Optional[str]:
        """提取属性值"""
        try:
            element = self.page.locator(selector)
            return await element.get_attribute(attribute)
        except:
            return None

    async def extract_attributes(
        self,
        selector: str,
        attribute: str,
    ) -> List[str]:
        """提取多个属性值 - 批量优化版"""
        try:
            values = await self.page.locator(selector).evaluate_all(
                f"elements => elements.map(el => el.getAttribute('{attribute}')).filter(v => v)"
            )
            return values
        except:
            return []

    async def extract_links(self, base_selector: str = "body") -> List[Dict[str, str]]:
        """提取所有链接"""
        try:
            links = await self.page.evaluate(f"""
                () => {{
                    const links = [];
                    document.querySelectorAll('{base_selector} a[href]').forEach(a => {{
                        links.push({{
                            text: a.textContent.trim(),
                            href: a.href,
                            title: a.title || '',
                        }});
                    }});
                    return links;
                }}
            """)
            return links
        except:
            return []

    async def extract_images(self, base_selector: str = "body") -> List[Dict[str, str]]:
        """提取所有图片"""
        try:
            images = await self.page.evaluate(f"""
                () => {{
                    const images = [];
                    document.querySelectorAll('{base_selector} img[src]').forEach(img => {{
                        images.push({{
                            src: img.src,
                            alt: img.alt || '',
                            title: img.title || '',
                            width: img.naturalWidth,
                            height: img.naturalHeight,
                        }});
                    }});
                    return images;
                }}
            """)
            return images
        except:
            return []

    async def extract_table(
        self,
        selector: str = "table",
        header_row: int = 0,
    ) -> List[Dict[str, str]]:
        """
        提取表格数据

        Args:
            selector: 表格选择器
            header_row: 表头行索引

        Returns:
            List[Dict]: 表格数据列表
        """
        try:
            table_data = await self.page.evaluate(f"""
                () => {{
                    const table = document.querySelector('{selector}');
                    if (!table) return [];

                    const rows = table.querySelectorAll('tr');
                    const headers = [];
                    const data = [];

                    // 获取表头
                    const headerCells = rows[{header_row}].querySelectorAll('th, td');
                    headerCells.forEach(cell => {{
                        headers.push(cell.textContent.trim());
                    }});

                    // 获取数据行
                    for (let i = {header_row + 1}; i < rows.length; i++) {{
                        const cells = rows[i].querySelectorAll('td');
                        const rowData = {{}};
                        cells.forEach((cell, index) => {{
                            if (index < headers.length) {{
                                rowData[headers[index]] = cell.textContent.trim();
                            }}
                        }});
                        if (Object.keys(rowData).length > 0) {{
                            data.push(rowData);
                        }}
                    }}

                    return data;
                }}
            """)
            return table_data
        except:
            return []

    async def extract_list(
        self,
        container_selector: str,
        item_selector: str,
        fields: Dict[str, str],
    ) -> List[Dict[str, Any]]:
        """
        提取列表数据 - 批量优化版

        Args:
            container_selector: 容器选择器
            item_selector: 列表项选择器
            fields: 字段映射 {字段名: 选择器}

        Returns:
            List[Dict]: 数据列表
        """
        try:
            # 将字段映射转为 JS 可用的格式
            fields_json = json.dumps(fields)

            results = await self.page.evaluate(f"""
                () => {{
                    const container = document.querySelector('{container_selector}');
                    if (!container) return [];

                    const items = container.querySelectorAll('{item_selector}');
                    const fields = {fields_json};
                    const results = [];

                    items.forEach(item => {{
                        const row = {{}};
                        for (const [fieldName, fieldSelector] of Object.entries(fields)) {{
                            try {{
                                const el = item.querySelector(fieldSelector);
                                row[fieldName] = el ? el.textContent.trim() : null;
                            }} catch {{
                                row[fieldName] = null;
                            }}
                        }}
                        results.push(row);
                    }});

                    return results;
                }}
            """)
            return results
        except:
            return []

    async def extract_with_rules(
        self,
        rules: List[ExtractionRule],
        base_selector: str = "",
    ) -> ExtractionResult:
        """
        使用规则提取数据 - 批量优化版

        Args:
            rules: 提取规则列表
            base_selector: 基础选择器

        Returns:
            ExtractionResult: 提取结果
        """
        data = {}
        errors = []

        for rule in rules:
            try:
                selector = f"{base_selector} {rule.selector}" if base_selector else rule.selector
                locator = self.page.locator(selector)

                if rule.multiple:
                    # 多值提取 - 批量优化
                    if rule.attribute:
                        values = await locator.evaluate_all(
                            f"elements => elements.map(el => el.getAttribute('{rule.attribute}')).filter(v => v)"
                        )
                    else:
                        values = await locator.evaluate_all(
                            "elements => elements.map(el => el.textContent?.trim() || '').filter(t => t)"
                        )

                    # 应用转换
                    if rule.transform and rule.transform in self._transforms:
                        values = [self._transforms[rule.transform](v) for v in values]

                    data[rule.name] = values
                else:
                    # 单值提取
                    if rule.attribute:
                        value = await locator.get_attribute(rule.attribute)
                    else:
                        value = await locator.text_content()

                    if value:
                        value = value.strip()
                        if rule.transform and rule.transform in self._transforms:
                            value = self._transforms[rule.transform](value)
                        data[rule.name] = value
                    elif rule.default is not None:
                        data[rule.name] = rule.default
                    elif rule.required:
                        errors.append(f"必填字段缺失: {rule.name}")
                        data[rule.name] = None

            except Exception as e:
                if rule.required:
                    errors.append(f"提取失败: {rule.name}, error={e}")
                data[rule.name] = rule.default

        return ExtractionResult(
            success=len(errors) == 0,
            data=data,
            errors=errors,
        )

    async def extract_json_ld(self) -> List[Dict[str, Any]]:
        """提取页面中的JSON-LD结构化数据"""
        try:
            json_lds = await self.page.evaluate("""
                () => {
                    const results = [];
                    document.querySelectorAll('script[type="application/ld+json"]').forEach(script => {
                        try {
                            results.push(JSON.parse(script.textContent));
                        } catch (e) {}
                    });
                    return results;
                }
            """)
            return json_lds
        except:
            return []

    async def extract_meta_tags(self) -> Dict[str, str]:
        """提取页面meta标签"""
        try:
            meta = await self.page.evaluate("""
                () => {
                    const result = {};
                    document.querySelectorAll('meta').forEach(meta => {
                        const name = meta.getAttribute('name') || meta.getAttribute('property');
                        const content = meta.getAttribute('content');
                        if (name && content) {
                            result[name] = content;
                        }
                    });
                    return result;
                }
            """)
            return meta
        except:
            return {}

    async def extract_open_graph(self) -> Dict[str, str]:
        """提取Open Graph标签"""
        meta = await self.extract_meta_tags()
        return {
            k.replace('og:', ''): v
            for k, v in meta.items()
            if k.startswith('og:')
        }

    async def extract_twitter_cards(self) -> Dict[str, str]:
        """提取Twitter Card标签"""
        meta = await self.extract_meta_tags()
        return {
            k.replace('twitter:', ''): v
            for k, v in meta.items()
            if k.startswith('twitter:')
        }
