"""Web search + fetch for codex tasks, following the oh-my-pi pattern.

Zero third-party dependencies; Python standard library only.

Design notes (what is copied from oh-my-pi, and what is not):
- Copied: one unified ``web_search`` tool, a ``SearchProvider`` abstraction,
  and an ``auto`` chain that walks providers in order and falls back to
  the next one on failure. Pinning a single provider by name disables
  fallback, exactly like oh-my-pi's pinned provider.
- Copied: a two-step workflow -- ``web_search`` finds candidate URLs,
  ``web_fetch`` reads a specific page -- and site-agnostic extraction
  that returns clean text plus a link list, not raw HTML.
- NOT copied: oh-my-pi's 23 providers. v1 ships the five that matter
  for a codex task: Perplexity (synthesised answer + citations),
  Exa (semantic search + highlights), Tavily (research + extract),
  Brave (fast, independent index) and two keyless fallbacks
  (SearXNG, DuckDuckGo). Model-native OAuth search providers
  (Gemini / Anthropic / Codex / xAI) are deliberately excluded: inside
  a codex task they add no provider the task does not already have.
- NOT copied: oh-my-pi's TypeScript / Rust extraction stack. Fetch
  here is a small stdlib HTML parser: script/style stripped, text
  collapsed, long pages truncated head+tail like Hermes' deterministic
  budget (default 15000 chars, 75% head / 25% tail, ``[TRUNCATED]``).

Security notes:
- API keys come only from the process environment, never from
  arguments, config files, or this module's return values. Error
  strings contain provider names and HTTP statuses only; they never
  format headers or key values, and the chain test asserts this.
- ``web_fetch`` is not a security boundary: a codex task running
  danger-full-access can already curl anything, so fetch does not
  pretend otherwise. It does refuse non-http(s) URLs and literal
  localhost / private-IP URLs, which stops the most trivial SSRF
  accidents, but it is not a sandbox.
"""

import html as _html
import ipaddress
import json
import os
import re
import socket
import urllib.error
import urllib.parse
import urllib.request
from html.parser import HTMLParser

__all__ = [
    "DEFAULT_CHAIN",
    "WebSearchError",
    "SearchResult",
    "available_providers",
    "resolve_provider_chain",
    "search",
    "fetch_url",
    "ensure_codex_home_websearch",
    "DEFAULT_FETCH_CHAR_LIMIT",
]

DEFAULT_CHAIN = ("perplexity", "exa", "tavily", "brave", "searxng",
                 "duckduckgo")
DEFAULT_FETCH_CHAR_LIMIT = 15000
_USER_AGENT = "codex-websearch-mcp/0.1 (+codex-app-server)"

_KEY_ENV = {
    "perplexity": "PERPLEXITY_API_KEY",
    "exa": "EXA_API_KEY",
    "tavily": "TAVILY_API_KEY",
    "brave": "BRAVE_API_KEY",
}
_FREE_PROVIDERS = ("searxng", "duckduckgo")


class WebSearchError(Exception):
    """All providers in the chain failed (or the pinned one did)."""

    def __init__(self, message, attempted=None):
        super().__init__(message)
        self.attempted = list(attempted or [])


class SearchResult(dict):
    """A single search hit: title, url, snippet, optional score."""

    def __init__(self, title, url, snippet="", score=None):
        super().__init__(title=title or "", url=url or "",
                         snippet=snippet or "")
        if score is not None:
            self["score"] = score


# ---------------------------------------------------------------- HTTP

class HttpClient:
    """Thin urllib wrapper, injectable so tests never touch the net."""

    def __init__(self, timeout=20):
        self.timeout = timeout

    def post_json(self, url, headers, payload):
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            url, data=data,
            headers=dict(headers, **{"Content-Type": "application/json"}),
            method="POST")
        return self._read(req)

    def get(self, url, headers=None):
        req = urllib.request.Request(
            url, headers=headers or {}, method="GET")
        return self._read(req)

    def _read(self, req):
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                body = resp.read()
                return resp.status, body, resp.headers.get_content_type()
        except urllib.error.HTTPError as exc:
            # Never include headers (keys live there) in the error.
            body = exc.read()[:500]
            raise WebSearchError(
                "HTTP %s" % exc.code) from _ProviderHttpError(
                    exc.code, body)

    def get_text(self, url, headers=None):
        status, body, _ctype = self.get(url, headers)
        charset = "utf-8"
        try:
            text = body.decode(charset, errors="replace")
        except Exception:
            text = body.decode("utf-8", errors="replace")
        return status, text


class _ProviderHttpError(Exception):
    def __init__(self, status, body):
        super().__init__("HTTP %s" % status)
        self.status = status
        self.body = body


def _http_error_message(exc):
    """Render a provider failure without ever leaking credentials."""
    if isinstance(exc, _ProviderHttpError):
        return "HTTP %s" % exc.status
    if isinstance(exc, WebSearchError):
        return str(exc)
    return "%s: %s" % (type(exc).__name__, exc)


def _default_headers(env, provider):
    key_env = _KEY_ENV.get(provider)
    key = env.get(key_env) if key_env else None
    if provider == "exa":
        return {"x-api-key": key, "User-Agent": _USER_AGENT}
    if provider == "tavily":
        return {"Authorization": "Bearer %s" % key,
                "User-Agent": _USER_AGENT}
    if provider == "brave":
        return {"X-Subscription-Token": key, "Accept": "application/json",
                "User-Agent": _USER_AGENT}
    if provider == "perplexity":
        return {"Authorization": "Bearer %s" % key,
                "User-Agent": _USER_AGENT}
    return {"User-Agent": _USER_AGENT}


def _clamp_limit(limit, default=5, lo=1, hi=10):
    try:
        value = int(limit)
    except (TypeError, ValueError):
        return default
    return max(lo, min(hi, value))


# ------------------------------------------------------------ providers

def _provider_perplexity(query, limit, env, http):
    # Sonar is a synthesis API: one OpenAI-shaped chat call returns a
    # written, cited answer. We surface both the answer and the
    # underlying sources, so web_search keeps its "answer + citations"
    # shape instead of degrading to a link list.
    _status, body, _ctype = http.post_json(
        "https://api.perplexity.ai/chat/completions",
        _default_headers(env, "perplexity"),
        {"model": "sonar",
         "messages": [{"role": "user", "content": query}]})
    data = json.loads(body.decode("utf-8", errors="replace"))
    choices = data.get("choices") or []
    message = (choices[0].get("message") or {}) if choices else {}
    answer = message.get("content") or None
    results = []
    for item in data.get("search_results") or []:
        if isinstance(item, dict) and item.get("url"):
            results.append(SearchResult(item.get("title") or item["url"],
                                        item["url"],
                                        item.get("snippet") or ""))
    if not results:
        for url in data.get("citations") or []:
            if isinstance(url, str) and url.startswith("http"):
                results.append(SearchResult(url, url, ""))
    return results, answer


def _provider_exa(query, limit, env, http):
    _status, body, _ctype = http.post_json(
        "https://api.exa.ai/search",
        _default_headers(env, "exa"),
        {"query": query, "numResults": limit, "type": "auto",
         "contents": {"text": {"maxCharacters": 2000},
                      "highlights": True}})
    data = json.loads(body.decode("utf-8", errors="replace"))
    results = []
    for item in data.get("results") or []:
        if not isinstance(item, dict) or not item.get("url"):
            continue
        highlights = item.get("highlights") or []
        snippet = item.get("text") or ""
        if highlights and not snippet:
            snippet = " ".join(str(h) for h in highlights[:3])
        elif highlights:
            snippet = snippet[:2000]
        results.append(SearchResult(item.get("title") or item["url"],
                                    item["url"], snippet,
                                    item.get("score")))
    return results, None


def _provider_tavily(query, limit, env, http):
    _status, body, _ctype = http.post_json(
        "https://api.tavily.com/search",
        _default_headers(env, "tavily"),
        {"query": query, "search_depth": "basic",
         "max_results": limit, "include_answer": False,
         "include_raw_content": False})
    data = json.loads(body.decode("utf-8", errors="replace"))
    results = []
    for item in data.get("results") or []:
        if not isinstance(item, dict) or not item.get("url"):
            continue
        results.append(SearchResult(item.get("title") or item["url"],
                                    item["url"],
                                    item.get("content") or "",
                                    item.get("score")))
    return results, data.get("answer") or None


def _provider_brave(query, limit, env, http):
    url = ("https://api.search.brave.com/res/v1/web/search?"
           + urllib.parse.urlencode(
               {"q": query, "count": limit, "extra_snippets": "true"}))
    _status, body, _ctype = http.get(url, _default_headers(env, "brave"))
    data = json.loads(body.decode("utf-8", errors="replace"))
    results = []
    for item in (data.get("web") or {}).get("results") or []:
        if not isinstance(item, dict) or not item.get("url"):
            continue
        parts = [item.get("description") or ""]
        parts.extend(str(s) for s in (item.get("extra_snippets") or []))
        snippet = " ".join(p for p in parts if p)
        results.append(SearchResult(item.get("title") or item["url"],
                                    item["url"], snippet))
    return results, None


def _provider_searxng(query, limit, env, http):
    base = (env.get("SEARXNG_URL") or "").rstrip("/")
    if not base:
        raise WebSearchError("searxng not configured (SEARXNG_URL unset)")
    # Engines are pinned, not left to SearXNG's default "general"
    # category. Measured 2026-10-30 on the production host:
    # default general search took 20.0s in a Docker task container
    # (wikipedia / wikidata / brave / duckduckgo / google engines
    # all timing out, only bing + 360search returning), which hits
    # the MCP's 20s timeout and then falls through to DuckDuckGo,
    # which itself times out via the egress proxy -- so the whole
    # chain fails even though SearXNG can answer. Pinning to the
    # two proven-fast engines cuts the same query to 0.47s via
    # the same proxy. SEARXNG_ENGINES overrides the pin; an empty
    # value restores the old unpinned behaviour.
    params = {"q": query, "format": "json"}
    engines = env.get("SEARXNG_ENGINES", "bing,360search")
    if engines:
        params["engines"] = engines
    url = base + "/search?" + urllib.parse.urlencode(params)
    _status, body, _ctype = http.get(url, _default_headers(env, "searxng"))
    data = json.loads(body.decode("utf-8", errors="replace"))
    results = []
    for item in data.get("results") or []:
        if not isinstance(item, dict) or not item.get("url"):
            continue
        results.append(SearchResult(item.get("title") or item["url"],
                                    item["url"],
                                    item.get("content") or ""))
    return results[:limit], None


class _DuckDuckGoParser(HTMLParser):
    """Best-effort parser for DuckDuckGo's HTML results page.

    Only the last-resort free path uses this; it is deliberately
    tolerant of layout drift (missing classes just yield no results,
    which moves the chain on, it never crashes the tool).
    """

    def __init__(self):
        super().__init__()
        self.results = []
        self._current = None
        self._in_title = False
        self._in_snippet = False

    def handle_starttag(self, tag, attrs):
        attrs = dict(attrs)
        cls = attrs.get("class") or ""
        if tag == "a" and "result__a" in cls:
            href = attrs.get("href") or ""
            self._current = {"title": "", "url": _unwrap_ddg(href),
                             "snippet": ""}
            self._in_title = True
        elif tag == "a" and "result__snippet" in cls:
            self._in_snippet = True

    def handle_endtag(self, tag):
        if tag == "a" and self._in_title:
            self._in_title = False
            if self._current and self._current["url"]:
                self.results.append(self._current)
            self._current = None
        if tag == "a" and self._in_snippet:
            self._in_snippet = False

    def handle_data(self, data):
        if self._in_title and self._current is not None:
            self._current["title"] += data
        elif self._in_snippet and self.results:
            self.results[-1]["snippet"] += data


def _unwrap_ddg(href):
    """DuckDuckGo wraps result links as //duckduckgo.com/l/?uddg=<url>."""
    if not href:
        return ""
    parsed = urllib.parse.urlparse(href)
    if "duckduckgo.com" in (parsed.netloc or ""):
        target = urllib.parse.parse_qs(parsed.query).get("uddg")
        if target:
            return target[0]
        return ""
    if href.startswith("//"):
        return "https:" + href
    return href


def _provider_duckduckgo(query, limit, env, http):
    url = ("https://html.duckduckgo.com/html/?"
           + urllib.parse.urlencode({"q": query}))
    _status, body, _ctype = http.get(
        url, {"User-Agent": _USER_AGENT})
    parser = _DuckDuckGoParser()
    parser.feed(body.decode("utf-8", errors="replace"))
    results = []
    for item in parser.results[:limit]:
        results.append(SearchResult(item["title"].strip(),
                                    item["url"],
                                    item["snippet"].strip()))
    return results, None


_PROVIDERS = {
    "perplexity": _provider_perplexity,
    "exa": _provider_exa,
    "tavily": _provider_tavily,
    "brave": _provider_brave,
    "searxng": _provider_searxng,
    "duckduckgo": _provider_duckduckgo,
}


def available_providers(env=None):
    """Providers usable right now: keyed ones with a key, plus free ones."""
    env = env if env is not None else os.environ
    out = []
    for name in DEFAULT_CHAIN:
        if name in _FREE_PROVIDERS:
            if name == "searxng" and not env.get("SEARXNG_URL"):
                continue
            out.append(name)
        elif env.get(_KEY_ENV[name]):
            out.append(name)
    return out


def resolve_provider_chain(provider="auto", env=None):
    """Resolve the chain oh-my-pi style: auto walks, pinned does not.

    ``provider`` may be "auto" (default chain, filtered to what is
    configured), a single name (that provider only, even if it fails),
    or a comma-separated list (that order only). Unknown names raise
    ValueError before any HTTP happens.
    """
    env = env if env is not None else os.environ
    if provider is None or str(provider).strip() in ("", "auto"):
        chain = list(available_providers(env))
        if not chain:
            chain = list(_FREE_PROVIDERS)
            # duckduckgo needs no configuration; searxng without a URL
            # will fail fast with a clear message, which is better
            # than pretending no provider exists at all.
            chain = ["duckduckgo"]
        return chain
    names = [p.strip() for p in str(provider).split(",") if p.strip()]
    for name in names:
        if name not in _PROVIDERS:
            raise ValueError("unknown search provider: %r" % name)
    return names


def search(query, limit=5, provider="auto", env=None, http=None):
    """Search the web. Returns {provider, results, answer, attempted}.

    Raises WebSearchError (with .attempted) when every provider in
    the chain fails or returns no results. Error strings never contain
    API keys; see the module docstring.
    """
    if not query or not str(query).strip():
        raise ValueError("query must be a non-empty string")
    env = env if env is not None else os.environ
    http = http or HttpClient(
        timeout=int(env.get("WEBSEARCH_TIMEOUT", "20") or 20))
    limit = _clamp_limit(limit)
    chain = resolve_provider_chain(provider, env)
    attempted = []
    for name in chain:
        if name in _KEY_ENV and not env.get(_KEY_ENV[name]):
            attempted.append({"provider": name, "ok": False,
                              "error": "missing API key (%s unset)"
                                       % _KEY_ENV[name]})
            continue
        try:
            results, answer = _PROVIDERS[name](
                str(query), limit, env, http)
        except Exception as exc:  # noqa: BLE001 - chain must continue
            attempted.append({"provider": name, "ok": False,
                              "error": _http_error_message(exc)})
            continue
        if not results:
            attempted.append({"provider": name, "ok": False,
                              "error": "no results"})
            continue
        attempted.append({"provider": name, "ok": True, "error": None})
        return {"provider": name,
                "results": [dict(r) for r in results[:limit]],
                "answer": answer,
                "attempted": attempted}
    raise WebSearchError(
        "all search providers failed: %s"
        % "; ".join("%s (%s)" % (a["provider"], a["error"])
                    for a in attempted),
        attempted=attempted)


# ---------------------------------------------------------------- fetch

class _TextExtractor(HTMLParser):
    def __init__(self):
        super().__init__()
        self.title = ""
        self._in_title = False
        self._skip = 0
        self.blocks = []
        self.links = []
        self._current_block = []

    def handle_starttag(self, tag, attrs):
        if tag in ("script", "style", "noscript"):
            self._skip += 1
            return
        if tag == "title":
            self._in_title = True
        if tag == "a":
            href = dict(attrs).get("href") or ""
            if href.startswith("http"):
                self.links.append(href)
        if tag in ("p", "h1", "h2", "h3", "h4", "li", "pre", "blockquote"):
            self._flush()

    def handle_endtag(self, tag):
        if tag in ("script", "style", "noscript"):
            self._skip = max(0, self._skip - 1)
        if tag == "title":
            self._in_title = False
        if tag in ("p", "h1", "h2", "h3", "h4", "li", "pre", "blockquote"):
            self._flush()

    def handle_data(self, data):
        if self._skip:
            return
        if self._in_title:
            self.title += data
            return
        text = data.strip()
        if text:
            self._current_block.append(text)

    def _flush(self):
        if self._current_block:
            self.blocks.append(" ".join(self._current_block))
            self._current_block = []

    def text(self):
        self._flush()
        return "\n\n".join(self.blocks)


def _truncate(text, char_limit):
    if len(text) <= char_limit:
        return text, False
    head = int(char_limit * 0.75)
    tail = char_limit - head
    return (text[:head] + "\n\n[TRUNCATED]\n\n" + text[-tail:]
            if tail > 0 else text[:head] + "\n\n[TRUNCATED]"), True


def _is_blocked_fetch_url(url):
    parsed = urllib.parse.urlparse(url)
    if parsed.scheme not in ("http", "https"):
        return "only http(s) URLs can be fetched"
    host = (parsed.hostname or "").lower()
    if host in ("localhost",):
        return "localhost URLs are blocked"
    try:
        addr = ipaddress.ip_address(host)
    except ValueError:
        return None
    if addr.is_private or addr.is_loopback or addr.is_link_local:
        return "private / loopback URLs are blocked"
    return None


def fetch_url(url, char_limit=DEFAULT_FETCH_CHAR_LIMIT, env=None,
              http=None):
    """Fetch one URL and return clean text + links (truncated)."""
    if not url or not str(url).strip():
        raise ValueError("url must be a non-empty string")
    env = env if env is not None else os.environ
    if env.get("WEBSEARCH_ALLOW_PRIVATE_URLS") != "1":
        blocked = _is_blocked_fetch_url(str(url))
        if blocked:
            raise WebSearchError(blocked)
    http = http or HttpClient(
        timeout=int(env.get("WEBSEARCH_TIMEOUT", "20") or 20))
    try:
        char_limit = int(char_limit)
    except (TypeError, ValueError):
        char_limit = DEFAULT_FETCH_CHAR_LIMIT
    char_limit = max(2000, min(500000, char_limit))
    _status, body, ctype = http.get(str(url), {"User-Agent": _USER_AGENT})
    if ctype and not any(
            token in ctype for token in
            ("text/html", "text/plain", "application/json",
             "application/xml", "text/")):
        raise WebSearchError("unsupported content type: %s" % ctype)
    text_in = body.decode("utf-8", errors="replace")
    if "text/html" in (ctype or "") or "<html" in text_in[:2000].lower():
        parser = _TextExtractor()
        parser.feed(text_in)
        title = _html.unescape(parser.title.strip())
        text = _html.unescape(parser.text())
        links = parser.links
    else:
        title = ""
        text = text_in
        links = []
    text, truncated = _truncate(text, char_limit)
    return {"url": str(url), "title": title, "text": text,
            "links": links[:100], "truncated": truncated,
            "char_limit": char_limit}


# ------------------------------------------------------- codex-home glue

def ensure_codex_home_websearch(codex_home_dir, command_args=None,
                                script_dir=None):
    """Seed websearch MCP files + config into a tenant codex-home.

    Idempotent and non-destructive by design:
    - copies ``websearch.py`` and ``websearch_mcp.py`` only when the
      destination is missing or differs (so a tenant's own edits to
      the script are never silently overwritten -- differing content
      is left alone and reported);
    - appends the ``[mcp_servers.websearch]`` section only when that
      section header is absent; an existing section is never rewritten
      (keys are never written into config.toml by this helper).
    Returns {config_path, script_path, config_changed, script_changed}.
    ``command_args`` defaults to ``python3 <codex-home>/websearch_mcp.py``
    on the host; Docker tenants must pass the container path
    (``/var/lib/codex-app-server/codex-home/websearch_mcp.py``)
    instead, because the same config.toml cannot hold both paths.
    """
    import shutil

    home = os.path.abspath(codex_home_dir)
    os.makedirs(home, exist_ok=True)
    here = os.path.abspath(script_dir or os.path.dirname(__file__))
    script_changed = False
    script_differs = []
    for name in ("websearch.py", "websearch_mcp.py"):
        src = os.path.join(here, name)
        dst = os.path.join(home, name)
        if not os.path.exists(src):
            continue
        if not os.path.exists(dst):
            shutil.copyfile(src, dst)
            script_changed = True
        else:
            with open(src, "rb") as fh:
                src_bytes = fh.read()
            with open(dst, "rb") as fh:
                dst_bytes = fh.read()
            if src_bytes != dst_bytes:
                # Do not clobber a tenant-modified script; report it.
                script_differs.append(name)
                continue
    config_path = os.path.join(home, "config.toml")
    if command_args is None:
        command_args = ["python3",
                        os.path.join(home, "websearch_mcp.py")]
    section = ("\n[mcp_servers.websearch]\n"
               "command = %s\n"
               "args = [%s]\n"
               % (json.dumps(command_args[0]),
                  ", ".join(json.dumps(a) for a in command_args[1:])))
    config_changed = False
    existing = ""
    if os.path.exists(config_path):
        with open(config_path, encoding="utf-8") as fh:
            existing = fh.read()
    if "[mcp_servers.websearch]" not in existing:
        with open(config_path, "a", encoding="utf-8") as fh:
            if existing and not existing.endswith("\n"):
                fh.write("\n")
            fh.write(section)
        config_changed = True
    return {"config_path": config_path,
            "script_path": os.path.join(home, "websearch_mcp.py"),
            "config_changed": config_changed,
            "script_changed": script_changed,
            "script_differs": script_differs}


def ensure_websearch_mcp_env(codex_home_dir, env):
    """Write [mcp_servers.websearch.env] into a tenant config.toml.

    Why this exists (measured 2026-10-30, not guessed): Codex's
    MCP server subprocess does NOT inherit the Codex process env.
    A Docker task whose Codex bash env contained SEARXNG_URL and
    http_proxy still ran web_search with chain ["duckduckgo"]
    only, failing instantly with "Temporary failure in name
    resolution" -- SEARXNG_URL never reached the MCP process, and
    without a proxy the container has no direct DNS. Adding an
    explicit ``env`` table to the MCP config made the same task
    return real SearXNG results. Container env (docker -e) is
    therefore NOT sufficient for MCP tools; the env must live in
    config.toml, where Codex actually reads it for MCP servers.

    Idempotent: an existing env subsection is replaced in place,
    other sections are left untouched, and a missing websearch
    section is created first (command/args only, no env values
    invented). Values are written as TOML basic strings via
    json.dumps (safe for URLs and proxy URLs). Returns True when
    the file was written, False when env is empty or the file
    could not be written (caller logs, never crashes the task).
    """
    if not env:
        return False
    home = os.path.abspath(codex_home_dir)
    config_path = os.path.join(home, "config.toml")
    try:
        existing = ""
        if os.path.exists(config_path):
            with open(config_path, encoding="utf-8") as fh:
                existing = fh.read()
        if "[mcp_servers.websearch]" not in existing:
            existing += ("\n[mcp_servers.websearch]\n"
                         "command = \"python3\"\n"
                         "args = [\"/var/lib/codex-app-server/"
                         "codex-home/websearch_mcp.py\"]\n")
        env_lines = ["[mcp_servers.websearch.env]"]
        for key in sorted(env):
            env_lines.append("%s = %s" % (key, json.dumps(env[key])))
        env_block = "\n".join(env_lines) + "\n"
        lines = existing.splitlines(keepends=True)
        start = end = None
        for i, line in enumerate(lines):
            if line.strip() == "[mcp_servers.websearch.env]":
                start = i
            elif start is not None and line.startswith("["):
                end = i
                break
        if start is not None:
            if end is None:
                end = len(lines)
            lines[start:end] = [env_block]
            new_text = "".join(lines)
        else:
            insert_at = None
            in_section = False
            for i, line in enumerate(lines):
                if line.strip() == "[mcp_servers.websearch]":
                    in_section = True
                elif in_section and line.startswith("["):
                    insert_at = i
                    break
            if insert_at is None:
                insert_at = len(lines)
            lines.insert(insert_at, env_block)
            new_text = "".join(lines)
        if new_text == existing:
            return False
        with open(config_path, "w", encoding="utf-8") as fh:
            fh.write(new_text)
        return True
    except Exception:
        return False