DealDocumentScreening/src/contract_check/bot/client.py

"""HTTP client to the contract-check api.

The bot touches the api exclusively over HTTP (hexagonal boundary). Every
request carries the service-token bearer and a per-interaction `X-Correlation-ID`
so api/worker logs line up under one id across the whole pipeline
(docs/ARCHITECTURE.md §13).
"""

from __future__ import annotations

from dataclasses import dataclass

import httpx

from ..core.logging import get_logger
from .config import BotSettings

log = get_logger(__name__)

DISCLAIMER_MARKER = "Дисклеймер"
FALLBACK_DISCLAIMER = (
    "\n\n> ⚠️ **Дисклеймер.** Это первичный скрининг, а не юридическая консультация. "
    "Находки могут содержать ошибки. Перед подписанием договор проверяет юрист."
)


def ensure_disclaimer(markdown: str) -> str:
    """Return markdown guaranteed to carry a disclaimer (append only if missing).

    The api's `render_markdown` already appends one; this is the defensive net
    required by docs/ARCHITECTURE.md §17 ("bot appends if api omitted").
    """
    if DISCLAIMER_MARKER in markdown:
        return markdown
    return markdown.rstrip() + FALLBACK_DISCLAIMER


class ApiError(Exception):
    """Unexpected api response."""

    def __init__(self, status_code: int, detail: str) -> None:
        super().__init__(f"api error {status_code}: {detail}")
        self.status_code = status_code
        self.detail = detail


class NoCreditsError(ApiError):
    """api returned 402 — the user has no credits left."""


class UnsupportedFormatError(ApiError):
    """api returned 400 — bad mime/size/extension."""


@dataclass(slots=True)
class UploadResult:
    document_id: str
    correlation_id: str
    credits_left: int


@dataclass(slots=True)
class ReportStatus:
    document_id: str
    status: str
    stage: str | None
    markdown: str | None = None
    filename: str | None = None


def _extract_detail(response: httpx.Response) -> str:
    try:
        body = response.json()
    except Exception:
        return response.text
    detail = body.get("detail") if isinstance(body, dict) else None
    return detail if isinstance(detail, str) else str(detail)


class ApiClient:
    """Thin httpx wrapper around the api endpoints the bot needs.

    The bot authenticates adapter-level calls (POST /auth/telegram/bot) with its
    service token. All user-scoped calls carry a user JWT obtained from that
    endpoint, not a raw telegram_id query parameter.
    """

    def __init__(self, settings: BotSettings) -> None:
        self._settings = settings
        self._client: httpx.AsyncClient | None = None
        self._token_cache: dict[int, str] = {}

    async def start(self) -> None:
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=self._settings.api_url,
                timeout=httpx.Timeout(30.0, connect=5.0),
            )

    async def aclose(self) -> None:
        if self._client is not None:
            await self._client.aclose()
            self._client = None

    @property
    def client(self) -> httpx.AsyncClient:
        if self._client is None:
            raise RuntimeError("ApiClient not started; call start() first")
        return self._client

    def _service_headers(self, correlation_id: str) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self._settings.bot_service_token}",
            "X-Correlation-ID": correlation_id,
        }

    def _user_headers(self, telegram_id: int, correlation_id: str) -> dict[str, str]:
        token = self._token_cache.get(telegram_id)
        if not token:
            raise RuntimeError("No JWT cached for telegram_id; call login() first")
        return {
            "Authorization": f"Bearer {token}",
            "X-Correlation-ID": correlation_id,
        }

    async def login(self, telegram_id: int, correlation_id: str) -> None:
        """Exchange a verified telegram_id for a user JWT and cache it."""
        r = await self.client.post(
            "/api/v1/auth/telegram/bot",
            json={"telegram_id": telegram_id},
            headers=self._service_headers(correlation_id),
        )
        if r.status_code != 200:
            raise ApiError(r.status_code, _extract_detail(r))
        body = r.json()
        token = body.get("access_token")
        if not isinstance(token, str):
            raise ApiError(500, "auth response missing access_token")
        self._token_cache[telegram_id] = token

    def forget(self, telegram_id: int) -> None:
        self._token_cache.pop(telegram_id, None)

    async def get_credits(self, telegram_id: int, correlation_id: str) -> int:
        r = await self.client.get(
            "/api/v1/me",
            headers=self._user_headers(telegram_id, correlation_id),
        )
        if r.status_code != 200:
            raise ApiError(r.status_code, _extract_detail(r))
        return int(r.json().get("credits_left", 0))

    async def upload_document(
        self,
        telegram_id: int,
        correlation_id: str,
        filename: str,
        data: bytes,
        content_type: str,
    ) -> UploadResult:
        files = {"file": (filename, data, content_type)}
        r = await self.client.post(
            "/api/v1/documents",
            files=files,
            headers=self._user_headers(telegram_id, correlation_id),
        )
        if r.status_code == 202:
            body = r.json()
            return UploadResult(
                document_id=body["document_id"],
                correlation_id=body["correlation_id"],
                credits_left=int(body.get("credits_left", 0)),
            )
        detail = _extract_detail(r)
        if r.status_code == 402:
            raise NoCreditsError(r.status_code, detail)
        if r.status_code == 400:
            raise UnsupportedFormatError(r.status_code, detail)
        raise ApiError(r.status_code, detail)

    async def get_report(
        self, telegram_id: int, correlation_id: str, document_id: str
    ) -> ReportStatus:
        r = await self.client.get(
            f"/api/v1/reports/{document_id}",
            headers=self._user_headers(telegram_id, correlation_id),
        )
        if r.status_code == 404:
            raise ApiError(404, "report not found")
        if r.status_code != 200:
            raise ApiError(r.status_code, _extract_detail(r))
        body = r.json()
        return ReportStatus(
            document_id=body["document_id"],
            status=body["status"],
            stage=body.get("stage"),
            markdown=body.get("markdown"),
            filename=body.get("filename"),
        )