"""通用搜索客户端（Tavily / Exa / Brave / Serper 路由）。

为 build_glossary.py 这类术语核查场景服务。

关键设计：
- `trust_env=False` 绕开系统 socks 代理（Clash on macOS 配 socks5 时 httpx 会 TLS EOF）
- 专利 / Scholar / News 优先 Serper，保证 Google Patents / Google Scholar 路径被真正调用
- 通用网页 Tavily 优先，Exa/Brave fallback
- 证据发现 Exa 优先，用 highlights/text 摘录喂给 evidence packet
- 遇到配额问题自动降级或返回 empty
- 不做深度 crawl，只要摘要
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any

import httpx


@dataclass
class SearchHit:
    title: str
    url: str
    snippet: str


class SearchError(RuntimeError):
    pass


class ExaClient:
    def __init__(self, api_key: str | None = None, timeout: float = 30.0) -> None:
        self.api_key = api_key or os.environ.get("EXA_API_KEY")
        if not self.api_key:
            raise SearchError("EXA_API_KEY not set")
        # trust_env=False 关键：不吃系统代理，避免 TLS EOF
        self._client = httpx.Client(trust_env=False, timeout=timeout)

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "ExaClient":
        return self

    def __exit__(self, *_args: Any) -> None:
        self.close()

    def search(
        self,
        query: str,
        *,
        num_results: int = 5,
        search_type: str = "auto",
        category: str | None = None,
        use_highlights: bool = False,
        max_characters: int = 800,
    ) -> list[SearchHit]:
        body = {
            "query": query,
            "numResults": num_results,
            "type": search_type,
            "contents": {"text": {"maxCharacters": max_characters}},
        }
        if category:
            body["category"] = category
        if use_highlights:
            body["contents"]["highlights"] = {
                "numSentences": 2,
                "highlightsPerUrl": 3,
            }
        r = self._client.post(
            "https://api.exa.ai/search",
            json=body,
            headers={"x-api-key": self.api_key, "Content-Type": "application/json"},
        )
        if r.status_code != 200:
            raise SearchError(f"Exa HTTP {r.status_code}: {r.text[:200]}")
        data = r.json()
        out: list[SearchHit] = []
        for item in data.get("results", [])[:num_results]:
            highlights = item.get("highlights") or []
            text = item.get("text") or item.get("snippet") or ""
            if highlights:
                text = " | ".join(str(h).strip() for h in highlights if str(h).strip())
            out.append(
                SearchHit(
                    title=(item.get("title") or "")[:200],
                    url=item.get("url") or "",
                    snippet=text[:1000],
                )
            )
        return out


class TavilyClient:
    def __init__(self, api_key: str | None = None, timeout: float = 30.0) -> None:
        self.api_key = api_key or os.environ.get("TAVILY_API_KEY")
        if not self.api_key:
            raise SearchError("TAVILY_API_KEY not set")
        self._client = httpx.Client(trust_env=False, timeout=timeout)

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "TavilyClient":
        return self

    def __exit__(self, *_args: Any) -> None:
        self.close()

    def search(self, query: str, *, num_results: int = 5) -> list[SearchHit]:
        body = {
            "api_key": self.api_key,
            "query": query,
            "search_depth": "basic",
            "max_results": num_results,
            "include_answer": False,
            "include_raw_content": False,
        }
        r = self._client.post("https://api.tavily.com/search", json=body)
        if r.status_code != 200:
            raise SearchError(f"Tavily HTTP {r.status_code}: {r.text[:200]}")
        data = r.json()
        out: list[SearchHit] = []
        for item in data.get("results", [])[:num_results]:
            out.append(
                SearchHit(
                    title=(item.get("title") or "")[:200],
                    url=item.get("url") or "",
                    snippet=(item.get("content") or "")[:600],
                )
            )
        return out


class BraveClient:
    def __init__(self, api_key: str | None = None, timeout: float = 30.0) -> None:
        self.api_key = api_key or os.environ.get("BRAVE_API_KEY")
        if not self.api_key:
            raise SearchError("BRAVE_API_KEY not set")
        self._client = httpx.Client(trust_env=False, timeout=timeout)

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "BraveClient":
        return self

    def __exit__(self, *_args: Any) -> None:
        self.close()

    def search(self, query: str, *, num_results: int = 5) -> list[SearchHit]:
        r = self._client.get(
            "https://api.search.brave.com/res/v1/web/search",
            params={"q": query, "count": min(max(num_results, 1), 20)},
            headers={
                "X-Subscription-Token": self.api_key,
                "Accept": "application/json",
            },
        )
        if r.status_code != 200:
            raise SearchError(f"Brave HTTP {r.status_code}: {r.text[:200]}")
        data = r.json()
        out: list[SearchHit] = []
        for item in (data.get("web") or {}).get("results", [])[:num_results]:
            out.append(
                SearchHit(
                    title=(item.get("title") or "")[:200],
                    url=item.get("url") or "",
                    snippet=(item.get("description") or "")[:600],
                )
            )
        return out


class SearchClient:
    """统一搜索门面，支持多路由：

    - `search(query)`：通用网页搜索，优先 Tavily → Exa → Brave
    - `evidence(query)`：证据发现，优先 Exa highlights → Tavily → Brave
    - `patents(query)`：专利检索，走 Serper（Google Patents）；失败则通用搜索补刀
    - `scholar(query)`：学术论文，走 Serper Scholar；失败则通用搜索补刀
    - `news(query)`：新闻检索，走 Serper News；失败则通用搜索补刀

    所有客户端都延迟导入 serper_client，避免没装 SERPAPI_KEY 时 import 炸。
    """

    def __init__(self, *, strict_specialized: bool = True) -> None:
        self._exa: ExaClient | None = None
        self._tavily: TavilyClient | None = None
        self._brave: BraveClient | None = None
        self._serper = None  # 惰性实例化
        self.strict_specialized = strict_specialized
        try:
            self._exa = ExaClient()
        except SearchError:
            pass
        try:
            self._tavily = TavilyClient()
        except SearchError:
            pass
        try:
            self._brave = BraveClient()
        except SearchError:
            pass
        self._has_serper_key = bool(os.environ.get("SERPER_API_KEY") or os.environ.get("SERPAPI_KEY"))
        if not (self._exa or self._tavily or self._brave or self._has_serper_key):
            raise SearchError("no search API key available: set SERPER_API_KEY, SERPAPI_KEY, EXA_API_KEY, TAVILY_API_KEY, or BRAVE_API_KEY")

    def _get_serper(self):
        """惰性创建 SerperClient。没 key 时返回 None。"""
        if self._serper is False:
            return None
        if self._serper is None:
            try:
                from scripts.lib.serper_client import SerperClient
                self._serper = SerperClient()
            except Exception:
                self._serper = False
                return None
        return self._serper

    def close(self) -> None:
        if self._exa:
            self._exa.close()
        if self._tavily:
            self._tavily.close()
        if self._brave:
            self._brave.close()
        if self._serper and self._serper is not False:
            self._serper.close()

    def __enter__(self) -> "SearchClient":
        return self

    def __exit__(self, *_args: Any) -> None:
        self.close()

    def search(self, query: str, *, num_results: int = 5) -> list[SearchHit]:
        """通用网页搜索。Tavily 首选，Exa/Brave 备选。"""
        if self._tavily:
            try:
                return self._tavily.search(query, num_results=num_results)
            except SearchError:
                pass
        if self._exa:
            try:
                return self._exa.search(query, num_results=num_results)
            except SearchError:
                pass
        if self._brave:
            try:
                return self._brave.search(query, num_results=num_results)
            except SearchError:
                pass
        return []

    def evidence(
        self,
        query: str,
        *,
        num_results: int = 10,
        category: str | None = None,
    ) -> list[SearchHit]:
        """Evidence discovery route.

        Exa is better suited for agent-facing evidence discovery because it can
        return concise highlights/text per URL. Results are still candidate
        sources only; downstream packets must score and trace important hits
        back to original Tier 1-2 sources before making final claims.
        """
        if self._exa:
            try:
                return self._exa.search(
                    query,
                    num_results=num_results,
                    search_type="auto",
                    category=category,
                    use_highlights=True,
                    max_characters=1200,
                )
            except SearchError:
                pass
        if self._tavily:
            try:
                return self._tavily.search(query, num_results=num_results)
            except SearchError:
                pass
        if self._brave:
            try:
                return self._brave.search(query, num_results=num_results)
            except SearchError:
                pass
        return []

    def patents(self, query: str, *, num_results: int = 10) -> list[SearchHit]:
        """专利检索：Serper 走 Google Patents 最准。降级到通用搜索 + site 限定。"""
        serper = self._get_serper()
        if serper:
            try:
                hits = serper.patents(query, num_results=num_results)
                return [SearchHit(h.title, h.url, h.snippet) for h in hits]
            except Exception as exc:
                if self.strict_specialized:
                    raise SearchError(f"serper patents failed: {exc}") from exc
        # 降级：通用搜索加 site 限定
        if self.strict_specialized:
            raise SearchError("serper unavailable for patents route; refusing silent fallback")
        return self.search(f"site:patents.google.com {query}", num_results=num_results)

    def scholar(
        self,
        query: str,
        *,
        num_results: int = 10,
        year_low: int | None = None,
    ) -> list[SearchHit]:
        """学术论文：Serper Scholar 带引用数。降级到通用搜索。"""
        serper = self._get_serper()
        if serper:
            try:
                hits = serper.scholar(query, num_results=num_results, year_low=year_low)
                return [
                    SearchHit(
                        title=h.title,
                        url=h.url,
                        snippet=f"{h.snippet}  | {h.source} | 引用 {h.cited_by}" if h.cited_by else h.snippet,
                    )
                    for h in hits
                ]
            except Exception as exc:
                if self.strict_specialized:
                    raise SearchError(f"serper scholar failed: {exc}") from exc
        if self.strict_specialized:
            raise SearchError("serper unavailable for scholar route; refusing silent fallback")
        return self.search(query, num_results=num_results)

    def news(
        self,
        query: str,
        *,
        num_results: int = 10,
        time_range: str | None = None,
    ) -> list[SearchHit]:
        """新闻检索：Serper News。降级到通用搜索。"""
        serper = self._get_serper()
        if serper:
            try:
                hits = serper.news(query, num_results=num_results, time_range=time_range)
                return [
                    SearchHit(
                        title=h.title,
                        url=h.url,
                        snippet=f"{h.snippet}  | {h.source} | {h.date}" if h.date else h.snippet,
                    )
                    for h in hits
                ]
            except Exception as exc:
                if self.strict_specialized:
                    raise SearchError(f"serper news failed: {exc}") from exc
        if self.strict_specialized:
            raise SearchError("serper unavailable for news route; refusing silent fallback")
        return self.search(query, num_results=num_results)

    def fda(self, query: str, *, num_results: int = 10) -> list[SearchHit]:
        """FDA-focused discovery for warning letters and meeting records.

        FDA enforcement examples are often more useful for GMP remediation than
        generic web pages, so this route biases discovery toward warning
        letters, inspection/enforcement pages, and meeting materials/minutes.
        """
        def fda_only(hits: list[SearchHit]) -> list[SearchHit]:
            return [hit for hit in hits if "fda.gov" in (hit.url or "").lower()]

        focused_queries = [
            f'site:fda.gov "Warning Letter" GMP pharmaceutical {query}',
            f'site:fda.gov/inspections-compliance-enforcement-and-criminal-investigations "Warning Letter" {query}',
            f'site:fda.gov "meeting materials" "pharmaceutical quality" {query}',
            f'site:fda.gov "meeting minutes" FDA pharmaceutical quality {query}',
        ]
        hits: list[SearchHit] = []
        seen: set[str] = set()
        per_query = max(2, min(num_results, 4))
        for focused_query in focused_queries:
            route_hits: list[SearchHit] = []
            serper = self._get_serper()
            if serper:
                try:
                    route_hits = [
                        SearchHit(h.title, h.url, h.snippet)
                        for h in serper.search(focused_query, num_results=per_query)
                    ]
                except Exception as exc:
                    if self.strict_specialized:
                        raise SearchError(f"serper FDA search failed: {exc}") from exc
            if not route_hits and not self.strict_specialized:
                route_hits = self.search(focused_query, num_results=per_query)
            for hit in fda_only(route_hits):
                key = hit.url or hit.title
                if not key or key in seen:
                    continue
                seen.add(key)
                hits.append(hit)
                if len(hits) >= num_results:
                    return hits
        return hits


if __name__ == "__main__":
    from scripts.lib.zenmux_client import load_secrets
    load_secrets()
    with SearchClient() as c:
        print("--- 通用: Mabwell 迈威生物 ---")
        for h in c.search("Mabwell 迈威生物 biopharmaceutical", num_results=3):
            print(f"  {h.title[:80]}")
            print(f"    {h.url}")
        print("\n--- 专利: dual-target siRNA ---")
        for h in c.patents("dual-target siRNA GalNAc", num_results=3):
            print(f"  {h.title[:80]}")
            print(f"    {h.url}")
        print("\n--- Scholar: dual-target RNAi 2024 ---")
        for h in c.scholar("dual-target RNAi drug", num_results=3, year_low=2023):
            print(f"  {h.title[:80]}")
            print(f"    {h.url}")