src/modules/web_search.py

"""
霜云(Shimokumo) - 联网搜索模块

支持多搜索引擎（Google, Bing, Baidu, DuckDuckGo），
提供搜索结果解析、排序、关键词提取和结果缓存功能。
"""

import hashlib
import json
import os
import re
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from urllib.parse import quote_plus, urljoin

from utils.logger import get_logger

logger = get_logger("Shimokumo.WebSearch")


@dataclass
class SearchResult:
    """搜索结果数据类"""
    title: str
    """结果标题"""
    url: str
    """结果URL"""
    snippet: str
    """结果摘要"""
    source: str
    """搜索引擎来源"""
    score: float = 0.0
    """相关性评分"""
    timestamp: float = field(default_factory=time.time)
    """获取时间戳"""
    rank: int = 0
    """原始排名"""

    def to_dict(self) -> Dict[str, any]:
        """转为字典格式"""
        return {
            "title": self.title,
            "url": self.url,
            "snippet": self.snippet,
            "source": self.source,
            "score": self.score,
            "timestamp": self.timestamp,
            "rank": self.rank,
        }


class SearchCache:
    """搜索结果缓存管理器"""

    def __init__(self, cache_dir: str = "/data/user/work/shimokumo_source/cache/search", ttl: int = 3600):
        """
        初始化缓存。

        Args:
            cache_dir: 缓存目录路径
            ttl: 缓存有效期（秒）
        """
        self.cache_dir = cache_dir
        self.ttl = ttl
        os.makedirs(cache_dir, exist_ok=True)
        self._memory_cache: Dict[str, Tuple[float, List[SearchResult]]] = {}

    def _get_cache_key(self, query: str, engine: str) -> str:
        """生成缓存键"""
        raw = f"{engine}:{query}".lower().strip()
        return hashlib.md5(raw.encode("utf-8")).hexdigest()

    def get(self, query: str, engine: str) -> Optional[List[SearchResult]]:
        """
        从缓存获取搜索结果。

        Args:
            query: 搜索查询
            engine: 搜索引擎名称

        Returns:
            缓存的搜索结果列表，过期返回None
        """
        key = self._get_cache_key(query, engine)

        # 先查内存缓存
        if key in self._memory_cache:
            cached_time, results = self._memory_cache[key]
            if time.time() - cached_time < self.ttl:
                return results
            else:
                del self._memory_cache[key]

        # 再查文件缓存
        cache_file = os.path.join(self.cache_dir, f"{key}.json")
        if os.path.exists(cache_file):
            try:
                with open(cache_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                cached_time = data.get("timestamp", 0)
                if time.time() - cached_time < self.ttl:
                    results = [
                        SearchResult(**item) for item in data.get("results", [])
                    ]
                    self._memory_cache[key] = (cached_time, results)
                    return results
                else:
                    os.remove(cache_file)
            except (json.JSONDecodeError, IOError):
                pass

        return None

    def put(self, query: str, engine: str, results: List[SearchResult]) -> None:
        """
        将搜索结果写入缓存。

        Args:
            query: 搜索查询
            engine: 搜索引擎名称
            results: 搜索结果列表
        """
        key = self._get_cache_key(query, engine)
        now = time.time()

        # 写入内存缓存
        self._memory_cache[key] = (now, results)

        # 写入文件缓存
        cache_file = os.path.join(self.cache_dir, f"{key}.json")
        try:
            data = {
                "query": query,
                "engine": engine,
                "timestamp": now,
                "results": [r.to_dict() for r in results],
            }
            with open(cache_file, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
        except IOError as e:
            logger.warning(f"写入搜索缓存失败: {e}")


class WebSearchModule:
    """联网搜索模块

    支持多搜索引擎聚合搜索，提供结果排序和关键词提取。

    用法：
        search = WebSearchModule()
        results = search.search("人工智能最新进展", engines=["google", "bing"])
        for r in results:
            print(f"[{r.source}] {r.title} - {r.url}")
    """

    # 搜索引擎配置
    ENGINE_CONFIGS = {
        "google": {
            "name": "Google",
            "search_url": "https://www.google.com/search?q={query}&num={num}&hl=zh-CN",
            "user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        },
        "bing": {
            "name": "Bing",
            "search_url": "https://www.bing.com/search?q={query}&count={num}&setlang=zh-Hans",
            "user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        },
        "baidu": {
            "name": "百度",
            "search_url": "https://www.baidu.com/s?wd={query}&rn={num}",
            "user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        },
        "duckduckgo": {
            "name": "DuckDuckGo",
            "search_url": "https://duckduckgo.com/html/?q={query}",
            "user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        },
    }

    def __init__(
        self,
        default_engines: Optional[List[str]] = None,
        max_results_per_engine: int = 10,
        cache_ttl: int = 3600,
        timeout: int = 10,
    ):
        """
        初始化搜索模块。

        Args:
            default_engines: 默认搜索引擎列表
            max_results_per_engine: 每个引擎最大结果数
            cache_ttl: 缓存有效期（秒）
            timeout: 请求超时时间（秒）
        """
        self.default_engines = default_engines or ["google", "bing", "baidu", "duckduckgo"]
        self.max_results = max_results_per_engine
        self.timeout = timeout
        self.cache = SearchCache(ttl=cache_ttl)

    def extract_keywords(self, query: str) -> List[str]:
        """
        从查询文本中提取关键词。

        Args:
            query: 查询文本

        Returns:
            关键词列表（按重要性排序）
        """
        # 移除常见停用词
        stop_words = {
            "的", "了", "是", "在", "和", "有", "不", "这", "那", "我",
            "你", "他", "她", "它", "们", "一", "个", "上", "下", "中",
            "大", "小", "来", "去", "说", "会", "能", "就", "对", "被",
            "把", "让", "给", "从", "用", "过", "也", "很", "最", "都",
            "what", "is", "the", "a", "an", "of", "in", "to", "for",
            "and", "or", "but", "with", "on", "at", "by", "how", "why",
            "搜索", "查找", "帮我", "请问", "想知道", "告诉我", "一下",
        }

        # 按空格和标点分割
        words = re.split(r"[\s,，。.!！?？;；:：、\"\"''（）()\[\]【】]+", query)

        keywords: List[str] = []
        for word in words:
            word = word.strip().lower()
            if len(word) >= 2 and word not in stop_words:
                keywords.append(word)

        # 去重并按长度排序（长的关键词通常更具体）
        seen: set = set()
        unique_keywords: List[str] = []
        for kw in sorted(keywords, key=len, reverse=True):
            if kw not in seen:
                seen.add(kw)
                unique_keywords.append(kw)

        return unique_keywords

    def _make_request(self, url: str, headers: Dict[str, str]) -> Optional[str]:
        """
        发送HTTP GET请求。

        Args:
            url: 请求URL
            headers: 请求头

        Returns:
            响应HTML文本，失败返回None
        """
        try:
            import requests
            resp = requests.get(
                url,
                headers=headers,
                timeout=self.timeout,
                allow_redirects=True,
            )
            resp.raise_for_status()
            resp.encoding = resp.apparent_encoding
            return resp.text
        except Exception as e:
            logger.warning(f"搜索请求失败 [{url[:50]}...]: {e}")
            return None

    def _parse_google(self, html: str) -> List[SearchResult]:
        """解析Google搜索结果"""
        results: List[SearchResult] = []
        try:
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(html, "html.parser")

            # Google搜索结果在div.g标签中
            for item in soup.select("div.g"):
                title_tag = item.select_one("h3")
                link_tag = item.select_one("a")
                snippet_tag = item.select_one("div[data-sncf], span[style]")

                if title_tag and link_tag:
                    title = title_tag.get_text(strip=True)
                    url = link_tag.get("href", "")
                    # 过滤Google内部链接
                    if url.startswith("/search") or url.startswith("#"):
                        continue
                    if not url.startswith("http"):
                        url = "https://www.google.com" + url

                    snippet = ""
                    if snippet_tag:
                        snippet = snippet_tag.get_text(strip=True)

                    results.append(SearchResult(
                        title=title,
                        url=url,
                        snippet=snippet,
                        source="google",
                    ))
        except Exception as e:
            logger.warning(f"解析Google结果失败: {e}")

        return results[:self.max_results]

    def _parse_bing(self, html: str) -> List[SearchResult]:
        """解析Bing搜索结果"""
        results: List[SearchResult] = []
        try:
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(html, "html.parser")

            for item in soup.select("li.b_algo"):
                title_tag = item.select_one("h2 a")
                snippet_tag = item.select_one("p, div.b_caption p")

                if title_tag:
                    title = title_tag.get_text(strip=True)
                    url = title_tag.get("href", "")
                    snippet = snippet_tag.get_text(strip=True) if snippet_tag else ""

                    results.append(SearchResult(
                        title=title,
                        url=url,
                        snippet=snippet,
                        source="bing",
                    ))
        except Exception as e:
            logger.warning(f"解析Bing结果失败: {e}")

        return results[:self.max_results]

    def _parse_baidu(self, html: str) -> List[SearchResult]:
        """解析百度搜索结果"""
        results: List[SearchResult] = []
        try:
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(html, "html.parser")

            for item in soup.select("div.result, div.c-container"):
                title_tag = item.select_one("h3 a, a[href]")
                snippet_tag = item.select_one("div.c-abstract, div.c-span9, span.content-right_8Zs80")

                if title_tag:
                    title = title_tag.get_text(strip=True)
                    url = title_tag.get("href", "")
                    snippet = ""
                    if snippet_tag:
                        snippet = snippet_tag.get_text(strip=True)

                    results.append(SearchResult(
                        title=title,
                        url=url,
                        snippet=snippet,
                        source="baidu",
                    ))
        except Exception as e:
            logger.warning(f"解析百度结果失败: {e}")

        return results[:self.max_results]

    def _parse_duckduckgo(self, html: str) -> List[SearchResult]:
        """解析DuckDuckGo搜索结果"""
        results: List[SearchResult] = []
        try:
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(html, "html.parser")

            for item in soup.select("div.result, div.web-result"):
                title_tag = item.select_one("a.result__a, h2 a")
                snippet_tag = item.select_one("a.result__snippet, p.result__snippet")

                if title_tag:
                    title = title_tag.get_text(strip=True)
                    url = title_tag.get("href", "")
                    snippet = snippet_tag.get_text(strip=True) if snippet_tag else ""

                    # DuckDuckGo使用跳转链接，需要提取实际URL
                    if "uddg=" in url:
                        import urllib.parse
                        parsed = urllib.parse.parse_qs(urllib.parse.urlparse(url).query)
                        if "uddg" in parsed:
                            url = parsed["uddg"][0]

                    results.append(SearchResult(
                        title=title,
                        url=url,
                        snippet=snippet,
                        source="duckduckgo",
                    ))
        except Exception as e:
            logger.warning(f"解析DuckDuckGo结果失败: {e}")

        return results[:self.max_results]

    def _search_engine(self, engine: str, query: str) -> List[SearchResult]:
        """
        使用指定搜索引擎进行搜索。

        Args:
            engine: 搜索引擎名称
            query: 搜索查询

        Returns:
            搜索结果列表
        """
        # 检查缓存
        cached = self.cache.get(query, engine)
        if cached is not None:
            logger.debug(f"命中搜索缓存: {engine} - {query[:20]}...")
            return cached

        config = self.ENGINE_CONFIGS.get(engine)
        if not config:
            logger.warning(f"未知搜索引擎: {engine}")
            return []

        # 构建搜索URL
        encoded_query = quote_plus(query)
        url = config["search_url"].format(query=encoded_query, num=self.max_results)

        # 发送请求
        headers = {
            "User-Agent": config["user_agent"],
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
        }

        html = self._make_request(url, headers)
        if not html:
            return []

        # 解析结果
        parsers = {
            "google": self._parse_google,
            "bing": self._parse_bing,
            "baidu": self._parse_baidu,
            "duckduckgo": self._parse_duckduckgo,
        }

        parser = parsers.get(engine)
        if parser:
            results = parser(html)
        else:
            results = []

        # 设置排名并缓存
        for i, r in enumerate(results):
            r.rank = i + 1

        if results:
            self.cache.put(query, engine, results)
            logger.info(f"[{config['name']}] 搜索 '{query[:20]}...' 获取到 {len(results)} 条结果")

        return results

    def search(
        self,
        query: str,
        engines: Optional[List[str]] = None,
        max_total_results: int = 20,
    ) -> List[SearchResult]:
        """
        执行聚合搜索。

        Args:
            query: 搜索查询文本
            engines: 使用的搜索引擎列表，为None则使用默认列表
            max_total_results: 最大总结果数

        Returns:
            排序后的搜索结果列表
        """
        engines = engines or self.default_engines
        keywords = self.extract_keywords(query)

        logger.info(f"开始搜索: '{query}' | 关键词: {keywords} | 引擎: {engines}")

        all_results: List[SearchResult] = []

        # 逐引擎搜索
        for engine in engines:
            try:
                results = self._search_engine(engine, query)
                all_results.extend(results)
            except Exception as e:
                logger.warning(f"搜索引擎 {engine} 搜索失败: {e}")

        # 去重（基于URL）
        seen_urls: set = set()
        unique_results: List[SearchResult] = []
        for r in all_results:
            url_key = r.url.rstrip("/")
            if url_key not in seen_urls:
                seen_urls.add(url_key)
                unique_results.append(r)

        # 计算相关性评分并排序
        scored_results = self._score_results(unique_results, keywords)

        # 限制结果数量
        return scored_results[:max_total_results]

    def _score_results(self, results: List[SearchResult], keywords: List[str]) -> List[SearchResult]:
        """
        计算搜索结果的相关性评分。

        评分规则：
        - 标题包含完整关键词 +3分
        - 摘要包含完整关键词 +2分
        - 标题包含部分关键词 +1分
        - 摘要包含部分关键词 +0.5分
        - 原始排名分值（排名越前分值越高）

        Args:
            results: 搜索结果列表
            keywords: 关键词列表

        Returns:
            按评分排序的结果列表
        """
        for result in results:
            score = 0.0
            query_lower = "".join(keywords).lower()
            title_lower = result.title.lower()
            snippet_lower = result.snippet.lower()

            # 关键词匹配评分
            for kw in keywords:
                kw_lower = kw.lower()
                if kw_lower in title_lower:
                    score += 3.0
                if kw_lower in snippet_lower:
                    score += 2.0

            # 完整查询匹配加分
            if query_lower in title_lower:
                score += 5.0
            if query_lower in snippet_lower:
                score += 3.0

            # 排名分值
            if result.rank > 0:
                score += max(0, (10 - result.rank) * 0.5)

            result.score = score

        # 按评分降序排列
        results.sort(key=lambda x: x.score, reverse=True)

        # 重新分配排名
        for i, r in enumerate(results):
            r.rank = i + 1

        return results

    def format_results(self, results: List[SearchResult], max_items: int = 10) -> str:
        """
        将搜索结果格式化为文本。

        Args:
            results: 搜索结果列表
            max_items: 最大显示条数

        Returns:
            格式化的结果文本
        """
        if not results:
            return "没有找到相关的搜索结果的说..."

        parts: List[str] = [f"霜云帮你找到了 {len(results[:max_items])} 条搜索结果：\n"]

        for i, r in enumerate(results[:max_items], 1):
            source_tag = f"[{r.source}]" if r.source else ""
            parts.append(f"{i}. {r.title} {source_tag}")
            parts.append(f"   链接: {r.url}")
            if r.snippet:
                parts.append(f"   摘要: {r.snippet}")
            parts.append("")

        return "\n".join(parts)