path: root/src/cepis/client.py

"""The HTTP client, built on `urllib.request`.

`urllib` opens a fresh TCP+TLS connection per request and does not pool them.
For a one-off lookup that is irrelevant. For thousands, **do not write a
loop** — use :meth:`Client.lookup_many`, which resolves up to 500 CEPs per
request. A loop of 500 calls pays 500 TLS handshakes; the batch pays one.

Retry policy: 429 and the transient 5xx in :data:`RETRIABLE_STATUS`, plus
network failures, are retried with exponentially backed-off, jittered waits.
400 and 404 are deliberately absent from that set — both are definitive
answers about the request, so retrying returns the same thing more slowly
while burning someone else's quota.
"""

from __future__ import annotations

import csv
import io
import json
import random
import re
import time
import urllib.error
import urllib.parse
import urllib.request
from email.utils import parsedate_to_datetime
from typing import Any

from . import __version__
from .errors import (
    CEPNotFound,
    InvalidCEP,
    InvalidResponse,
    NetworkError,
    RateLimited,
    ServerError,
)
from .models import Address, Faixa, Health

DEFAULT_URL = "https://cep.is"

MAX_BATCH = 500

RETRIABLE_STATUS = frozenset({429, 500, 502, 503, 504})

_DIGITS_ONLY = re.compile(r"^\d{8}$")


def normalize_cep(value: str) -> str:
    """Return the CEP's 8 digits, accepting the mask.

    `"01001-000"`, `"01001000"` and `" 01001-000 "` all become `"01001000"`.

    Raises :class:`InvalidCEP` unless exactly 8 digits remain. This runs
    **before** any request: a 7-digit CEP is the caller's bug, and spending a
    round trip for the server to say so wastes both sides' time.
    """
    if not isinstance(value, str):
        raise InvalidCEP(f"CEP must be a str, got {type(value).__name__}")
    cleaned = value.strip().replace("-", "").replace(".", "").replace(" ", "")
    if not _DIGITS_ONLY.match(cleaned):
        raise InvalidCEP(f"CEP must be 8 digits, got {value!r}")
    return cleaned


def _retry_after_seconds(value: str | None) -> float | None:
    """Parse the `Retry-After` header in both RFC 9110 forms.

    It may arrive as a number of seconds (`"120"`) or as an HTTP date
    (`"Wed, 21 Oct 2026 07:28:00 GMT"`). Handling only the first is the common
    bug, and it makes a client ignore exactly the wait request that matters
    most.
    """
    if not value:
        return None
    value = value.strip()
    try:
        return max(0.0, float(value))
    except ValueError:
        pass
    try:
        target = parsedate_to_datetime(value)
    except (TypeError, ValueError):
        return None
    if target is None:
        return None
    return max(0.0, target.timestamp() - time.time())


class Client:
    """Client for the cep.is API.

    Reuse one instance: the configuration is immutable and no method keeps
    state between calls, so it is safe to share across threads.

    :param base_url:
        Where the API lives. Point it at a local instance (the API repo brings
        one up on `http://localhost:8080` with `just up`) or at a test server.
        It is the only parameter that usually changes between environments. A
        trailing slash is stripped, so it cannot produce doubled separators
        that some proxies normalise and others 404.
    :param timeout:
        Seconds to wait **per socket operation** (connect, and each read) —
        this is not a cap on the whole call. That cap is `total_timeout`.
    :param max_attempts:
        How many attempts in total, counting the first. `1` disables retrying.
        The default of `3` covers the common case (one transient failure)
        without turning a real outage into a long wait.
    :param backoff_base:
        Base of the exponential wait, in seconds. The actual wait is drawn
        from `[0, base * 2**(n-1)]` — *full jitter*. The randomness is not
        decoration: without it, N clients that fail together come back
        together and reproduce the spike that took the service down.
    :param backoff_max:
        Cap on each individual wait, in seconds. It applies to a server-sent
        `Retry-After` too, or a `Retry-After: 3600` from a misconfigured proxy
        would park the thread for an hour.
    :param total_timeout:
        Cap on the whole call, attempts and waits included. `None` leaves
        `max_attempts` as the only bound. Set it in code serving a web
        request: without it the worst case is `max_attempts * timeout` plus
        the sum of the backoffs.
    :param user_agent:
        Overrides the `User-Agent`. The default identifies the library and its
        version, which is what lets the server diagnose an old client
        misbehaving.
    """

    def __init__(
        self,
        base_url: str = DEFAULT_URL,
        *,
        timeout: float = 10.0,
        max_attempts: int = 3,
        backoff_base: float = 0.5,
        backoff_max: float = 8.0,
        total_timeout: float | None = None,
        user_agent: str | None = None,
    ) -> None:
        if max_attempts < 1:
            raise ValueError("max_attempts must be >= 1")
        if timeout <= 0:
            raise ValueError("timeout must be > 0")

        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.max_attempts = max_attempts
        self.backoff_base = backoff_base
        self.backoff_max = backoff_max
        self.total_timeout = total_timeout
        self.user_agent = user_agent or f"cepis-python/{__version__} (+https://cep.is)"

    def lookup(self, cep: str) -> Address:
        """Look up a single CEP.

        >>> client.lookup("01001-000").localidade
        'São Paulo'

        A 400 from the API is surfaced as :class:`InvalidCEP` carrying the
        API's own message: the format already passed local validation, so the
        server knows something this client does not.

        A 404 may carry the municipality the CEP would have belonged to, in
        `CEPNotFound.faixa` (see :class:`cepis.Faixa`) — still a 404, still an
        error, but usually enough to ask a better question than "invalid CEP"::

            try:
                address = client.lookup(cep)
            except CEPNotFound as exc:
                if exc.faixa:
                    print(f"no such CEP in {exc.faixa.label}")

        :raises InvalidCEP: the CEP is not 8 digits (no request is made).
        :raises CEPNotFound: well-formed, but not in the database.
        :raises RateLimited: 429 after every attempt.
        :raises ServerError: 5xx after every attempt.
        :raises NetworkError: the request never completed.
        """
        number = normalize_cep(cep)
        status, body, _ = self._request(f"/{number}")

        if status == 404:
            faixa = Faixa.from_json(self._error_json(body).get("faixa"))
            message = f"CEP {number} does not exist in the DNE database"
            if faixa is not None:
                message = f"{message} (it would be in {faixa.label})"
            raise CEPNotFound(message, faixa=faixa)
        if status == 400:
            raise InvalidCEP(self._error_message(body, f"CEP {number} rejected"))

        return Address.from_json(self._json(body))

    def lookup_many(self, ceps: list[str]) -> list[Address | None]:
        """Look up many CEPs in one batch.

        Returns a list of **the same length and in the same order** as the
        input. A CEP that does not exist becomes `None` at its position — it
        does not vanish from the list, and it does not raise. That is the right
        difference from :meth:`lookup`: in a batch of a thousand, one bad CEP
        must not cost the other nine hundred and ninety-nine.

        A malformed CEP also becomes `None`, for the same reason, rather than
        raising `InvalidCEP`. If you need to know *which* one was malformed,
        validate first with :func:`normalize_cep`.

        Input is chunked into blocks of 500 (the API's limit) automatically.

        .. warning::
           The batch carries **no coordinate**. The endpoint behind it is CSV,
           and the CSV has no `geo` column. Every `Address` returned here has
           `geo=None`, including for CEPs that do have one. Need `geo`? Use
           :meth:`lookup`, one at a time.
        """
        if not ceps:
            return []

        result: list[Address | None] = []
        for start in range(0, len(ceps), MAX_BATCH):
            result.extend(self._lookup_chunk(ceps[start : start + MAX_BATCH]))
        return result

    def health(self) -> Health:
        """Query `GET /health`: does the API answer, and with which database.

        Does not raise for an API reporting itself degraded — that arrives in
        `Health.status`, and checking health is exactly when you want the
        diagnosis rather than an exception. Network failures still raise.
        """
        _, body, _ = self._request("/health", accepted_status=(200, 503))
        return Health.from_json(self._json(body))

    def _lookup_chunk(self, ceps: list[str]) -> list[Address | None]:
        """One chunk of at most `MAX_BATCH` CEPs, in a single request.

        A blank CSV row means "not found": the API fills only the `cep` column
        and leaves the rest empty, precisely so it does not misalign whoever is
        consuming the batch in a spreadsheet.

        The API guarantees one row per requested CEP. Should that ever fail to
        hold — a proxy truncating, a cut response — the result is padded to the
        input length rather than returned short, because a shifted list
        silently pairs the wrong address with the wrong CEP.
        """
        query = urllib.parse.urlencode({"ceps": ",".join(str(c) for c in ceps)})
        _, body, _ = self._request(f"/csv?{query}")

        text = body.decode("utf-8-sig", errors="replace")
        reader = csv.DictReader(io.StringIO(text))

        rows: list[Address | None] = []
        for row in reader:
            if not (row.get("localidade") or "").strip():
                rows.append(None)
                continue
            rows.append(Address.from_json({k: (v or "") for k, v in row.items()}))

        if len(rows) != len(ceps):
            rows = (rows + [None] * len(ceps))[: len(ceps)]
        return rows

    def _json(self, body: bytes) -> dict[str, Any]:
        """Decode a JSON object body, or raise :class:`InvalidResponse`."""
        try:
            data = json.loads(body.decode("utf-8"))
        except (ValueError, UnicodeDecodeError) as exc:
            raise InvalidResponse(
                "response is not valid JSON — something between you and the "
                f"API is answering instead of it: {body[:120]!r}"
            ) from exc
        if not isinstance(data, dict):
            raise InvalidResponse(f"expected a JSON object, got {type(data).__name__}")
        return data

    @staticmethod
    def _error_json(body: bytes) -> dict[str, Any]:
        """Decode an error body, or return `{}`.

        Unlike :meth:`_json` this never raises. An error response is already
        the bad news; whether its body could be parsed only decides how much
        detail the exception carries, never whether one is raised.
        """
        try:
            data = json.loads(body.decode("utf-8"))
        except (ValueError, UnicodeDecodeError):
            return {}
        return data if isinstance(data, dict) else {}

    def _error_message(self, body: bytes, fallback: str) -> str:
        """Pull `mensagem` out of an error body, falling back to a default."""
        message = self._error_json(body).get("mensagem")
        return str(message) if message else fallback

    def _backoff(self, attempt: int, retry_after: float | None) -> float:
        """How long to wait before the next attempt.

        The server's `Retry-After` wins over the local calculation when
        present: it knows when it will be back, we are guessing. It still goes
        through the `backoff_max` cap. Otherwise the wait is drawn from
        `[0, ceiling]` rather than being the ceiling — see `backoff_base` on
        why that randomness matters.
        """
        if retry_after is not None:
            return min(retry_after, self.backoff_max)
        ceiling = min(self.backoff_base * (2 ** (attempt - 1)), self.backoff_max)
        return random.uniform(0, ceiling)

    def _request(
        self,
        path: str,
        *,
        accepted_status: tuple[int, ...] = (200,),
    ) -> tuple[int, bytes, Any]:
        """Perform the GET, with retries. Returns `(status, body, headers)`.

        *Definitive* error statuses (400, 404) are returned normally, for the
        caller to turn into the right exception — only transient failures
        become exceptions here, and only once the attempts are exhausted.

        The two ``except`` clauses below must stay in this order.
        ``HTTPError`` subclasses ``URLError``, so catching ``URLError`` first
        would classify every 404 as a network failure and retry it.
        """
        url = f"{self.base_url}{path}"
        deadline = None if self.total_timeout is None else time.monotonic() + self.total_timeout

        last_status: int | None = None
        last_retry_after: float | None = None
        last_network_error: BaseException | None = None

        for attempt in range(1, self.max_attempts + 1):
            if deadline is not None and time.monotonic() >= deadline:
                break

            request = urllib.request.Request(
                url,
                method="GET",
                headers={
                    "User-Agent": self.user_agent,
                    "Accept": "application/json, text/csv;q=0.9, */*;q=0.1",
                    "Accept-Encoding": "identity",
                },
            )

            try:
                with urllib.request.urlopen(request, timeout=self.timeout) as response:
                    return response.status, response.read(), response.headers

            except urllib.error.HTTPError as exc:
                body = exc.read()
                status = exc.code
                if status in accepted_status or status not in RETRIABLE_STATUS:
                    return status, body, exc.headers
                last_status = status
                last_retry_after = _retry_after_seconds(
                    exc.headers.get("Retry-After") if exc.headers else None
                )
                last_network_error = None

            except (urllib.error.URLError, TimeoutError, OSError) as exc:
                last_network_error = exc
                last_status = None
                last_retry_after = None

            if attempt == self.max_attempts:
                break

            wait = self._backoff(attempt, last_retry_after)
            if deadline is not None:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    break
                wait = min(wait, remaining)
            time.sleep(wait)

        if last_status == 429:
            raise RateLimited(
                f"rate limit reached at {url} (after {self.max_attempts} attempt(s))",
                retry_after=last_retry_after,
            )
        if last_status is not None:
            raise ServerError(
                f"the API answered {last_status} at {url} (after {self.max_attempts} attempt(s))",
                status=last_status,
            )
        raise NetworkError(
            f"could not reach {url} (after {self.max_attempts} attempt(s)): {last_network_error}"
        ) from last_network_error