"""
天目 OCR HTTP 客户端。

封装对天目微服务的 HTTP 调用，提供同步/异步 OCR 接口。
后端其他模块通过此客户端调用 OCR 能力，不直接依赖 HTTP 细节。

特性：
  - 自动重试（网络故障 / 5xx 错误，指数退避）
  - 连接池复用（httpx.AsyncClient 自动管理）
  - 端点级别超时配置（/healthz 短超时，/v1/ocr 适中，/v1/ocr/vl 长超时）
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import httpx

from app.core.logging import logger

TIANMU_DEFAULT_BASE = "http://localhost:8101"

# 各端点默认超时（秒）
_TIMEOUTS = {
    "health": 5.0,
    "ocr": 60.0,
    "ocr_vl": 600.0,
}

# 重试配置
_MAX_RETRIES = 2
_RETRYABLE_STATUSES = {502, 503, 504}
_RETRY_BASE_DELAY = 0.5  # 秒


@dataclass
class OCRBlock:
    """单行 OCR 识别结果。"""

    text: str
    confidence: float
    bbox: list[list[float]] = field(default_factory=list)


@dataclass
class OCRResult:
    """统一 OCR 响应。"""

    engine: str
    text: str
    blocks: list[OCRBlock] = field(default_factory=list)
    duration_ms: float = 0.0
    error: str | None = None
    saved_md_path: str | None = None

    @property
    def success(self) -> bool:
        return self.error is None


class OCRClient:
    """天目 OCR 服务 HTTP 客户端。

    用法::

        client = OCRClient()
        result = await client.ocr(image_bytes)
        await client.close()
    """

    def __init__(
        self,
        base_url: str = "",
        http_client: httpx.AsyncClient | None = None,
        timeout: float = 600.0,
        max_retries: int = _MAX_RETRIES,
    ) -> None:
        self.base_url = (base_url or TIANMU_DEFAULT_BASE).rstrip("/")
        self._max_retries = max_retries
        self._http = http_client or httpx.AsyncClient(timeout=timeout)
        self._own_http = http_client is None

    async def close(self) -> None:
        """释放 HTTP 客户端。"""
        if self._own_http:
            await self._http.aclose()

    # ── 通用 OCR（PP-OCRv6 优先）─────────────────────────

    async def ocr(
        self,
        image: bytes | str,
        *,
        engine: str = "auto",
        lang: list[str] | None = None,
        save_md: bool = False,
    ) -> OCRResult:
        """通用文字识别。

        Args:
            image: 图片二进制数据或 HTTP URL。
            engine: 引擎选择 auto / ppocrv6 / paddleocr-vl。
            lang: 语言列表（默认中英文，仅 ppocrv6 生效）。
            save_md: 是否保存结果为 Markdown 文件。

        Returns:
            OCRResult: 识别结果。
        """
        body: dict[str, Any] = {"engine": engine}
        if lang:
            body["lang"] = lang
        if save_md:
            body["save_md"] = True

        body["image"] = _encode_image(image)

        data = await self._request_with_retry(
            "POST",
            f"{self.base_url}/v1/ocr",
            json=body,
            timeout=_TIMEOUTS["ocr"],
        )
        if data is None:
            return OCRResult(engine=engine, error="天目服务不可达（重试耗尽）")

        blocks = [
            OCRBlock(
                text=b["text"],
                confidence=b["confidence"],
                bbox=b.get("bbox", {}).get("points", []),
            )
            for b in data.get("blocks", [])
        ]

        return OCRResult(
            engine=data.get("engine", ""),
            text=data.get("text", ""),
            blocks=blocks,
            duration_ms=data.get("duration_ms", 0),
            error=data.get("error"),
            saved_md_path=data.get("saved_md_path"),
        )

    # ── 复杂文档解析（PaddleOCR-VL）──────────────────────

    async def ocr_vl(
        self,
        image: bytes | str,
        *,
        output_format: str = "markdown",
        save_md: bool = False,
    ) -> OCRResult:
        """复杂文档解析（PaddleOCR-VL）。

        Args:
            image: 图片二进制数据或 HTTP URL。
            output_format: 输出格式 markdown / json / text。
            save_md: 是否保存结果为 Markdown 文件。

        Returns:
            OCRResult: 结构化解析结果。
        """
        body: dict[str, Any] = {"output_format": output_format}
        if save_md:
            body["save_md"] = True

        body["image"] = _encode_image(image)

        data = await self._request_with_retry(
            "POST",
            f"{self.base_url}/v1/ocr/vl",
            json=body,
            timeout=_TIMEOUTS["ocr_vl"],
        )
        if data is None:
            return OCRResult(engine="paddleocr-vl", error="天目服务不可达（重试耗尽）")

        return OCRResult(
            engine=data.get("engine", "paddleocr-vl"),
            text=data.get("text", ""),
            duration_ms=data.get("duration_ms", 0),
            error=data.get("error"),
            saved_md_path=data.get("saved_md_path"),
        )

    # ── 健康检查 ──────────────────────────────────────────

    async def health(self) -> dict:
        """检查天目服务健康状态。"""
        resp = await self._http.get(
            f"{self.base_url}/healthz",
            timeout=_TIMEOUTS["health"],
        )
        resp.raise_for_status()
        return resp.json()

    # ── 内部：重试逻辑 ────────────────────────────────────

    async def _request_with_retry(
        self,
        method: str,
        url: str,
        *,
        json: dict[str, Any] | None = None,
        timeout: float = 60.0,
    ) -> dict[str, Any] | None:
        """带指数退避重试的 HTTP 请求。

        Args:
            method: HTTP 方法。
            url: 请求 URL。
            json: JSON 请求体。
            timeout: 请求超时（秒）。

        Returns:
            响应 JSON 数据，重试耗尽返回 None。
        """

        for attempt in range(self._max_retries + 1):
            try:
                resp = await self._http.request(method, url, json=json, timeout=timeout)

                # 可重试的服务器错误
                if resp.status_code in _RETRYABLE_STATUSES and attempt < self._max_retries:
                    delay = _RETRY_BASE_DELAY * (2 ** attempt)
                    logger.warning(
                        "天目返回 {}，{}s 后重试 ({}/{})",
                        resp.status_code, delay, attempt + 1, self._max_retries,
                    )
                    await asyncio.sleep(delay)
                    continue

                resp.raise_for_status()
                return resp.json()

            except (httpx.TimeoutException, httpx.ConnectError, httpx.RemoteProtocolError) as e:
                if attempt < self._max_retries:
                    delay = _RETRY_BASE_DELAY * (2 ** attempt)
                    logger.warning(
                        "天目请求失败 ({}，{}s 后重试): {}",
                        type(e).__name__, delay, e,
                    )
                    await asyncio.sleep(delay)
                    continue

                logger.error("天目请求重试耗尽 ({}): {}", type(e).__name__, e)

        return None


def _encode_image(image: bytes | str | Path) -> str:
    """将图片编码为 base64 字符串。

    Args:
        image: 图片二进制数据、文件路径或 HTTP URL。

    Returns:
        base64 编码字符串（二进制输入）或原 URL。
    """
    if isinstance(image, Path):
        image = image.read_bytes()
    if isinstance(image, str):
        return image

    import base64
    return base64.b64encode(image).decode("ascii")