"""Execute the `fetch-page` tool: read one web page and return its text as Markdown.

`fetch-page` sits next to `web-search` and `calculator` on the website and in
the tool-use training data, so a model sees the same observation in both. HTML is
reduced to its main content by trafilatura, PDFs to their text, and plain
text types pass through. Each observation holds at most `PAGE_CHARS`
characters; a longer page ends with a note naming the `offset` that reads on.

Only public addresses are fetched. Every hostname is resolved before the
request, every address it resolves to must be globally routable, and the
connection goes to the vetted address itself, so a second DNS answer cannot
point it inside the network. Redirects are followed by hand and vetted the
same way.
"""

from __future__ import annotations

import ipaddress
import socket
import threading
import time
from collections import OrderedDict
from collections.abc import Mapping
from dataclasses import dataclass
from io import BytesIO
from typing import Any
from urllib.parse import urljoin, urlsplit, urlunsplit

import trafilatura
import urllib3
from bs4 import BeautifulSoup
from pypdf import PdfReader
from pypdf.errors import PdfReadError

FETCH_PAGE_TOOL_NAME = "fetch-page"

FETCH_PAGE_TOOL: dict[str, Any] = {
    "type": "function",
    "function": {
        "name": FETCH_PAGE_TOOL_NAME,
        "description": "Fetch a web page or PDF by URL and return its readable text as Markdown.",
        "parameters": {
            "type": "object",
            "properties": {
                "url": {
                    "type": "string",
                    "description": "The http or https URL to fetch.",
                },
                "offset": {
                    "type": "integer",
                    "description": "Character offset to start reading from, for pages longer than one result.",
                },
            },
            "required": ["url"],
        },
    },
}

PAGE_CHARS = 12_000
MAX_URL_CHARS = 2_048
MAX_BODY_BYTES = 8 * 1024 * 1024
MAX_REDIRECTS = 5
MAX_PDF_PAGES = 200
# Longer extracted text is cut here before caching; offsets past it read nothing.
MAX_DOCUMENT_CHARS = 400_000
CONNECT_SECONDS = 10.0
READ_SECONDS = 15.0
DEADLINE_SECONDS = 30.0
CACHE_ENTRIES = 128
CACHE_SECONDS = 15 * 60
USER_AGENT = (
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/131.0.0.0 Safari/537.36 KerzgrrAI-fetch/1.0 (+https://kerzgrr.ai)"
)
REQUEST_HEADERS = {
    "User-Agent": USER_AGENT,
    "Accept": "text/html,application/xhtml+xml,application/pdf;q=0.9,text/plain;q=0.8,*/*;q=0.5",
    "Accept-Language": "en-US,en;q=0.9",
    # urllib3 decodes these two without extra packages.
    "Accept-Encoding": "gzip, deflate",
}
REDIRECT_STATUSES = {301, 302, 303, 307, 308}
TEXT_TYPES = {
    "application/json",
    "application/ld+json",
    "application/xml",
    "application/rss+xml",
    "application/atom+xml",
    "application/javascript",
    "application/x-yaml",
    "application/yaml",
    "application/toml",
}
HTML_TYPES = {"text/html", "application/xhtml+xml"}


class FetchError(Exception):
    """A page that could not be fetched; the message is shown to the model."""


@dataclass(frozen=True)
class FetchOutcome:
    observation: str
    ok: bool
    url: str
    status: int | None = None
    content_type: str | None = None
    # Length of the whole extracted page, not of this observation.
    chars: int = 0

    def public(self) -> dict[str, Any]:
        return {"url": self.url, "status": self.status, "contentType": self.content_type, "chars": self.chars}


@dataclass(frozen=True)
class _Document:
    url: str
    title: str
    text: str
    status: int
    content_type: str


def is_fetch_page_tool_name(name: str) -> bool:
    return name.strip().lower().replace("_", "-") == FETCH_PAGE_TOOL_NAME


def url_from_arguments(arguments: Mapping[str, Any] | None) -> str:
    if not arguments:
        return ""
    raw = arguments.get("url")
    return raw.strip() if isinstance(raw, str) else ""


def offset_from_arguments(arguments: Mapping[str, Any] | None) -> int:
    if not arguments:
        return 0
    raw = arguments.get("offset")
    if isinstance(raw, bool):
        return 0
    if isinstance(raw, (int, float)):
        return max(0, int(raw))
    if isinstance(raw, str) and raw.strip().isdigit():
        return int(raw.strip())
    return 0


def fetch_page(url: str, *, offset: int = 0, page_chars: int = PAGE_CHARS) -> FetchOutcome:
    """One `fetch-page` observation: the page's text from `offset`, or an error line."""

    try:
        target = normalize_url(url)
    except FetchError as error:
        return FetchOutcome(f"Error: {error}", False, url.strip())
    try:
        document = _cache.get(target)
        if document is None:
            document = _load(target)
            _cache.put(target, document)
    except FetchError as error:
        return FetchOutcome(f"Error: {error}", False, target)
    except Exception as error:  # noqa: BLE001 — a page that breaks a parser is data, not a server fault
        return FetchOutcome(f"Error: could not read {target}: {type(error).__name__}", False, target)
    return FetchOutcome(
        render_observation(document, offset=max(0, int(offset)), page_chars=page_chars),
        True,
        document.url,
        document.status,
        document.content_type,
        len(document.text),
    )


def render_observation(document: _Document, *, offset: int, page_chars: int) -> str:
    total = len(document.text)
    lines = [f"URL: {document.url}"]
    if document.title:
        lines.append(f"Title: {document.title}")
    if total == 0:
        return "\n".join(lines) + "\n\nThe page has no readable text."
    if offset >= total:
        return "\n".join(lines) + f"\n\nThe offset {offset:,} is past the end of the page, which has {total:,} characters."
    end = min(total, offset + page_chars)
    if end < total:
        # Cut at a paragraph or line break near the limit rather than mid-word.
        window = document.text[offset:end]
        cut = max(window.rfind("\n\n"), window.rfind("\n"))
        if cut > page_chars // 2:
            end = offset + cut
    body = document.text[offset:end].strip("\n")
    lines.append("")
    lines.append(body)
    if end < total:
        lines.append("")
        lines.append(
            f"[Showing characters {offset:,}–{end:,} of {total:,}. "
            f"Call fetch-page with offset {end} to read on.]"
        )
    elif offset > 0:
        lines.append("")
        lines.append(f"[Showing characters {offset:,}–{end:,} of {total:,}, the end of the page.]")
    return "\n".join(lines)


def normalize_url(raw: str) -> str:
    text = raw.strip()
    if not text:
        raise FetchError("the URL is empty")
    if len(text) > MAX_URL_CHARS:
        raise FetchError(f"the URL is longer than {MAX_URL_CHARS:,} characters")
    if "://" not in text:
        text = f"https://{text}"
    parts = urlsplit(text)
    scheme = parts.scheme.lower()
    if scheme not in {"http", "https"}:
        raise FetchError(f"only http and https URLs can be fetched, not {parts.scheme}")
    if parts.username is not None or parts.password is not None:
        raise FetchError("URLs with a user name or password cannot be fetched")
    host = (parts.hostname or "").rstrip(".")
    if not host:
        raise FetchError("the URL has no host")
    try:
        ascii_host = host.encode("idna").decode("ascii") if not host.isascii() else host.lower()
    except UnicodeError as error:
        raise FetchError(f"the host {host!r} is not a valid domain name") from error
    try:
        port = parts.port
    except ValueError as error:
        raise FetchError("the URL has an invalid port") from error
    netloc = f"[{ascii_host}]" if ":" in ascii_host else ascii_host
    if port is not None and port != (443 if scheme == "https" else 80):
        netloc = f"{netloc}:{port}"
    return urlunsplit((scheme, netloc, parts.path or "/", parts.query, ""))


def public_addresses(host: str, port: int) -> list[str]:
    """Every address `host` resolves to; raises unless all of them are public."""

    try:
        infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM)
    except socket.gaierror as error:
        raise FetchError(f"could not resolve {host}") from error
    addresses: list[str] = []
    for info in infos:
        address = str(info[4][0]).split("%", 1)[0]
        if address not in addresses:
            addresses.append(address)
    if not addresses:
        raise FetchError(f"could not resolve {host}")
    for address in addresses:
        if not is_public_address(address):
            raise FetchError(f"{host} resolves to a private or reserved address, which cannot be fetched")
    return addresses


def is_public_address(address: str) -> bool:
    try:
        ip = ipaddress.ip_address(address)
    except ValueError:
        return False
    if isinstance(ip, ipaddress.IPv6Address):
        mapped = ip.ipv4_mapped or ip.sixtofour
        if mapped is not None:
            ip = mapped
        elif ip.teredo is not None:
            return False
    return ip.is_global and not ip.is_multicast and not ip.is_reserved and not ip.is_unspecified


def _load(url: str) -> _Document:
    deadline = time.monotonic() + DEADLINE_SECONDS
    current = url
    for _hop in range(MAX_REDIRECTS + 1):
        status, headers, body = _request(current, deadline)
        if status in REDIRECT_STATUSES:
            location = headers.get("location")
            if not location:
                raise FetchError(f"{current} redirected without a location")
            current = normalize_url(urljoin(current, location))
            continue
        if status >= 400:
            reason = _REASONS.get(status, "")
            raise FetchError(f"{current} returned HTTP {status}{f' {reason}' if reason else ''}")
        content_type = (headers.get("content-type") or "").split(";", 1)[0].strip().lower()
        title, text = _extract(body, content_type, headers.get("content-type") or "", current)
        return _Document(current, title, text[:MAX_DOCUMENT_CHARS], status, content_type or "unknown")
    raise FetchError(f"{url} redirected more than {MAX_REDIRECTS} times")


def _request(url: str, deadline: float) -> tuple[int, Mapping[str, str], bytes]:
    parts = urlsplit(url)
    host = parts.hostname or ""
    https = parts.scheme == "https"
    port = parts.port or (443 if https else 80)
    address = public_addresses(host, port)[0]
    remaining = deadline - time.monotonic()
    if remaining <= 0:
        raise FetchError(f"{url} took too long to answer")
    timeout = urllib3.Timeout(connect=min(CONNECT_SECONDS, remaining), read=min(READ_SECONDS, remaining))
    bracketed = f"[{host}]" if ":" in host else host
    host_header = bracketed if port == (443 if https else 80) else f"{bracketed}:{port}"
    pool: urllib3.HTTPConnectionPool
    if https:
        pool = urllib3.HTTPSConnectionPool(
            address,
            port=port,
            timeout=timeout,
            retries=False,
            maxsize=1,
            cert_reqs="CERT_REQUIRED",
            server_hostname=host,
            assert_hostname=host,
        )
    else:
        pool = urllib3.HTTPConnectionPool(address, port=port, timeout=timeout, retries=False, maxsize=1)
    path = urlunsplit(("", "", parts.path or "/", parts.query, ""))
    try:
        response = pool.urlopen(
            "GET",
            path,
            headers={**REQUEST_HEADERS, "Host": host_header},
            redirect=False,
            preload_content=False,
            decode_content=True,
        )
    except urllib3.exceptions.SSLError as error:
        pool.close()
        raise FetchError(f"{url} has an invalid TLS certificate") from error
    except (urllib3.exceptions.TimeoutError, socket.timeout) as error:
        pool.close()
        raise FetchError(f"{url} took too long to answer") from error
    except urllib3.exceptions.HTTPError as error:
        pool.close()
        raise FetchError(f"could not connect to {url}") from error
    try:
        headers = {key.lower(): value for key, value in response.headers.items()}
        if response.status in REDIRECT_STATUSES or response.status >= 400:
            return response.status, headers, b""
        chunks: list[bytes] = []
        size = 0
        while size < MAX_BODY_BYTES:
            if time.monotonic() > deadline:
                if not chunks:
                    raise FetchError(f"{url} took too long to answer")
                break
            chunk = response.read(64 * 1024)
            if not chunk:
                break
            chunks.append(chunk)
            size += len(chunk)
        return response.status, headers, b"".join(chunks)[:MAX_BODY_BYTES]
    except (urllib3.exceptions.TimeoutError, socket.timeout) as error:
        raise FetchError(f"{url} took too long to answer") from error
    except urllib3.exceptions.HTTPError as error:
        raise FetchError(f"the connection to {url} broke while reading") from error
    finally:
        response.release_conn()
        response.close()
        pool.close()


def _extract(body: bytes, content_type: str, raw_content_type: str, url: str) -> tuple[str, str]:
    if not content_type:
        content_type = _sniff(body)
    if content_type == "application/pdf":
        return _pdf_text(body, url)
    if content_type in HTML_TYPES:
        return _html_text(body, url)
    if content_type.startswith("text/") or content_type in TEXT_TYPES or content_type.endswith(("+json", "+xml")):
        return "", _decode(body, raw_content_type).strip()
    raise FetchError(f"{url} is {content_type}, which fetch-page cannot read")


def _sniff(body: bytes) -> str:
    head = body[:1024].lstrip()
    if head.startswith(b"%PDF"):
        return "application/pdf"
    lowered = head.lower()
    if lowered.startswith((b"<!doctype html", b"<html")) or b"<body" in lowered:
        return "text/html"
    try:
        head.decode("utf-8")
    except UnicodeDecodeError:
        return "application/octet-stream"
    return "text/plain"


def _decode(body: bytes, raw_content_type: str) -> str:
    charset = ""
    for part in raw_content_type.split(";")[1:]:
        name, _, value = part.partition("=")
        if name.strip().lower() == "charset":
            charset = value.strip().strip('"').strip("'")
    for encoding in (charset, "utf-8"):
        if not encoding:
            continue
        try:
            return body.decode(encoding)
        except (LookupError, UnicodeDecodeError):
            continue
    return body.decode("utf-8", errors="replace")


def _html_text(body: bytes, url: str) -> tuple[str, str]:
    metadata = trafilatura.extract_metadata(body, default_url=url)
    title = (metadata.title or "").strip() if metadata is not None else ""
    text = trafilatura.extract(
        body,
        url=url,
        output_format="markdown",
        include_comments=False,
        include_tables=True,
        # Inline links roughly double a page's length, and trafilatura garbles some of them.
        include_links=False,
        include_images=False,
        include_formatting=True,
        favor_recall=True,
        # No `deduplicate`: its cache spans calls, so a page fetched twice would come back emptied.
    )
    if text and text.strip():
        return title, _tidy(text)
    soup = BeautifulSoup(body, "lxml")
    if not title and soup.title and soup.title.string:
        title = soup.title.string.strip()
    for tag in soup(["script", "style", "noscript", "template", "svg", "nav", "footer", "header", "form"]):
        tag.decompose()
    lines = [line.strip() for line in soup.get_text("\n").splitlines()]
    return title, _tidy("\n".join(line for line in lines if line))


def _pdf_text(body: bytes, url: str) -> tuple[str, str]:
    try:
        reader = PdfReader(BytesIO(body))
        if reader.is_encrypted:
            try:
                reader.decrypt("")
            except Exception as error:  # noqa: BLE001 — pypdf raises several types for locked files
                raise FetchError(f"{url} is a password-protected PDF") from error
        title = ""
        if reader.metadata is not None and reader.metadata.title:
            title = str(reader.metadata.title).strip()
        pages: list[str] = []
        size = 0
        for page in reader.pages[:MAX_PDF_PAGES]:
            text = (page.extract_text() or "").strip()
            if text:
                pages.append(text)
                size += len(text)
            if size >= MAX_DOCUMENT_CHARS:
                break
    except FetchError:
        raise
    except (PdfReadError, ValueError, KeyError, TypeError) as error:
        raise FetchError(f"{url} is a PDF that could not be read") from error
    return title, _tidy("\n\n".join(pages))


_EDIT_MARKERS = {"[edit]", "\\[edit\\]"}


def _tidy(text: str) -> str:
    lines = [line.rstrip() for line in text.replace("\r\n", "\n").replace("\r", "\n").split("\n")]
    tidied: list[str] = []
    blank = 0
    for line in lines:
        if line.strip() in _EDIT_MARKERS:
            continue
        if line.strip():
            blank = 0
            tidied.append(line)
        else:
            blank += 1
            if blank == 1:
                tidied.append("")
    return "\n".join(tidied).strip()


class _DocumentCache:
    """Recently fetched pages, so reading on with `offset` does not fetch again."""

    def __init__(self, entries: int, seconds: float) -> None:
        self._entries = entries
        self._seconds = seconds
        self._items: OrderedDict[str, tuple[float, _Document]] = OrderedDict()
        self._lock = threading.Lock()

    def get(self, url: str) -> _Document | None:
        with self._lock:
            item = self._items.get(url)
            if item is None:
                return None
            if time.monotonic() - item[0] > self._seconds:
                del self._items[url]
                return None
            self._items.move_to_end(url)
            return item[1]

    def put(self, url: str, document: _Document) -> None:
        with self._lock:
            self._items[url] = (time.monotonic(), document)
            self._items.move_to_end(url)
            while len(self._items) > self._entries:
                self._items.popitem(last=False)

    def clear(self) -> None:
        with self._lock:
            self._items.clear()


_cache = _DocumentCache(CACHE_ENTRIES, CACHE_SECONDS)

_REASONS = {
    400: "Bad Request",
    401: "Unauthorized",
    403: "Forbidden",
    404: "Not Found",
    405: "Method Not Allowed",
    406: "Not Acceptable",
    408: "Request Timeout",
    410: "Gone",
    429: "Too Many Requests",
    451: "Unavailable For Legal Reasons",
    500: "Internal Server Error",
    502: "Bad Gateway",
    503: "Service Unavailable",
    504: "Gateway Timeout",
}