plugins/content/core.py

"""
内容提取核心逻辑

集成 run/log 服务，记录处理过程日志。
"""

import logging
import tempfile
from pathlib import Path
from typing import Optional, Dict, Any

from .extractors.xiaohongshu import XiaohongshuExtractor
from .extractors.weibo import WeiboExtractor
from plugins.ocr.core import get_ocr_service
from app.plugins.run_log import get_run_log_service, RunStatus

logger = logging.getLogger(__name__)


class ContentExtractorCore:
    """内容提取核心逻辑"""

    def __init__(self):
        self.xhs_extractor = XiaohongshuExtractor()
        self.weibo_extractor = WeiboExtractor()
        self.ocr_service = get_ocr_service()

    def validate_url(self, url: str) -> dict:
        """
        验证URL

        Args:
            url: 待验证的URL

        Returns:
            {
                "valid": bool,
                "source_type": str,
                "error": str  # 错误信息（如果无效）
            }
        """
        if not url or not url.strip():
            return {
                "valid": False,
                "source_type": "unknown",
                "error": "链接不能为空"
            }

        # 按平台逐个解析，避免把其他平台链接误判成小红书错误。
        for extractor in (self.xhs_extractor, self.weibo_extractor):
            result = extractor.parse_url(url)
            if result["valid"]:
                return {
                    "valid": True,
                    "source_type": result["source_type"],
                    "error": None
                }

        return {
            "valid": False,
            "source_type": "unknown",
            "error": "无效或暂不支持的链接"
        }

    async def extract_with_run(
        self,
        url: str,
        run_id: str,
        include_ocr: bool = True,
        cookies: Optional[dict] = None
    ) -> dict:
        """
        提取内容，记录 run 事件

        Args:
            url: 内容链接
            run_id: 运行 ID
            include_ocr: 是否对图片进行OCR识别
            cookies: Cookie（可选，用于应对反爬）

        Returns:
            {
                "success": bool,
                "title": str,
                "content": str,
                "raw_html": str,
                "raw_text": str,
                "normalized_content": str,
                "images": list,
                "images_text": str,
                "source_type": str,
                "error": str
            }
        """
        run_service = get_run_log_service()

        # 1. 验证URL
        run_service.add_event(
            run_id=run_id,
            stage="validate_url",
            message=f"验证 URL: {url}",
        )

        validation = self.validate_url(url)
        if not validation["valid"]:
            run_service.add_event(
                run_id=run_id,
                stage="validate_url",
                message=f"URL 验证失败: {validation['error']}",
                level="error",
            )
            return {
                "success": False,
                "title": "",
                "content": "",
                "raw_html": "",
                "raw_text": "",
                "normalized_content": "",
                "images": [],
                "images_text": "",
                "source_type": validation["source_type"],
                "error": validation["error"]
            }

        source_type = validation["source_type"]
        run_service.add_event(
            run_id=run_id,
            stage="validate_url",
            message=f"URL 验证通过: {source_type}",
        )

        # 2. 根据来源类型选择提取器
        if source_type == "xiaohongshu":
            return await self._extract_xiaohongshu_with_run(url, run_id, include_ocr, cookies)
        if source_type == "weibo":
            return await self._extract_weibo_with_run(url, run_id, include_ocr, cookies)
        else:
            error_msg = f"不支持的内容来源: {source_type}"
            run_service.add_event(
                run_id=run_id,
                stage="extract",
                message=error_msg,
                level="error",
            )
            return {
                "success": False,
                "title": "",
                "content": "",
                "raw_html": "",
                "raw_text": "",
                "normalized_content": "",
                "images": [],
                "images_text": "",
                "source_type": source_type,
                "error": error_msg
            }

    async def _extract_xiaohongshu_with_run(
        self,
        url: str,
        run_id: str,
        include_ocr: bool,
        cookies: Optional[dict] = None
    ) -> dict:
        """提取小红书内容，记录 run 事件"""
        run_service = get_run_log_service()

        try:
            # 1. 获取页面内容
            run_service.add_event(
                run_id=run_id,
                stage="fetch_page",
                message="获取页面内容",
            )

            page_result = await self.xhs_extractor.fetch_page(url, cookies)
            if not page_result["success"]:
                run_service.add_event(
                    run_id=run_id,
                    stage="fetch_page",
                    message=f"获取页面失败: {page_result['error']}",
                    level="error",
                )
                return {
                    "success": False,
                    "title": "",
                    "content": "",
                    "raw_html": "",
                    "raw_text": "",
                    "normalized_content": "",
                    "images": [],
                    "images_text": "",
                    "source_type": "xiaohongshu",
                    "error": page_result["error"]
                }

            raw_html = page_result.get("html", "")
            run_service.add_event(
                run_id=run_id,
                stage="fetch_page",
                message=f"获取页面成功: {len(raw_html)} 字符",
            )

            # 2. 解析页面内容
            run_service.add_event(
                run_id=run_id,
                stage="parse_content",
                message="解析页面内容",
            )

            content_result = self.xhs_extractor.parse_content(raw_html)
            if not content_result["success"]:
                run_service.add_event(
                    run_id=run_id,
                    stage="parse_content",
                    message=f"解析内容失败: {content_result['error']}",
                    level="error",
                )
                return {
                    "success": False,
                    "title": "",
                    "content": "",
                    "raw_html": raw_html,
                    "raw_text": "",
                    "normalized_content": "",
                    "images": [],
                    "images_text": "",
                    "source_type": "xiaohongshu",
                    "error": content_result["error"]
                }

            images = content_result.get("images", [])
            run_service.add_event(
                run_id=run_id,
                stage="parse_content",
                message=f"解析内容成功: {len(images)} 张图片",
            )

            # 3. OCR识别图片（如果需要）
            images_text = ""
            if include_ocr and images:
                run_service.add_event(
                    run_id=run_id,
                    stage="ocr",
                    message=f"开始 OCR 识别: {len(images)} 张图片",
                )
                images_text = await self._ocr_images_with_run(
                    images,
                    self.xhs_extractor.download_image,
                    run_id,
                )
                run_service.add_event(
                    run_id=run_id,
                    stage="ocr",
                    message=f"OCR 识别完成: {len(images_text)} 字符",
                )

            # 4. 格式化结果
            run_service.add_event(
                run_id=run_id,
                stage="format",
                message="格式化结果",
            )

            return {
                "success": True,
                "title": content_result.get("title", ""),
                "content": content_result.get("content", ""),
                "raw_html": raw_html,
                "raw_text": content_result.get("raw_text", ""),
                "normalized_content": content_result.get("content", ""),
                "images": images,
                "images_text": images_text,
                "source_type": "xiaohongshu",
                "error": None
            }

        except Exception as e:
            logger.error(f"提取小红书内容失败: {e}", exc_info=True)
            run_service.add_event(
                run_id=run_id,
                stage="extract",
                message=f"提取异常: {str(e)}",
                level="error",
            )
            return {
                "success": False,
                "title": "",
                "content": "",
                "raw_html": "",
                "raw_text": "",
                "normalized_content": "",
                "images": [],
                "images_text": "",
                "source_type": "xiaohongshu",
                "error": str(e)
            }

    async def _extract_weibo_with_run(
        self,
        url: str,
        run_id: str,
        include_ocr: bool,
        cookies: Optional[dict] = None
    ) -> dict:
        """提取微博内容，记录 run 事件"""
        run_service = get_run_log_service()

        try:
            # 1. 获取页面内容
            run_service.add_event(
                run_id=run_id,
                stage="fetch_page",
                message="获取页面内容",
            )

            page_result = await self.weibo_extractor.fetch_page(url, cookies)
            if not page_result["success"]:
                run_service.add_event(
                    run_id=run_id,
                    stage="fetch_page",
                    message=f"获取页面失败: {page_result['error']}",
                    level="error",
                )
                return {
                    "success": False,
                    "title": "",
                    "content": "",
                    "raw_html": "",
                    "raw_text": "",
                    "normalized_content": "",
                    "images": [],
                    "images_text": "",
                    "source_type": "weibo",
                    "error": page_result["error"]
                }

            raw_html = page_result.get("html", "")
            run_service.add_event(
                run_id=run_id,
                stage="fetch_page",
                message=f"获取页面成功: {len(raw_html)} 字符",
            )

            # 2. 解析页面内容
            run_service.add_event(
                run_id=run_id,
                stage="parse_content",
                message="解析页面内容",
            )

            content_result = self.weibo_extractor.parse_content(raw_html)
            if not content_result["success"]:
                run_service.add_event(
                    run_id=run_id,
                    stage="parse_content",
                    message=f"解析内容失败: {content_result['error']}",
                    level="error",
                )
                return {
                    "success": False,
                    "title": "",
                    "content": "",
                    "raw_html": raw_html,
                    "raw_text": "",
                    "normalized_content": "",
                    "images": [],
                    "images_text": "",
                    "source_type": "weibo",
                    "error": content_result["error"]
                }

            images = content_result.get("images", [])
            run_service.add_event(
                run_id=run_id,
                stage="parse_content",
                message=f"解析内容成功: {len(images)} 张图片",
            )

            # 3. OCR识别图片（如果需要）
            images_text = ""
            if include_ocr and images:
                run_service.add_event(
                    run_id=run_id,
                    stage="ocr",
                    message=f"开始 OCR 识别: {len(images)} 张图片",
                )
                images_text = await self._ocr_images_with_run(
                    images,
                    self.weibo_extractor.download_image,
                    run_id,
                )
                run_service.add_event(
                    run_id=run_id,
                    stage="ocr",
                    message=f"OCR 识别完成: {len(images_text)} 字符",
                )

            # 4. 格式化结果
            run_service.add_event(
                run_id=run_id,
                stage="format",
                message="格式化结果",
            )

            return {
                "success": True,
                "title": content_result.get("title", ""),
                "content": content_result.get("content", ""),
                "raw_html": raw_html,
                "raw_text": content_result.get("raw_text", ""),
                "normalized_content": content_result.get("content", ""),
                "images": images,
                "images_text": images_text,
                "source_type": "weibo",
                "error": None
            }

        except Exception as e:
            logger.error(f"提取微博内容失败: {e}", exc_info=True)
            run_service.add_event(
                run_id=run_id,
                stage="extract",
                message=f"提取异常: {str(e)}",
                level="error",
            )
            return {
                "success": False,
                "title": "",
                "content": "",
                "raw_html": "",
                "raw_text": "",
                "normalized_content": "",
                "images": [],
                "images_text": "",
                "source_type": "weibo",
                "error": str(e)
            }

    async def _ocr_images_with_run(self, image_urls: list, download_image, run_id: str) -> str:
        """对图片进行OCR识别，记录 run 事件"""
        run_service = get_run_log_service()

        if not image_urls:
            run_service.add_event(
                run_id=run_id,
                stage="ocr",
                message="跳过 OCR：页面未解析到可下载图片",
            )
            return ""

        all_text = []
        for idx, image_url in enumerate(image_urls[:5]):  # 限制最多5张图片
            try:
                # 下载图片
                run_service.add_event(
                    run_id=run_id,
                    stage="download_image",
                    message=f"下载图片 {idx + 1}: {image_url[:50]}...",
                )

                download_result = await download_image(image_url)
                if not download_result["success"]:
                    run_service.add_event(
                        run_id=run_id,
                        stage="download_image",
                        message=f"下载图片失败: {download_result['error']}",
                        level="warning",
                    )
                    continue

                # 保存到临时文件
                with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
                    tmp_file.write(download_result["data"])
                    tmp_path = tmp_file.name

                try:
                    # 执行OCR识别
                    run_service.add_event(
                        run_id=run_id,
                        stage="ocr",
                        message=f"OCR 识别图片 {idx + 1}",
                    )

                    ocr_result = self.ocr_service.extract_text_from_file(tmp_path)
                    if ocr_result["success"] and ocr_result["text"]:
                        all_text.append(f"[图片{idx + 1}]\n{ocr_result['text']}")
                        run_service.add_event(
                            run_id=run_id,
                            stage="ocr",
                            message=f"OCR 识别成功: {len(ocr_result['text'])} 字符",
                        )
                    elif ocr_result["success"]:
                        run_service.add_event(
                            run_id=run_id,
                            stage="ocr",
                            message=f"OCR 未识别到文字: {image_url[:50]}...",
                        )
                    else:
                        run_service.add_event(
                            run_id=run_id,
                            stage="ocr",
                            message=f"OCR 识别失败: {ocr_result.get('error')}",
                            level="warning",
                        )
                finally:
                    # 清理临时文件
                    Path(tmp_path).unlink(missing_ok=True)

            except Exception as e:
                logger.error(f"OCR识别图片失败: {e}")
                run_service.add_event(
                    run_id=run_id,
                    stage="ocr",
                    message=f"OCR 识别异常: {str(e)}",
                    level="error",
                )
                continue

        return "\n\n".join(all_text)

    async def extract(
        self,
        url: str,
        include_ocr: bool = True,
        cookies: Optional[dict] = None
    ) -> dict:
        """
        提取内容（兼容旧接口）

        Args:
            url: 内容链接
            include_ocr: 是否对图片进行OCR识别
            cookies: Cookie（可选，用于应对反爬）

        Returns:
            {
                "success": bool,
                "title": str,
                "content": str,
                "images_text": str,  # 图片OCR识别的文本
                "source_type": str,
                "error": str  # 错误信息（如果失败）
            }
        """
        # 1. 验证URL
        validation = self.validate_url(url)
        if not validation["valid"]:
            return {
                "success": False,
                "title": "",
                "content": "",
                "images_text": "",
                "source_type": validation["source_type"],
                "error": validation["error"]
            }

        source_type = validation["source_type"]

        # 2. 根据来源类型选择提取器
        if source_type == "xiaohongshu":
            return await self._extract_xiaohongshu(url, include_ocr, cookies)
        if source_type == "weibo":
            return await self._extract_weibo(url, include_ocr, cookies)
        else:
            return {
                "success": False,
                "title": "",
                "content": "",
                "images_text": "",
                "source_type": source_type,
                "error": f"不支持的内容来源: {source_type}"
            }

    async def _extract_xiaohongshu(
        self,
        url: str,
        include_ocr: bool,
        cookies: Optional[dict] = None
    ) -> dict:
        """提取小红书内容"""
        try:
            # 1. 获取页面内容
            page_result = await self.xhs_extractor.fetch_page(url, cookies)
            if not page_result["success"]:
                return {
                    "success": False,
                    "title": "",
                    "content": "",
                    "images_text": "",
                    "source_type": "xiaohongshu",
                    "error": page_result["error"]
                }

            # 2. 解析页面内容
            content_result = self.xhs_extractor.parse_content(page_result["html"])
            if not content_result["success"]:
                return {
                    "success": False,
                    "title": "",
                    "content": "",
                    "images_text": "",
                    "source_type": "xiaohongshu",
                    "error": content_result["error"]
                }

            # 3. OCR识别图片（如果需要）
            images_text = ""
            if include_ocr and content_result["images"]:
                images_text = await self._ocr_images(
                    content_result["images"],
                    self.xhs_extractor.download_image,
                )

            return {
                "success": True,
                "title": content_result["title"],
                "content": content_result["content"],
                "images_text": images_text,
                "source_type": "xiaohongshu",
                "error": None
            }

        except Exception as e:
            logger.error(f"提取小红书内容失败: {e}", exc_info=True)
            return {
                "success": False,
                "title": "",
                "content": "",
                "images_text": "",
                "source_type": "xiaohongshu",
                "error": str(e)
            }

    async def _extract_weibo(
        self,
        url: str,
        include_ocr: bool,
        cookies: Optional[dict] = None
    ) -> dict:
        """提取微博内容。"""
        try:
            page_result = await self.weibo_extractor.fetch_page(url, cookies)
            if not page_result["success"]:
                return {
                    "success": False,
                    "title": "",
                    "content": "",
                    "images_text": "",
                    "source_type": "weibo",
                    "error": page_result["error"]
                }

            content_result = self.weibo_extractor.parse_content(page_result["html"])
            if not content_result["success"]:
                return {
                    "success": False,
                    "title": "",
                    "content": "",
                    "images_text": "",
                    "source_type": "weibo",
                    "error": content_result["error"]
                }

            images_text = ""
            if include_ocr and content_result["images"]:
                images_text = await self._ocr_images(
                    content_result["images"],
                    self.weibo_extractor.download_image,
                )

            return {
                "success": True,
                "title": content_result["title"],
                "content": content_result["content"],
                "images_text": images_text,
                "source_type": "weibo",
                "error": None
            }

        except Exception as e:
            logger.error(f"提取微博内容失败: {e}", exc_info=True)
            return {
                "success": False,
                "title": "",
                "content": "",
                "images_text": "",
                "source_type": "weibo",
                "error": str(e)
            }

    async def _ocr_images(self, image_urls: list, download_image) -> str:
        """对图片进行OCR识别"""
        if not image_urls:
            logger.info("跳过 OCR：页面未解析到可下载图片")
            return ""

        all_text = []
        for idx, image_url in enumerate(image_urls[:5]):  # 限制最多5张图片
            try:
                # 下载图片
                download_result = await download_image(image_url)
                if not download_result["success"]:
                    logger.warning(f"下载图片失败: {download_result['error']}")
                    continue

                # 保存到临时文件
                with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
                    tmp_file.write(download_result["data"])
                    tmp_path = tmp_file.name

                try:
                    # 执行OCR识别
                    ocr_result = self.ocr_service.extract_text_from_file(tmp_path)
                    if ocr_result["success"] and ocr_result["text"]:
                        all_text.append(f"[图片{idx + 1}]\n{ocr_result['text']}")
                    elif ocr_result["success"]:
                        logger.info(f"OCR 未识别到文字: {image_url}")
                    else:
                        logger.warning(f"OCR 识别失败: {ocr_result.get('error')}")
                finally:
                    # 清理临时文件
                    Path(tmp_path).unlink(missing_ok=True)

            except Exception as e:
                logger.error(f"OCR识别图片失败: {e}")
                continue

        return "\n\n".join(all_text)