"""
基础数据提取器
从网页中提取结构化数据
"""
import re
import json
import asyncio
from typing import Optional, List, Dict, Any, Union, Callable
from playwright.async_api import Page, Locator, Response
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: Union[str, Locator] = "table",
        header_row: int = 0,
        nth: Optional[int] = None,
    ) -> List[Dict[str, str]]:
        """
        提取表格数据

        Args:
            selector: 表格选择器字符串或 Locator 对象
            header_row: 表头行索引
            nth: 表格索引（当有多个匹配时选择第 N 个，从 0 开始）

        Returns:
            List[Dict]: 表格数据列表

        Examples:
            # 提取第一个表格
            data = await extractor.extract_table("table.data")

            # 提取第三个表格（索引从 0 开始）
            data = await extractor.extract_table("table", nth=2)

            # 使用 Locator
            locator = page.locator("table").nth(2)
            data = await extractor.extract_table(locator)
        """
        try:
            # 处理 Locator 对象
            if isinstance(selector, Locator):
                return await self._extract_table_from_locator(selector, header_row)

            # 处理字符串选择器
            return await self._extract_table_from_selector(selector, header_row, nth)
        except Exception as e:
            logger.error(f"提取表格失败: {e}")
            return []

    async def _extract_table_from_locator(
        self,
        locator: Locator,
        header_row: int = 0,
    ) -> List[Dict[str, str]]:
        """从 Locator 对象提取表格数据"""
        table_data = await locator.evaluate("""
            (table) => {
                const rows = table.querySelectorAll('tr');
                if (rows.length === 0) return [];

                const headers = [];
                const data = [];

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

                // 获取数据行
                for (let i = %d; 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;
            }
        """ % (header_row, header_row + 1))
        return table_data

    async def _extract_table_from_selector(
        self,
        selector: str,
        header_row: int = 0,
        nth: Optional[int] = None,
    ) -> List[Dict[str, str]]:
        """从选择器字符串提取表格数据"""
        # 构建选择表达式
        if nth is not None:
            select_expr = f"document.querySelectorAll('{selector}')[{nth}]"
        else:
            select_expr = f"document.querySelector('{selector}')"

        table_data = await self.page.evaluate(f"""
            () => {{
                const table = {select_expr};
                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

    async def extract_split_table(
        self,
        header_selector: str,
        body_selector: str,
        header_row: int = 0,
    ) -> List[Dict[str, str]]:
        """
        提取分离表格数据（表头和内容在不同的 table 中）

        适用于表头固定、内容可滚动的场景：
        <table class="header-table">  <!-- 表头 -->
            <tr><th>列1</th><th>列2</th></tr>
        </table>
        <table class="body-table">    <!-- 内容 -->
            <tr><td>数据1</td><td>数据2</td></tr>
        </table>

        Args:
            header_selector: 表头表格选择器
            body_selector: 内容表格选择器
            header_row: 表头行索引（默认第0行）

        Returns:
            List[Dict]: 表格数据列表
        """
        try:
            table_data = await self.page.evaluate(f"""
                () => {{
                    const headerTable = document.querySelector('{header_selector}');
                    const bodyTable = document.querySelector('{body_selector}');

                    if (!headerTable || !bodyTable) return [];

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

                    // 获取数据行
                    const bodyRows = bodyTable.querySelectorAll('tr');
                    const data = [];

                    for (let i = 0; i < bodyRows.length; i++) {{
                        const cells = bodyRows[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;
                }}
            """)
            logger.debug(f"提取分离表格: headers={len(table_data[0]) if table_data else 0}, rows={len(table_data)}")
            return table_data
        except Exception as e:
            logger.error(f"提取分离表格失败: {e}")
            return []

    async def extract_tables(
        self,
        selector: str = "table",
        split_mode: bool = False,
        header_selector: Optional[str] = None,
        body_selector: Optional[str] = None,
    ) -> List[List[Dict[str, str]]]:
        """
        提取页面上所有表格数据

        Args:
            selector: 表格选择器
            split_mode: 是否使用分离模式
            header_selector: 表头表格选择器（split_mode=True 时使用）
            body_selector: 内容表格选择器（split_mode=True 时使用）

        Returns:
            List[List[Dict]]: 所有表格数据列表
        """
        if split_mode and header_selector and body_selector:
            # 分离表格模式
            data = await self.extract_split_table(header_selector, body_selector)
            return [data] if data else []

        # 提取所有普通表格
        try:
            tables_data = await self.page.evaluate(f"""
                () => {{
                    const tables = document.querySelectorAll('{selector}');
                    const results = [];

                    tables.forEach(table => {{
                        const rows = table.querySelectorAll('tr');
                        if (rows.length === 0) return;

                        const headers = [];
                        const data = [];
                        let headerFound = false;

                        // 查找表头行（第一个包含 th 或 td 的行）
                        for (let i = 0; i < rows.length; i++) {{
                            const cells = rows[i].querySelectorAll('th, td');
                            if (cells.length > 0 && !headerFound) {{
                                cells.forEach(cell => headers.push(cell.textContent.trim()));
                                headerFound = true;
                            }} else if (headerFound && cells.length > 0) {{
                                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);
                                }}
                            }}
                        }})

                        if (data.length > 0) {{
                            results.push(data);
                        }}
                    }});

                    return results;
                }}
            """)
            logger.debug(f"提取所有表格: {len(tables_data)} 个表格")
            return tables_data
        except Exception as e:
            logger.error(f"提取所有表格失败: {e}")
            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:')
        }

    # ============================================================
    # 表格提取加速方案
    # ============================================================

    async def extract_table_via_html(
        self,
        selector: str = "table",
        parser: str = "auto",
        header_row: int = 0,
    ) -> List[Dict[str, str]]:
        """
        方案1: HTML解析方案 - 在Python端解析表格

        优势: 纯Python处理，避免JS引擎开销，适合大数据量
        劣势: 无法获取动态渲染后的数据

        Args:
            selector: 表格选择器
            parser: HTML解析器
                - 'auto': 自动降级 lxml → html.parser → evaluate（默认）
                - 'lxml': 最快，需要安装 lxml
                - 'html.parser': 内置，无需额外依赖
            header_row: 表头行索引

        Returns:
            List[Dict]: 表格数据列表

        性能对比 (10000行表格):
            - evaluate: ~200ms
            - html解析: ~50ms (lxml) / ~100ms (html.parser)
        """
        try:
            from bs4 import BeautifulSoup
        except ImportError:
            logger.warning("BeautifulSoup4 未安装，回退到 evaluate 方案")
            return await self.extract_table(selector, header_row=header_row)

        # 单次获取整个表格HTML
        try:
            html = await self.page.locator(selector).outer_html()
        except Exception as e:
            logger.error(f"获取表格HTML失败: {e}")
            return []

        # 解析器降级链
        parsers_to_try = self._get_parser_chain(parser)

        for current_parser in parsers_to_try:
            try:
                soup = BeautifulSoup(html, current_parser)
                table = soup.find('table')
                if not table:
                    continue

                rows = table.find_all('tr')
                if len(rows) <= header_row:
                    continue

                # 获取表头
                header_cells = rows[header_row].find_all(['th', 'td'])
                headers = [cell.get_text(strip=True) for cell in header_cells]

                # 获取数据行
                data = []
                for tr in rows[header_row + 1:]:
                    cells = tr.find_all('td')
                    if not cells:
                        continue
                    row_data = {}
                    for i, td in enumerate(cells):
                        if i < len(headers):
                            row_data[headers[i]] = td.get_text(strip=True)
                    if row_data:
                        data.append(row_data)

                if data:
                    logger.debug(f"HTML解析表格: {len(data)} 行, parser={current_parser}")
                    return data

            except Exception as e:
                logger.debug(f"解析器 {current_parser} 失败: {e}")
                continue

        # 所有解析器都失败，回退到 evaluate
        logger.debug("HTML解析无数据，回退到 evaluate 方案")
        return await self.extract_table(selector, header_row=header_row)

    def _get_parser_chain(self, parser: str) -> List[str]:
        """获取解析器尝试链"""
        if parser == "auto":
            # 自动降级：lxml → html.parser
            chain = []
            try:
                import lxml  # noqa: F401
                chain.append("lxml")
            except ImportError:
                pass
            chain.append("html.parser")
            return chain
        elif parser == "lxml":
            try:
                import lxml  # noqa: F401
                return ["lxml"]
            except ImportError:
                logger.warning("lxml 未安装，使用 html.parser")
                return ["html.parser"]
        else:
            return [parser]

    async def extract_table_from_response(
        self,
        url_pattern: str,
        table_selector: str = "table",
        parser: str = "lxml",
        timeout: float = 30.0,
    ) -> List[Dict[str, str]]:
        """
        方案2: 响应拦截方案 - 从网络响应中直接解析

        优势: 无需等待页面完整渲染，适合服务端渲染的表格
        劣势: 需要触发表格数据的网络请求

        Args:
            url_pattern: 匹配的URL模式（支持通配符）
            table_selector: 表格选择器
            parser: HTML解析器
            timeout: 超时时间（秒）

        Returns:
            List[Dict]: 表格数据列表

        使用示例:
            # 表格数据来自API请求
            data = await extractor.extract_table_from_response(
                url_pattern="/api/data",
                table_selector="table.result"
            )
        """
        from bs4 import BeautifulSoup
        from fnmatch import fnmatch

        results = []
        event_received = asyncio.Event()

        async def handle_response(response: Response):
            nonlocal results
            url = response.url

            # 匹配URL模式
            if fnmatch(url, f"*{url_pattern}*") or url_pattern in url:
                try:
                    if response.ok and 'text/html' in (response.headers.get('content-type', '')):
                        html = await response.text()
                        soup = BeautifulSoup(html, parser)
                        table = soup.select_one(table_selector)

                        if table:
                            rows = table.find_all('tr')
                            if len(rows) > 0:
                                headers = [cell.get_text(strip=True)
                                          for cell in rows[0].find_all(['th', 'td'])]

                                for tr in rows[1:]:
                                    cells = tr.find_all('td')
                                    if cells:
                                        row_data = {}
                                        for i, td in enumerate(cells):
                                            if i < len(headers):
                                                row_data[headers[i]] = td.get_text(strip=True)
                                        if row_data:
                                            results.append(row_data)

                                event_received.set()
                except Exception as e:
                    logger.debug(f"响应解析失败: {e}")

        # 注册监听器
        self.page.on("response", handle_response)

        try:
            # 等待目标响应
            await asyncio.wait_for(event_received.wait(), timeout=timeout)
            logger.debug(f"响应拦截表格: {len(results)} 行")
            return results
        except asyncio.TimeoutError:
            logger.warning(f"响应拦截超时: pattern={url_pattern}")
            return []
        finally:
            self.page.remove_listener("response", handle_response)

    async def extract_table_via_cdp(
        self,
        selector: str = "table",
        header_row: int = 0,
    ) -> List[Dict[str, str]]:
        """
        方案3: CDP协议方案 - 使用Chrome DevTools Protocol直接访问DOM

        优势: 绕过Playwright的抽象层，直接与浏览器通信，性能最优
        劣势: 仅支持Chromium内核浏览器，API较低层

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

        Returns:
            List[Dict]: 表格数据列表

        性能对比 (10000行表格):
            - evaluate: ~200ms
            - CDP: ~150ms (减少约25%)

        注意: 仅在Chromium浏览器上可用
        """
        try:
            # 获取CDP会话
            cdp = await self.page.context.new_cdp_session(self.page)

            # 获取文档节点
            document = await cdp.send('DOM.getDocument')

            # 查找表格元素
            result = await cdp.send('DOM.querySelector', {
                'nodeId': document['root']['nodeId'],
                'selector': selector
            })

            if 'nodeId' not in result:
                return []

            table_node_id = result['nodeId']

            # 获取表格HTML
            html_result = await cdp.send('DOM.getOuterHTML', {
                'nodeId': table_node_id
            })

            # 解析HTML
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(html_result['outerHTML'], 'lxml')
            table = soup.find('table')

            if not table:
                return []

            rows = table.find_all('tr')
            if len(rows) <= header_row:
                return []

            # 获取表头和数据
            header_cells = rows[header_row].find_all(['th', 'td'])
            headers = [cell.get_text(strip=True) for cell in header_cells]

            data = []
            for tr in rows[header_row + 1:]:
                cells = tr.find_all('td')
                if not cells:
                    continue
                row_data = {}
                for i, td in enumerate(cells):
                    if i < len(headers):
                        row_data[headers[i]] = td.get_text(strip=True)
                if row_data:
                    data.append(row_data)

            logger.debug(f"CDP表格提取: {len(data)} 行")
            return data

        except Exception as e:
            # CDP失败时回退到evaluate
            logger.debug(f"CDP不可用，回退到evaluate: {e}")
            return await self.extract_table(selector, header_row=header_row)

    async def extract_table_fast(
        self,
        selector: str = "table",
        method: str = "auto",
        header_row: int = 0,
        **kwargs,
    ) -> List[Dict[str, str]]:
        """
        智能快速表格提取 - 自动选择最优方案

        Args:
            selector: 表格选择器
            method: 提取方法
                - 'auto': 自动选择（默认，推荐）
                - 'evaluate': JS evaluate方案
                - 'html': HTML解析方案（自动降级 lxml→html.parser→evaluate）
                - 'cdp': CDP协议方案（仅Chromium）
                - 'response': 响应拦截方案（需要url_pattern参数）
            header_row: 表头行索引
            **kwargs: 传递给具体方法的额外参数
                - parser: HTML解析器 ('auto', 'lxml', 'html.parser')，method='html'时有效
                - url_pattern: URL匹配模式，method='response'时必需
                - timeout: 超时时间，method='response'时有效

        Returns:
            List[Dict]: 表格数据列表

        使用示例:
            # 自动选择最优方案（推荐）
            data = await extractor.extract_table_fast("table.data")

            # 强制使用HTML解析
            data = await extractor.extract_table_fast("table", method="html")

            # 从响应中拦截（需要触发请求）
            data = await extractor.extract_table_fast(
                "table", method="response", url_pattern="/api/data"
            )
        """
        if method == "html":
            return await self.extract_table_via_html(
                selector, header_row=header_row, **kwargs
            )
        elif method == "cdp":
            return await self.extract_table_via_cdp(
                selector, header_row=header_row
            )
        elif method == "response":
            url_pattern = kwargs.get('url_pattern', '')
            if not url_pattern:
                raise ValueError("response方法需要url_pattern参数")
            return await self.extract_table_from_response(
                url_pattern, table_selector=selector, **kwargs
            )
        elif method == "evaluate":
            return await self.extract_table(
                selector, header_row=header_row
            )
        else:  # auto
            # 自动选择逻辑：
            # extract_table_via_html 内部已实现完整降级链
            # lxml → html.parser → evaluate
            return await self.extract_table_via_html(
                selector, header_row=header_row
            )
