api_client.py

"""HTTP client for HIRO Translation API (public connect gateway or internal stage)."""

from __future__ import annotations

import json
import os
from collections.abc import Iterator
from typing import Any

import requests

# Public gateway (matches curl examples against connect.zhihuiya.com).
DEFAULT_BASE = "https://connect.zhihuiya.com/hiro_translation"

API_KEY_ENV = "HIRO_API_KEY"


def lang_to_codes(lang: str) -> tuple[str, str]:
    src, _, tgt = lang.partition("2")
    if not src or not tgt:
        raise ValueError(f"invalid lang {lang!r}; expected format like zh2en")
    return src, tgt


def translate_payload(
    content: str,
    lang: str,
    *,
    mode: str,
) -> dict[str, Any]:
    source, target = lang_to_codes(lang)
    return {
        "content": content,
        "sourceLanguageCode": source,
        "targetLanguageCode": target,
        "mode": mode,
    }


def _api_url(base_url: str, path: str) -> str:
    return f"{base_url.rstrip('/')}{path}"


def _request_headers(api_key: str | None = None) -> dict[str, str]:
    headers = {"Content-Type": "application/json"}
    key = (api_key if api_key is not None else os.environ.get(API_KEY_ENV, "")).strip()
    if key:
        headers["Authorization"] = f"Bearer {key}"
    return headers


def _translated_text(data: dict[str, Any]) -> str:
    value = data.get("textTranslated", data.get("text_translated", ""))
    return value if isinstance(value, str) else ""


def _original_text(data: dict[str, Any], fallback: str = "") -> str:
    value = data.get("textOriginal", data.get("text_original", fallback))
    return value if isinstance(value, str) else fallback


def health_ok(
    base_url: str = DEFAULT_BASE,
    *,
    api_key: str | None = None,
) -> tuple[bool, str]:
    try:
        r = requests.get(
            _api_url(base_url, "/health"),
            headers=_request_headers(api_key),
            timeout=12,
        )
        r.raise_for_status()
        data = r.json()
        if data.get("status") != "OK":
            return False, f"Unexpected response: {r.text[:200]}"
        upstream = data.get("upstream", "UNKNOWN")
        if upstream == "OK":
            return True, "Healthy · upstream OK"
        return False, f"Gateway up, upstream unavailable ({upstream})"
    except requests.RequestException as exc:
        return False, str(exc)


def translate_fast(
    text: str,
    lang: str,
    *,
    base_url: str = DEFAULT_BASE,
    api_key: str | None = None,
    timeout: int = 1800,
) -> dict[str, Any]:
    r = requests.post(
        _api_url(base_url, "/translate"),
        json=translate_payload(text, lang, mode="fast"),
        headers=_request_headers(api_key),
        timeout=timeout,
    )
    if r.status_code >= 400:
        try:
            err = r.json().get("error", r.text[:500])
        except json.JSONDecodeError:
            err = r.text[:500]
        raise RuntimeError(err)
    data = r.json()
    if data.get("state") not in (None, "success") and "error" in data:
        raise RuntimeError(str(data.get("error", data)))
    # Normalize to keys the UI already understands (snake_case aliases).
    return {
        "state": data.get("state", "success"),
        "text_original": _original_text(data, text),
        "text_translated": _translated_text(data),
        "translated_character_count": data.get("translatedCharacterCount"),
        "billing_amount": r.headers.get("X-Openapi-Amount"),
        "raw": data,
    }


def _iter_sse_json(resp: requests.Response) -> Iterator[dict[str, Any]]:
    for raw in resp.iter_lines(decode_unicode=True):
        if not raw:
            continue
        line = raw.strip()
        if not line.startswith("data:"):
            continue
        payload = line[5:].lstrip()
        if not payload:
            continue
        chunk = json.loads(payload)
        if isinstance(chunk, dict) and chunk.get("error"):
            raise RuntimeError(str(chunk["error"]))
        if isinstance(chunk, dict):
            yield {
                "state": chunk.get("state", "success"),
                "text_original": _original_text(chunk),
                "text_translated": _translated_text(chunk),
                "translated_character_count": chunk.get("translatedCharacterCount"),
                "progress": chunk.get("progress"),
                "raw": chunk,
            }


def stream_translate(
    text: str,
    lang: str,
    *,
    base_url: str = DEFAULT_BASE,
    api_key: str | None = None,
    timeout: int = 1800,
) -> Iterator[dict[str, Any]]:
    with requests.post(
        _api_url(base_url, "/translate"),
        json=translate_payload(text, lang, mode="stream"),
        headers=_request_headers(api_key),
        stream=True,
        timeout=timeout,
    ) as resp:
        if resp.status_code >= 400:
            try:
                err = resp.json().get("error", resp.text[:500])
            except (json.JSONDecodeError, ValueError):
                err = resp.text[:500]
            raise RuntimeError(err)
        yield from _iter_sse_json(resp)