"""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)")
    url = (base + "/search?"
           + urllib.parse.urlencode({"q": query, "format": "json"}))
    _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}