super-search/server.py

#!/usr/bin/env python3
"""
Super Search MCP Server v2.0.0 — Unified web search + extraction via FastMCP.

Tools (10 total):
  1.  web_search(query, limit, category, time_range, site_filter)
        → SearXNG → Exa → OpenCorporates → CourtListener → DuckDuckGo → Firecrawl
  2.  web_extract(url, char_limit)
        → Trafilatura → Firecrawl
  3.  web_search_premium(query, limit, include_domains, exclude_domains, date_start, date_end)
        → Exa direct (DRE-grade)
  4.  web_search_news(query, limit, days)
        → SearXNG (news) → Exa
  5.  web_search_academic(query, limit)
        → SearXNG (science) → Exa
  6.  web_extract_batch(urls, char_limit)
        → Parallel Trafilatura → Firecrawl per URL
  7.  web_suggest_queries(query, limit)
        → SearXNG suggestions
  8.  health_check()
        → Provider statuses + cache stats
  9.  web_lookup(query, limit)
        → DuckDuckGo Instant Answers + Wikipedia
  10. web_search_images(query, limit)
        → SearXNG (images) → Exa

Providers: SearXNG, Exa, OpenCorporates, CourtListener, DuckDuckGo, Wikipedia, Firecrawl
Features: token-bucket rate limiting, TTL caching, retry with exponential backoff, provider telemetry
Endpoint: http://127.0.0.1:8899/mcp
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import sys
import time
from pathlib import Path
from typing import Any

import httpx
import trafilatura
from dotenv import load_dotenv
from fastmcp import FastMCP
from firecrawl import FirecrawlApp

from ratelimit import (
    CACHE_TTL,
    RATE_LIMITS,
    TTLCache,
    TokenBucket,
    cache_key,
    rate_limited_search,
    get_bucket,
)

# ── Logging ──────────────────────────────────────────────────────────────────────

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s: %(message)s")
logger = logging.getLogger("super-search")

# ── Environment ─────────────────────────────────────────────────────────────────

env_path = Path.home() / ".hermes" / ".env"
if env_path.exists():
    load_dotenv(env_path)

FIRECRAWL_API_KEY = os.getenv("FIRECRAWL_API_KEY", "")
EXA_API_KEY = os.getenv("EXA_API_KEY", "")

# ── Provider URLs ───────────────────────────────────────────────────────────────

SEARXNG_URL = "http://127.0.0.1:8888/search"
EXA_SEARCH_URL = "https://api.exa.ai/search"
OPENCORPORATES_URL = "https://api.opencorporates.com/v0.4/companies/search"
COURTLISTENER_URL = "https://www.courtlistener.com/api/rest/v3/opinions/"
DUCKDUCKGO_URL = "https://api.duckduckgo.com/"
WIKIPEDIA_API_URL = "https://en.wikipedia.org/w/api.php"
WIKIPEDIA_REST_URL = "https://en.wikipedia.org/api/rest_v1/page/summary/"

# ── SearXNG category map ────────────────────────────────────────────────────────

SEARXNG_CATEGORIES = {
    "general": "general",
    "news": "news",
    "science": "science",
    "images": "images",
    "videos": "videos",
    "files": "files",
    "social_media": "social+media",
    "music": "music",
    "map": "map",
    "it": "it",
}

# ── Server ──────────────────────────────────────────────────────────────────────

mcp = FastMCP("Super Search", version="2.0.0")

# ── Globals ─────────────────────────────────────────────────────────────────────

_cache = TTLCache()
_server_start = time.time()

# ═══════════════════════════════════════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════════════════════════════════════

def _get_firecrawl():
    """Lazily create FirecrawlApp — avoids crash if key is missing until needed."""
    if not FIRECRAWL_API_KEY:
        raise RuntimeError("FIRECRAWL_API_KEY not set in ~/.hermes/.env")
    return FirecrawlApp(api_key=FIRECRAWL_API_KEY)


def _normalize_results(raw: list[dict]) -> list[dict]:
    """Ensure every result has title, url, description keys."""
    out = []
    for r in raw:
        if not isinstance(r, dict):
            continue
        out.append({
            "title": r.get("title", "") or "",
            "url": r.get("url", "") or "",
            "description": r.get("description", "") or r.get("snippet", "") or "",
        })
    return out


# ═══════════════════════════════════════════════════════════════════════════════
# Provider: SearXNG
# ═══════════════════════════════════════════════════════════════════════════════

async def _searxng_raw(
    query: str,
    limit: int = 5,
    category: str = "general",
    time_range: str = "",
) -> dict[str, Any]:
    """Query SearXNG and return results plus health telemetry."""
    cat = SEARXNG_CATEGORIES.get(category, "general")
    params: dict[str, Any] = {
        "q": query,
        "format": "json",
        "language": "en-US",
        "categories": cat,
    }
    if time_range:
        params["time_range"] = time_range  # day, week, month, year

    async with httpx.AsyncClient(timeout=15) as client:
        resp = await client.get(SEARXNG_URL, params=params)
        resp.raise_for_status()
        data = resp.json()

    results = data.get("results", [])[:limit]
    normalized = _normalize_results(results)

    return {
        "results": normalized,
        "health": {
            "result_count": len(normalized),
            "unresponsive_engines": data.get("unresponsive_engines", []),
            "answers_count": len((data.get("answers") or [])),
            "suggestions_count": len((data.get("suggestions") or [])),
        },
    }


async def _searxng_search(
    query: str,
    limit: int = 5,
    category: str = "general",
    time_range: str = "",
) -> list[dict]:
    return (await _searxng_raw(query, limit, category, time_range))["results"]


async def _searxng_suggestions(query: str) -> list[str]:
    """Return suggested queries from SearXNG."""
    async with httpx.AsyncClient(timeout=15) as client:
        resp = await client.get(
            SEARXNG_URL,
            params={"q": query, "format": "json", "categories": "general"},
        )
        resp.raise_for_status()
        data = resp.json()
    return data.get("suggestions", [])


# ═══════════════════════════════════════════════════════════════════════════════
# Provider: Exa
# ═══════════════════════════════════════════════════════════════════════════════

async def _exa_search(
    query: str,
    limit: int = 5,
    search_type: str = "keyword",
    include_domains: list[str] | None = None,
    exclude_domains: list[str] | None = None,
    start_date: str = "",
    end_date: str = "",
) -> list[dict]:
    """Search via Exa API with optional filters."""
    if not EXA_API_KEY:
        raise RuntimeError("EXA_API_KEY not set in ~/.hermes/.env")

    body: dict[str, Any] = {
        "query": query,
        "numResults": limit,
        "type": search_type,
    }
    if include_domains:
        body["includeDomains"] = include_domains
    if exclude_domains:
        body["excludeDomains"] = exclude_domains
    if start_date:
        body["startPublishedDate"] = start_date
    if end_date:
        body["endPublishedDate"] = end_date

    async with httpx.AsyncClient(timeout=15) as client:
        resp = await client.post(
            EXA_SEARCH_URL,
            headers={
                "Authorization": f"Bearer {EXA_API_KEY}",
                "Content-Type": "application/json",
            },
            json=body,
        )
        resp.raise_for_status()
        data = resp.json()

    results = data.get("results", [])
    return [
        {
            "title": r.get("title", ""),
            "url": r.get("url", ""),
            "description": r.get("text", "")[:500] if r.get("text") else r.get("snippet", ""),
            "published_date": r.get("publishedDate", ""),
        }
        for r in results
    ]


# ═══════════════════════════════════════════════════════════════════════════════
# Provider: Firecrawl
# ═══════════════════════════════════════════════════════════════════════════════

async def _firecrawl_search(query: str, limit: int = 5) -> list[dict]:
    """Search via Firecrawl API."""
    fc = _get_firecrawl()
    response = fc.search(query=query, limit=limit)

    results = []
    if isinstance(response, dict):
        results = response.get("data") or response.get("results") or []
    elif hasattr(response, "data"):
        results = getattr(response, "data") or []
    elif hasattr(response, "results"):
        results = getattr(response, "results") or []
    else:
        results = response or []

    normalized = []
    for r in results:
        if not isinstance(r, dict):
            r = getattr(r, "__dict__", {})
        normalized.append({
            "title": r.get("title") or r.get("markdown", "")[:80] or r.get("url", ""),
            "url": r.get("url", ""),
            "description": r.get("description") or r.get("content") or "",
        })
    return normalized


def _firecrawl_extract(url: str) -> str:
    """Extract via Firecrawl scrape API."""
    fc = _get_firecrawl()
    response = fc.scrape(url, formats=["markdown"])
    if isinstance(response, dict):
        return response.get("markdown") or response.get("content") or json.dumps(response)
    if hasattr(response, "markdown"):
        return getattr(response, "markdown") or ""
    if hasattr(response, "content"):
        return getattr(response, "content") or ""
    return str(response)


# ═══════════════════════════════════════════════════════════════════════════════
# Provider: DuckDuckGo Instant Answers
# ═══════════════════════════════════════════════════════════════════════════════

@rate_limited_search("duckduckgo")
async def _duckduckgo_lookup(query: str, limit: int = 5) -> list[dict]:
    """Instant answers + related topics from DuckDuckGo (free, no key)."""
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.get(
            DUCKDUCKGO_URL,
            params={"q": query, "format": "json", "no_html": "1", "skip_disambig": "1"},
        )
        resp.raise_for_status()
        data = resp.json()

    results = []

    # Instant answer / abstract
    abstract = data.get("AbstractText", "")
    abstract_url = data.get("AbstractURL", "")
    abstract_source = data.get("AbstractSource", "")
    heading = data.get("Heading", "")

    if abstract:
        results.append({
            "title": heading or abstract_source or "Instant Answer",
            "url": abstract_url or "",
            "description": abstract,
        })

    # Related topics
    for topic in data.get("RelatedTopics", [])[:limit]:
        if isinstance(topic, dict) and "Text" in topic:
            results.append({
                "title": topic.get("FirstURL", "").split("/")[-1].replace("_", " "),
                "url": topic.get("FirstURL", ""),
                "description": topic.get("Text", ""),
            })

    # Infobox / definition
    answer = data.get("Answer", "")
    answer_type = data.get("AnswerType", "")
    if answer and answer_type not in ("calc",):
        results.append({
            "title": f"{answer_type.title() if answer_type else 'Answer'}",
            "url": "",
            "description": answer,
        })

    return results[:limit]


# ═══════════════════════════════════════════════════════════════════════════════
# Provider: Wikipedia
# ═══════════════════════════════════════════════════════════════════════════════

@rate_limited_search("wikipedia")
async def _wikipedia_search(query: str, limit: int = 5) -> list[dict]:
    """Search Wikipedia articles."""
    async with httpx.AsyncClient(timeout=10) as client:
        # Search for page titles
        resp = await client.get(
            WIKIPEDIA_API_URL,
            params={
                "action": "query",
                "list": "search",
                "srsearch": query,
                "format": "json",
                "srlimit": limit,
            },
        )
        resp.raise_for_status()
        data = resp.json()

    results = []
    for hit in data.get("query", {}).get("search", [])[:limit]:
        title = hit.get("title", "")
        snippet = hit.get("snippet", "")
        # Strip HTML from snippet
        snippet = snippet.replace("<span class='searchmatch'>", "").replace("</span>", "")
        results.append({
            "title": title,
            "url": f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}",
            "description": snippet,
        })
    return results


async def _wikipedia_summary(title: str) -> str | None:
    """Get the extract/summary for a Wikipedia article."""
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.get(
            WIKIPEDIA_REST_URL + title.replace(" ", "_"),
            headers={"User-Agent": "SuperSearch/2.0.0"},
        )
        if resp.status_code == 404:
            return None
        resp.raise_for_status()
        data = resp.json()
    return data.get("extract", "")


# ═══════════════════════════════════════════════════════════════════════════════
# Provider: OpenCorporates
# ═══════════════════════════════════════════════════════════════════════════════

@rate_limited_search("opencorporates")
async def _opencorporates_search(query: str, limit: int = 5) -> list[dict]:
    """Search company records via OpenCorporates API (free, no key needed)."""
    async with httpx.AsyncClient(timeout=15) as client:
        resp = await client.get(
            OPENCORPORATES_URL,
            params={"q": query, "per_page": limit},
            headers={"User-Agent": "SuperSearch/2.0.0"},
        )
        resp.raise_for_status()
        data = resp.json()

    companies = data.get("results", {}).get("companies", [])[:limit]
    results = []
    for entry in companies:
        company = entry.get("company", {})
        name = company.get("name", "")
        company_number = company.get("company_number", "")
        jurisdiction = company.get("jurisdiction_code", "")
        raw_address = company.get("registered_address_in_full", "")
        url = f"https://opencorporates.com/companies/{jurisdiction}/{company_number}"

        desc_parts = []
        if raw_address:
            desc_parts.append(raw_address)
        if jurisdiction:
            desc_parts.append(f"Jurisdiction: {jurisdiction.upper()}")

        results.append({
            "title": name or company_number,
            "url": url,
            "description": "; ".join(desc_parts) if desc_parts else f"Company #{company_number}",
        })
    return results


# ═══════════════════════════════════════════════════════════════════════════════
# Provider: CourtListener
# ═══════════════════════════════════════════════════════════════════════════════

@rate_limited_search("courtlistener")
async def _courtlistener_search(query: str, limit: int = 5) -> list[dict]:
    """Search US court opinions via CourtListener API (free, no key needed)."""
    async with httpx.AsyncClient(timeout=15) as client:
        resp = await client.get(
            COURTLISTENER_URL,
            params={"q": query, "format": "json"},
            headers={"User-Agent": "SuperSearch/2.0.0"},
        )
        resp.raise_for_status()
        data = resp.json()

    opinions = data.get("results", [])[:limit]
    results = []
    for opinion in opinions:
        case_name = opinion.get("caseName", "")
        court = opinion.get("court", "")
        date_filed = opinion.get("dateFiled", "")[:10] if opinion.get("dateFiled") else ""
        snippet = (opinion.get("snippet") or opinion.get("plain_text") or "")[:200]

        title = case_name or "Untitled"
        abs_url = opinion.get("absolute_url", "")

        desc_parts = []
        if court:
            desc_parts.append(court)
        if date_filed:
            desc_parts.append(date_filed)
        if snippet:
            desc_parts.append(snippet)

        results.append({
            "title": title,
            "url": abs_url,
            "description": " | ".join(desc_parts) if desc_parts else "",
        })
    return results


# ═══════════════════════════════════════════════════════════════════════════════
# Extraction
# ═══════════════════════════════════════════════════════════════════════════════

async def _trafilatura_extract(url: str) -> str:
    """Extract clean text from a URL using trafilatura."""
    async with httpx.AsyncClient(timeout=30, follow_redirects=True) as client:
        resp = await client.get(
            url,
            headers={"User-Agent": "Mozilla/5.0 (compatible; SuperSearch/2.0)"},
        )
        resp.raise_for_status()
        text = trafilatura.extract(
            resp.text,
            include_comments=False,
            include_tables=True,
            include_images=False,
            include_links=True,
            output_format="markdown",
        )
        if not text:
            raise ValueError("Trafilatura could not extract content")
        return text


async def _extract_one(url: str, char_limit: int = 0) -> dict[str, Any]:
    """Extract a single URL, returning {url, content, provider, error?}."""
    # Try Trafilatura
    try:
        text = await _trafilatura_extract(url)
        if char_limit and len(text) > char_limit:
            text = text[:char_limit] + f"\n\n[...truncated at {char_limit} chars]"
        return {"url": url, "content": text, "provider": "trafilatura"}
    except Exception as e:
        traf_err = str(e)

    # Fallback to Firecrawl
    try:
        text = _firecrawl_extract(url)
        if char_limit and len(text) > char_limit:
            text = text[:char_limit] + f"\n\n[...truncated at {char_limit} chars]"
        return {"url": url, "content": text, "provider": "firecrawl"}
    except Exception as e:
        return {"url": url, "content": "", "provider": "none", "error": f"Trafilatura: {traf_err}; Firecrawl: {e}"}


# ═══════════════════════════════════════════════════════════════════════════════
# Cache stats helper
# ═══════════════════════════════════════════════════════════════════════════════

def _cache_stats() -> dict:
    """Return cache statistics."""
    # Access internals of the global _cache for monitoring
    cache_dict = getattr(_cache, "_cache", {})
    total_entries = len(cache_dict)
    active_entries = 0
    now = time.monotonic()
    for key, (_, expires_at) in cache_dict.items():
        if expires_at > now:
            active_entries += 1
    return {
        "total_entries": total_entries,
        "active_entries": active_entries,
        "expired_entries": total_entries - active_entries,
    }


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 1: web_search (enhanced)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Search the web. Tries SearXNG first, falls back to Exa, OpenCorporates, CourtListener, DuckDuckGo, then Firecrawl. Supports category, time_range, and site_filter.")
async def web_search(
    query: str,
    limit: int = 5,
    category: str = "general",
    time_range: str = "",
    site_filter: str = "",
) -> str:
    """Search the web with provider fallback.

    Args:
        query: Search query string.
        limit: Max results (1-20, default 5).
        category: One of general, news, science, images, videos, files, social_media, music, map, it.
        time_range: Optional — day, week, month, year.
        site_filter: Optional — limit to domain (e.g. "github.com").

    Returns JSON with provider, results, and telemetry.
    """
    limit = max(1, min(limit, 20))
    telemetry: list[dict] = []

    # Build query with site filter
    search_query = query
    if site_filter:
        search_query = f"{query} site:{site_filter}"

    # 1. Try SearXNG
    try:
        searx = await _searxng_raw(search_query, limit, category, time_range)
        results = searx["results"]
        telemetry.append({"provider": "searxng", "status": "ok", **searx.get("health", {})})
        if results:
            return json.dumps(
                {"provider": "searxng", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "searxng", "status": "error", "error": str(e)[:200]})

    # 2. Fallback to Exa
    try:
        results = await _exa_search(search_query, limit)
        telemetry.append({"provider": "exa", "status": "ok", "result_count": len(results)})
        if results:
            return json.dumps(
                {"provider": "exa", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "exa", "status": "error", "error": str(e)[:200]})

    # 3. Fallback to OpenCorporates
    try:
        results = await _opencorporates_search(query, limit)
        telemetry.append({"provider": "opencorporates", "status": "ok", "result_count": len(results)})
        if results:
            return json.dumps(
                {"provider": "opencorporates", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "opencorporates", "status": "error", "error": str(e)[:200]})

    # 4. Fallback to CourtListener
    try:
        results = await _courtlistener_search(query, limit)
        telemetry.append({"provider": "courtlistener", "status": "ok", "result_count": len(results)})
        if results:
            return json.dumps(
                {"provider": "courtlistener", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "courtlistener", "status": "error", "error": str(e)[:200]})

    # 5. Fallback to DuckDuckGo
    try:
        results = await _duckduckgo_lookup(query, limit)
        telemetry.append({"provider": "duckduckgo", "status": "ok", "result_count": len(results)})
        if results:
            return json.dumps(
                {"provider": "duckduckgo", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "duckduckgo", "status": "error", "error": str(e)[:200]})

    # 6. Last resort: Firecrawl
    try:
        results = await _firecrawl_search(query, limit)
        telemetry.append({"provider": "firecrawl", "status": "ok", "result_count": len(results)})
        return json.dumps(
            {"provider": "firecrawl", "results": results, "telemetry": telemetry},
            indent=2,
        )
    except Exception as e:
        telemetry.append({"provider": "firecrawl", "status": "error", "error": str(e)[:200]})
        return json.dumps(
            {"error": "All search providers failed", "telemetry": telemetry},
            indent=2,
        )


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 2: web_extract (enhanced)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Extract clean text (markdown) from a URL. Tries Trafilatura first, falls back to Firecrawl.")
async def web_extract(url: str, char_limit: int = 0) -> str:
    """Extract clean markdown text from a URL.

    Args:
        url: The URL to extract content from.
        char_limit: Optional max characters (0 = no limit).

    Returns markdown text or error message.
    """
    result = await _extract_one(url, char_limit)
    if result.get("error"):
        return f"Extraction failed for {url}: {result['error']}"
    return result["content"]


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 3: web_search_premium (enhanced)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Premium web search using Exa API directly. Best for DRE/debt-recovery research and skip tracing. Supports domain and date filters.")
async def web_search_premium(
    query: str,
    limit: int = 5,
    include_domains: str = "",
    exclude_domains: str = "",
    date_start: str = "",
    date_end: str = "",
) -> str:
    """Search the web using Exa's neural search API with advanced filters.

    Args:
        query: Search query.
        limit: Max results (1-20, default 5).
        include_domains: Comma-separated domains to include (e.g. "linkedin.com,github.com").
        exclude_domains: Comma-separated domains to exclude.
        date_start: Start date in YYYY-MM-DD format for filtering results.
        date_end: End date in YYYY-MM-DD format.

    Returns JSON with provider and results.
    """
    limit = max(1, min(limit, 20))
    inc = [d.strip() for d in include_domains.split(",") if d.strip()] if include_domains else None
    exc = [d.strip() for d in exclude_domains.split(",") if d.strip()] if exclude_domains else None

    try:
        results = await _exa_search(
            query,
            limit=limit,
            include_domains=inc,
            exclude_domains=exc,
            start_date=date_start,
            end_date=date_end,
        )
        return json.dumps(
            {"provider": "exa-premium", "results": results, "result_count": len(results)},
            indent=2,
        )
    except Exception as e:
        return json.dumps({"error": f"Exa search failed: {e}", "results": []}, indent=2)


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 4: web_search_news (NEW)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Search news articles. Uses SearXNG news category, falls back to Exa with date filtering for recent results.")
async def web_search_news(query: str, limit: int = 5, days: int = 7) -> str:
    """Search for recent news articles.

    Args:
        query: Search query.
        limit: Max results (1-20, default 5).
        days: Look back N days (default 7, max 365).

    Returns JSON with provider, results, and telemetry.
    """
    limit = max(1, min(limit, 20))
    days = max(1, min(days, 365))
    telemetry: list[dict] = []

    # 1. Try SearXNG news category
    try:
        searx = await _searxng_raw(query, limit, category="news")
        results = searx["results"]
        telemetry.append({"provider": "searxng-news", "status": "ok", **searx.get("health", {})})
        if results:
            return json.dumps(
                {"provider": "searxng-news", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "searxng-news", "status": "error", "error": str(e)[:200]})

    # 2. Fallback to Exa with date range
    try:
        from datetime import date, timedelta
        end = date.today().isoformat()
        start = (date.today() - timedelta(days=days)).isoformat()
        results = await _exa_search(query, limit=limit, start_date=start, end_date=end)
        telemetry.append({"provider": "exa", "status": "ok", "result_count": len(results)})
        if results:
            return json.dumps(
                {"provider": "exa", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "exa", "status": "error", "error": str(e)[:200]})

    return json.dumps(
        {"error": "News search failed — all providers unavailable", "telemetry": telemetry},
        indent=2,
    )


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 5: web_search_academic (NEW)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Search academic/scholarly articles. Uses SearXNG science category, falls back to Exa.")
async def web_search_academic(query: str, limit: int = 5) -> str:
    """Search for academic papers and scholarly articles.

    Args:
        query: Search query (e.g., author name, paper title, topic).
        limit: Max results (1-20, default 5).

    Returns JSON with provider, results, and telemetry.
    """
    limit = max(1, min(limit, 20))
    telemetry: list[dict] = []

    # 1. Try SearXNG science category
    try:
        searx = await _searxng_raw(query, limit, category="science")
        results = searx["results"]
        telemetry.append({"provider": "searxng-science", "status": "ok", **searx.get("health", {})})
        if results:
            return json.dumps(
                {"provider": "searxng-science", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "searxng-science", "status": "error", "error": str(e)[:200]})

    # 2. Fallback to Exa
    try:
        results = await _exa_search(f"{query} research paper", limit=limit)
        telemetry.append({"provider": "exa", "status": "ok", "result_count": len(results)})
        if results:
            return json.dumps(
                {"provider": "exa", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "exa", "status": "error", "error": str(e)[:200]})

    return json.dumps(
        {"error": "Academic search failed — all providers unavailable", "telemetry": telemetry},
        indent=2,
    )


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 6: web_extract_batch (NEW)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Extract clean text from multiple URLs in parallel. Each URL tries Trafilatura first, then falls back to Firecrawl.")
async def web_extract_batch(urls: str, char_limit: int = 0) -> str:
    """Extract content from multiple URLs concurrently.

    Args:
        urls: Comma-separated or newline-separated list of URLs (max 10).
        char_limit: Optional max characters per extraction (0 = no limit).

    Returns JSON array with {url, content, provider} per URL.
    """
    # Parse URLs
    url_list = []
    for part in urls.replace("\n", ",").split(","):
        stripped = part.strip()
        if stripped:
            url_list.append(stripped)

    url_list = url_list[:10]  # Hard limit at 10

    if not url_list:
        return json.dumps({"error": "No URLs provided"}, indent=2)

    # Extract all in parallel
    tasks = [_extract_one(url, char_limit) for url in url_list]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    output = []
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            output.append({"url": url_list[i], "content": "", "provider": "none", "error": str(result)})
        else:
            output.append(result)

    return json.dumps(output, indent=2)


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 7: web_suggest_queries (NEW)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Get related search query suggestions from SearXNG. Useful for expanding or refining a search.")
async def web_suggest_queries(query: str, limit: int = 5) -> str:
    """Get suggested related queries for a search term.

    Args:
        query: The search query to get suggestions for.
        limit: Max suggestions to return (1-10, default 5).

    Returns JSON with suggestions array.
    """
    limit = max(1, min(limit, 10))
    try:
        suggestions = await _searxng_suggestions(query)
        return json.dumps(
            {"query": query, "suggestions": suggestions[:limit]},
            indent=2,
        )
    except Exception as e:
        return json.dumps({"query": query, "suggestions": [], "error": str(e)[:200]}, indent=2)


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 8: health_check (NEW)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Check the health and status of all search providers. Returns uptime, provider reachability, and cache statistics.")
async def health_check() -> str:
    """Check provider health and server status.

    Returns JSON with uptime, provider statuses, and cache stats.
    """
    uptime = time.time() - _server_start
    providers: dict[str, dict] = {}

    # Test each provider
    provider_tests = [
        ("searxng", lambda: httpx.get("http://127.0.0.1:8888/health", timeout=5)),
        ("exa", lambda: httpx.get("https://api.exa.ai/health", timeout=5) if EXA_API_KEY else None),
        ("duckduckgo", lambda: httpx.get("https://api.duckduckgo.com/?q=test&format=json&no_html=1", timeout=5)),
        ("wikipedia", lambda: httpx.get("https://en.wikipedia.org/w/api.php?action=query&meta=siteinfo&format=json", timeout=5)),
        ("opencorporates", lambda: httpx.get("https://api.opencorporates.com/v0.4/companies/search?q=test&per_page=1", timeout=5)),
        ("courtlistener", lambda: httpx.get("https://www.courtlistener.com/api/rest/v3/opinions/?format=json", timeout=5)),
        ("firecrawl", lambda: (lambda fc: fc.search(query="test", limit=1))(_get_firecrawl()) if FIRECRAWL_API_KEY else None),
    ]

    async def _test_one(name: str, factory) -> tuple[str, dict]:
        try:
            if factory is None:
                return name, {"status": "disabled", "reason": "no credentials"}
            result = factory()
            if asyncio.iscoroutine(result):
                result = await result
            return name, {"status": "ok"}
        except Exception as e:
            return name, {"status": "error", "error": str(e)[:200]}

    tests = [_test_one(name, factory) for name, factory in provider_tests]
    results = await asyncio.gather(*tests)

    for name, status in results:
        providers[name] = status

    # Rate limit bucket statuses
    from ratelimit import _buckets
    buckets = {}
    for name, bucket in _buckets.items():
        buckets[name] = {
            "tokens_available": round(bucket._tokens, 2),
            "capacity": round(bucket._capacity, 2),
            "refill_rate": round(bucket._refill_rate, 2),
        }

    return json.dumps({
        "server": {
            "version": "2.0.0",
            "uptime_seconds": round(uptime),
            "uptime_human": f"{int(uptime // 3600)}h {int((uptime % 3600) // 60)}m",
        },
        "providers": providers,
        "cache": _cache_stats(),
        "rate_limits": buckets,
    }, indent=2)


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 9: web_lookup (NEW)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Look up a factual query using DuckDuckGo Instant Answers and Wikipedia. Best for definitions, facts, entities, and encyclopedic information.")
async def web_lookup(query: str, limit: int = 5) -> str:
    """Look up factual information using DuckDuckGo + Wikipedia.

    Best for: definitions, entity lookups, encyclopedic facts, quick answers.
    Not for: current news, opinion pieces, product comparisons.

    Args:
        query: What to look up (person, place, concept, definition, etc.).
        limit: Max results from each source (1-10, default 5).

    Returns JSON with instant_answer, ddg_results, and wiki_results.
    """
    limit = max(1, min(limit, 10))
    output: dict[str, Any] = {"query": query}

    # DuckDuckGo instant answer + topics
    try:
        ddg = await _duckduckgo_lookup(query, limit)
        # Separate the first result (instant answer) from the rest
        if ddg:
            output["instant_answer"] = ddg[0]["description"] if ddg[0].get("description") else None
            output["ddg_source"] = ddg[0].get("title", "")
            output["ddg_url"] = ddg[0].get("url", "")
            output["ddg_results"] = ddg[1:limit] if len(ddg) > 1 else []
        else:
            output["ddg_results"] = []
    except Exception as e:
        output["ddg_results"] = []
        output["ddg_error"] = str(e)[:200]

    # Wikipedia search
    try:
        wiki = await _wikipedia_search(query, limit)
        output["wiki_results"] = wiki
        # If we have a good match, try to get the summary
        if wiki and wiki[0].get("title"):
            try:
                summary = await _wikipedia_summary(wiki[0]["title"])
                if summary:
                    output["wiki_summary"] = summary[:1000]
            except Exception:
                pass
    except Exception as e:
        output["wiki_results"] = []
        output["wiki_error"] = str(e)[:200]

    return json.dumps(output, indent=2)


# ═══════════════════════════════════════════════════════════════════════════════
# Tool 10: web_search_images (NEW)
# ═══════════════════════════════════════════════════════════════════════════════

@mcp.tool(description="Search for images. Uses SearXNG images category. Returns image URLs, thumbnails, source pages, and dimensions.")
async def web_search_images(query: str, limit: int = 5) -> str:
    """Search for images on the web.

    Args:
        query: What images to search for.
        limit: Max results (1-20, default 5).

    Returns JSON with image results (url, thumbnail, source, dimensions, title).
    """
    limit = max(1, min(limit, 20))
    telemetry: list[dict] = []

    # 1. Try SearXNG images category
    try:
        searx = await _searxng_raw(query, limit, category="images")
        results = searx["results"]
        # For images, the SearXNG result format includes img_src, thumbnail_src
        telemetry.append({"provider": "searxng-images", "status": "ok", **searx.get("health", {})})
        if results:
            return json.dumps(
                {"provider": "searxng-images", "results": results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "searxng-images", "status": "error", "error": str(e)[:200]})

    # 2. Fallback: SearXNG general search (image results embed in general too)
    try:
        searx = await _searxng_raw(query, limit, category="general")
        # Filter for results that have img_src
        img_results = [r for r in searx["results"] if r.get("img_src") or r.get("thumbnail")]
        telemetry.append({"provider": "searxng-general", "status": "ok", "result_count": len(img_results)})
        if img_results:
            return json.dumps(
                {"provider": "searxng-general", "results": img_results, "telemetry": telemetry},
                indent=2,
            )
    except Exception as e:
        telemetry.append({"provider": "searxng-general", "status": "error", "error": str(e)[:200]})

    return json.dumps(
        {"error": "Image search failed — all providers unavailable", "telemetry": telemetry},
        indent=2,
    )


# ═══════════════════════════════════════════════════════════════════════════════
# Entrypoint
# ═══════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    mcp.run(
        transport="http",
        host="127.0.0.1",
        port=8899,
        path="/mcp",
    )