"""The one way out of this container.

Every outbound request in the suite goes through here, which means every
outbound request goes through squid. There is no fallback path and there
should never be one - the isolation story depends on this process having
no other route.
"""

from __future__ import annotations

import logging
import re
import socket
import ssl

import httpx

from .settings import settings

log = logging.getLogger(__name__)

_client: httpx.AsyncClient | None = None
_local_client: httpx.AsyncClient | None = None

# Generous enough for a slow origin, short enough that a wedged fetch does
# not sit past Open WebUI's own tool timeout and lose the whole turn.
DEFAULT_TIMEOUT = httpx.Timeout(connect=10.0, read=30.0, write=10.0, pool=10.0)

# Hard byte ceiling on any download. A datasheet PDF is legitimately a few
# MB; nothing this assistant needs is 200MB, and streaming one into memory
# in a 1GB container is how you get an OOM kill instead of an answer.
MAX_DOWNLOAD_BYTES = 12 * 1024 * 1024

# A vendor reference manual is legitimately larger than that: ST's RM0008 is
# 13MB, and the 12MB ceiling was silently truncating it. The truncated bytes
# then failed PDF extraction, and the model was told the document "may be a
# scanned image, a video, or a page that renders entirely in JavaScript" -
# false in every particular. It went to forum posts instead and cited them as
# the manual. 40MB of PDF text extraction fits the container; the real fix is
# that going over now says so.
MAX_PDF_BYTES = 40 * 1024 * 1024

_PROXY_STATUS = re.compile(r"\b([45]\d\d)\b")


def get_client() -> httpx.AsyncClient:
    global _client
    if _client is None:
        _client = httpx.AsyncClient(
            proxy=settings.egress_proxy,
            timeout=DEFAULT_TIMEOUT,
            follow_redirects=True,
            max_redirects=5,
            headers={
                "User-Agent": settings.user_agent,
                "Accept-Encoding": "gzip, deflate",
            },
            limits=httpx.Limits(max_connections=24, max_keepalive_connections=8),
        )
    return _client


def get_local_client() -> httpx.AsyncClient:
    """Client for services inside our own namespace - searxng, and nothing else.

    These must NOT go through squid. squid's job is to deny private
    destinations, and searxng on loopback is exactly that: routing its
    requests through the proxy got them refused with a 403, which surfaced
    to the model as "this address is on a private network".

    Deliberately a separate client rather than a per-request override.
    httpx has no supported way to bypass a client's proxy for one call
    (`extensions={"proxy": None}` is silently ignored), so a single client
    would send every internal request through the proxy again the moment
    someone forgot.
    """
    global _local_client
    if _local_client is None:
        _local_client = httpx.AsyncClient(
            timeout=httpx.Timeout(connect=5.0, read=30.0, write=10.0, pool=5.0),
            follow_redirects=False,
            headers={"User-Agent": "owuitools-internal/1.0"},
        )
    return _local_client


async def close_client() -> None:
    global _client, _local_client
    if _client is not None:
        await _client.aclose()
        _client = None
    if _local_client is not None:
        await _local_client.aclose()
        _local_client = None


class FetchError(RuntimeError):
    """Outbound failure, phrased for the model rather than for a log."""


def _squid_error(exc: Exception) -> str | None:
    """The X-Squid-Error token, when squid answered instead of the origin.

    Squid stamps its own error pages with this header. Matching on the status
    code alone - which this used to do - reported an origin refusing a bot as
    an infrastructure block, and the model then told the user the source was
    unreachable and stopped trying. That is worse than the bare 403, because
    it sounds authoritative and it is wrong.
    """
    response = getattr(exc, "response", None)
    if response is None:
        return None
    for key, value in response.headers.items():
        if key.lower() == "x-squid-error":
            parts = value.split()
            return parts[0] if parts else value
    return None


def _proxy_status(exc: Exception) -> str | None:
    """The status squid gave a CONNECT, dug out of the exception text.

    httpx raises ProxyError for any non-200 CONNECT and puts the status line in
    the message. That status is the only thing separating "this name does not
    resolve" (503) from "policy refused this address" (403), and the two need
    opposite advice. Collapsing them cost a graded run: the model invented
    three plausible API hostnames, none of which existed, and was told each
    time that the address was on a private network - so it never suspected the
    hostname, and reported to the user that public services were blocked here.
    """
    if not isinstance(exc, httpx.ProxyError):
        return None
    found = _PROXY_STATUS.search(str(exc))
    return found.group(1) if found else None


def _host_of(target: str) -> str:
    try:
        return httpx.URL(target).host or target
    except Exception:
        return target


def explain(exc: Exception, target: str) -> FetchError:
    """Turn a transport exception into something the model can act on.

    Error bodies from a tool server are fed straight back to the model, so
    they are part of the prompt whether you treat them that way or not. What
    the message blames matters as much as the fact that it failed: name the
    wrong cause and the model adopts the wrong remedy, confidently.
    """
    token = (_squid_error(exc) or "").upper()
    status = _proxy_status(exc)
    if token or status:
        code = status or str(getattr(getattr(exc, "response", None), "status_code", ""))
        host = _host_of(target)
        if token == "ERR_ACCESS_DENIED" or code == "403":
            return FetchError(
                f"Access to {target} was refused by the egress proxy: the name resolves to a "
                "private, loopback or otherwise reserved address, or the port is not 80 or 443. "
                "This assistant reaches only the public internet. Tell the user you have no "
                "access to that host instead of guessing its contents."
            )
        if token == "ERR_DNS_FAIL" or code == "503":
            return FetchError(
                f"{host} does not resolve - there is no such host in DNS. The hostname is wrong, "
                "or it was guessed. Do not try variants of it, and do not report this as the host "
                "being blocked: search for the real address of the service you want, or tell the "
                "user that host does not exist."
            )
        if token in {"ERR_CONNECT_FAIL", "ERR_READ_TIMEOUT", "ERR_ZERO_SIZE_OBJECT"} \
                or code in {"502", "504"}:
            return FetchError(
                f"{host} resolves, but the connection failed - the server is down, refusing "
                "connections, or timing out. This is not a restriction on this assistant. Try "
                "again later, or use a different source."
            )
        return FetchError(
            f"{target} could not be reached through the egress proxy ({token or code}). Either the "
            "name does not resolve or the host is not accepting connections; rewording the URL "
            "will not change that."
        )
    if isinstance(exc, httpx.ConnectTimeout | httpx.ReadTimeout):
        return FetchError(
            f"{target} did not send data in time. For a large PDF from a slow server one retry is "
            "reasonable; a third is not. If it fails again, look for a smaller document - a single "
            "chapter, an application note, or an HTML version - rather than retrying this one."
        )
    if isinstance(exc, httpx.TooManyRedirects):
        return FetchError(f"{target} redirected too many times; the page may be behind a login.")
    if isinstance(exc, httpx.HTTPStatusError):
        code = exc.response.status_code
        if code == 401:
            return FetchError(f"{target} returned 401 - the content is behind a login. Find another source.")
        if code == 403:
            return FetchError(
                f"{target} returned 403. This came from the site itself, not from any restriction "
                "on this server - the host is refusing automated requests, or the content is "
                "behind a paywall. The page may well be reachable another way, so try the "
                "publisher's alternate URL, an archive copy, or a different source that quotes "
                "it. Do not report this as the assistant being unable to reach the internet."
            )
        if code == 404:
            return FetchError(f"{target} does not exist (404). Check the URL or search for the correct one.")
        if code == 429:
            return FetchError(f"{target} is rate-limiting this server (429). Use a different source rather than retrying.")
        return FetchError(f"{target} returned HTTP {code}.")
    return FetchError(f"Could not reach {target}: {type(exc).__name__}.")


async def get_bytes(url: str, **kwargs) -> tuple[bytes, str, str]:
    """Fetch a URL, returning (body, content_type, final_url).

    Streams so an oversized response is abandoned mid-flight rather than
    buffered and then rejected.
    """
    client = get_client()
    async with client.stream("GET", url, **kwargs) as response:
        response.raise_for_status()
        content_type = response.headers.get("content-type", "").split(";")[0].strip().lower()

        # PDFs get the larger allowance. A reference manual genuinely is tens
        # of megabytes, and the alternative to reading it is the model quoting
        # a forum post as though it were the manual.
        is_pdf = "pdf" in content_type or url.lower().split("?")[0].endswith(".pdf")
        limit = MAX_PDF_BYTES if is_pdf else MAX_DOWNLOAD_BYTES

        chunks: list[bytes] = []
        total = 0
        async for chunk in response.aiter_bytes():
            total += len(chunk)
            if total > limit:
                # This used to `break` and hand back the partial bytes. They
                # then failed to parse, and the parse failure was reported
                # with three suggested causes that were all untrue. Refusing
                # out loud is worse for one fetch and far better for the
                # answer.
                log.warning("aborting oversized download from %s at %d bytes", url, total)
                raise FetchError(
                    f"{url} is larger than the {limit // (1024 * 1024)} MB this assistant can "
                    f"download (stopped after {total // (1024 * 1024)} MB). None of it was read, "
                    "so nothing in it is available - do not describe its contents from memory or "
                    "from someone's summary of it. If you need one section of a large manual, "
                    "look for that section as its own document: search for its title, or for an "
                    "application note covering it, or ask the user to paste the page."
                )
            chunks.append(chunk)

        return b"".join(chunks), content_type, str(response.url)


async def get_json(url: str, **kwargs):
    client = get_client()
    response = await client.get(url, **kwargs)
    response.raise_for_status()
    return response.json()


def tls_socket(host: str, port: int = 443, *, verify: bool, timeout: float = 12.0) -> tuple[ssl.SSLSocket, socket.socket]:
    """Open a TLS connection through squid's CONNECT tunnel.

    httpx will not surface a peer certificate chain, so inspect_tls needs a
    real socket. Doing the CONNECT by hand keeps that traffic on the same
    single egress path as everything else rather than opening a second one.

    Caller closes both returned sockets.
    """
    proxy = httpx.URL(settings.egress_proxy)
    raw = socket.create_connection((proxy.host, proxy.port or 3128), timeout=timeout)
    try:
        raw.sendall(f"CONNECT {host}:{port} HTTP/1.1\r\nHost: {host}:{port}\r\n\r\n".encode())

        # Read just the tunnel response header, not a byte of the TLS record
        # that follows it.
        header = b""
        while b"\r\n\r\n" not in header:
            chunk = raw.recv(1)
            if not chunk:
                raise FetchError(f"The egress proxy closed the connection while opening a tunnel to {host}.")
            header += chunk
            if len(header) > 8192:
                raise FetchError("The egress proxy returned an oversized response.")

        status_line = header.split(b"\r\n", 1)[0].decode("latin-1")
        if " 200 " not in status_line:
            # Same discrimination as explain(): squid says 503 when it could
            # not resolve the name and 403 when policy refused the address,
            # and telling a model its guessed hostname is "on a private
            # network" stops it ever fixing the name.
            found = _PROXY_STATUS.search(status_line)
            code = found.group(1) if found else ""
            if code == "503":
                raise FetchError(
                    f"{host} does not resolve - there is no such host in DNS. The hostname is "
                    "wrong, or it was guessed; it is not being blocked."
                )
            if code in {"502", "504"}:
                raise FetchError(
                    f"{host} resolves, but the connection to port {port} failed - the server is "
                    "down, refusing connections, or timing out."
                )
            raise FetchError(
                f"The egress proxy refused a connection to {host}:{port} ({status_line.strip()}). "
                "That address is private, reserved, or on a port this assistant cannot reach."
            )

        if verify:
            context = ssl.create_default_context()
        else:
            # Second pass: we still want to *describe* an expired or
            # self-signed certificate rather than refuse to look at it.
            context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
            context.check_hostname = False
            context.verify_mode = ssl.CERT_NONE

        wrapped = context.wrap_socket(raw, server_hostname=host)
        return wrapped, raw
    except Exception:
        raw.close()
        raise