ForcePilot/backend/package/yuxi/channel/extensions/flock/utils.py

from __future__ import annotations

import asyncio
import logging
import random
import re

import httpx

from .constants import (
    API_TIMEOUT,
    BASE_URL,
    BOT_TOKEN_PATTERN,
    FLOCK_GROUP_ID_PATTERN,
    FLOCK_USER_ID_PATTERN,
    MAX_RETRIES,
    RATE_LIMIT_RETRY_AFTER_DEFAULT,
    RETRY_BASE_DELAY,
    RETRY_MAX_DELAY,
)

logger = logging.getLogger(__name__)

_bot_token_re = re.compile(BOT_TOKEN_PATTERN)
_user_id_re = re.compile(FLOCK_USER_ID_PATTERN)
_group_id_re = re.compile(FLOCK_GROUP_ID_PATTERN)


class FlockAPIError(Exception):
    def __init__(self, error: str, description: str = "", retryable: bool = False):
        super().__init__(f"Flock API Error: {error} - {description}")
        self.error = error
        self.description = description
        self.retryable = retryable


def create_http_client(timeout: float = API_TIMEOUT) -> httpx.AsyncClient:
    return httpx.AsyncClient(timeout=timeout)


def build_auth_headers(bot_token: str) -> dict[str, str]:
    return {
        "Authorization": f"Bearer {bot_token}",
        "Content-Type": "application/json",
    }


def validate_bot_token(token: str) -> bool:
    return bool(_bot_token_re.match(token))


def validate_flock_user_id(user_id: str) -> bool:
    return bool(_user_id_re.match(user_id))


def validate_flock_group_id(group_id: str) -> bool:
    return bool(_group_id_re.match(group_id))


def strip_flock_prefix(identifier: str) -> str:
    for prefix in ("u:", "g:", "t:", "b:"):
        if identifier.startswith(prefix):
            return identifier[len(prefix) :]
    return identifier


async def flock_api_call(
    client: httpx.AsyncClient,
    endpoint: str,
    token: str,
    payload: dict,
    max_retries: int = MAX_RETRIES,
) -> dict:
    url = f"{BASE_URL}{endpoint}"
    headers = build_auth_headers(token)

    for attempt in range(max_retries + 1):
        try:
            resp = await client.post(url, json=payload, headers=headers)

            if resp.status_code == 200:
                return resp.json()

            if resp.status_code == 429:
                retry_after = int(resp.headers.get("Retry-After", RATE_LIMIT_RETRY_AFTER_DEFAULT))
                if attempt < max_retries:
                    logger.warning(
                        "Flock API rate limited (429) for %s, retrying after %ds (attempt %d/%d)",
                        endpoint,
                        retry_after,
                        attempt + 1,
                        max_retries,
                    )
                    await asyncio.sleep(retry_after)
                    continue
                raise FlockAPIError("RateLimitExceeded", "Rate limit exceeded", retryable=True)

            if resp.status_code >= 500:
                if attempt < max_retries:
                    wait = RETRY_BASE_DELAY * (2**attempt) + random.uniform(0, 1)
                    wait = min(wait, RETRY_MAX_DELAY)
                    logger.warning(
                        "Flock API server error (%d) for %s, retrying in %.1fs (attempt %d/%d)",
                        resp.status_code,
                        endpoint,
                        wait,
                        attempt + 1,
                        max_retries,
                    )
                    await asyncio.sleep(wait)
                    continue
                error_data = _parse_error_body(resp)
                raise FlockAPIError(
                    error_data.get("error", "InternalError"),
                    error_data.get("errorDescription", str(resp.status_code)),
                    retryable=True,
                )

            error_data = _parse_error_body(resp)
            raise FlockAPIError(
                error_data.get("error", f"HTTP {resp.status_code}"),
                error_data.get("errorDescription", ""),
                retryable=False,
            )

        except httpx.TimeoutException:
            if attempt < max_retries:
                wait = RETRY_BASE_DELAY * (2**attempt)
                logger.warning(
                    "Flock API timeout for %s, retrying in %.1fs (attempt %d/%d)",
                    endpoint,
                    wait,
                    attempt + 1,
                    max_retries,
                )
                await asyncio.sleep(wait)
                continue
            raise FlockAPIError("Timeout", f"Request to {endpoint} timed out", retryable=True)

    raise FlockAPIError("MaxRetriesExceeded", f"Failed after {max_retries} retries", retryable=True)


def _parse_error_body(resp: httpx.Response) -> dict:
    try:
        return resp.json()
    except Exception:
        return {"error": f"HTTP {resp.status_code}", "errorDescription": resp.text[:200]}


def classify_api_error(resp: httpx.Response) -> tuple[str, str, bool]:
    error_code, error_description = "UnknownError", ""
    retryable = False
    try:
        body = resp.json()
        error_code = body.get("error", f"HTTP {resp.status_code}")
        error_description = body.get("errorDescription", "")
    except Exception:
        error_code = f"HTTP {resp.status_code}"

    if resp.status_code == 429 or resp.status_code >= 500:
        retryable = True

    return error_code, error_description, retryable


def is_retryable_http_status(status_code: int) -> bool:
    return status_code == 429 or status_code >= 500