"""services/extraction.py —— 抽取任务上游代理（8000 主 + 7900 回落，API.md §10）。

- 默认 `RAG_BASE_URL`（http://127.0.0.1:8000），每次请求失败回落
  `RAG_ALT_BASE_URL`（http://127.0.0.1:7900）。
- 回落失败状态不缓存超过 30s。
- progress_pct 由 BFF 按 stage 派生：queued=5 / mineru=35 / llm=75 /
  completed=100 / failed=null（上游没有进度数值字段）。
"""
from __future__ import annotations

import logging
import os
import tempfile
import time
from typing import Any, AsyncIterator

import anyio
import httpx

logger = logging.getLogger(__name__)

_UPSTREAM_PREFIX = "/api/v1/methodology"
_DEFAULT_PRIMARY = "http://127.0.0.1:8000"
_DEFAULT_ALT = "http://127.0.0.1:7900"

PRIMARY_BASE_URL = os.environ.get("RAG_BASE_URL", _DEFAULT_PRIMARY)
ALT_BASE_URL = os.environ.get("RAG_ALT_BASE_URL", _DEFAULT_ALT)

# 回落失败状态缓存上限（API.md §10：不缓存失败状态超过 30s）
ALT_FAILURE_TTL = 30.0

# 超时阶梯（API.md §10）
READ_TIMEOUT = 30.0
UPLOAD_READ_TIMEOUT = 600.0

# 上传限制（FR-EXTRACT-004）
MAX_PDF_BYTES = 50 * 1024 * 1024  # 50MB
_SPOOL_MAX_MEMORY = 5 * 1024 * 1024  # spool 驻内存上限，超出落临时文件

_UPLOAD_ROUTE = "/extract-upload"

# 上游成功上传形态：{success, job_id, status_url}；message 由 BFF 附加
_SUCCESS_UPLOAD_MESSAGE = "已提交抽取任务"

_PROGRESS_BY_STAGE: dict[str, int | None] = {
    "queued": 5,
    "mineru": 35,
    "llm": 75,
    "completed": 100,
    "failed": None,
}


def derive_progress_pct(stage: str | None) -> int | None:
    """按 stage 派生 progress_pct（API.md §10）。未知 stage 返回 None。"""
    return _PROGRESS_BY_STAGE.get(stage if stage is not None else "")


class ExtractionUnavailable(Exception):
    """两个抽取上游（8000 + 7900）均不可达。"""

    def __init__(self, message: str, *, cause: BaseException | None = None) -> None:
        super().__init__(message)
        self.cause = cause


class ExtractionUpstreamError(Exception):
    """两个上游均可连通但都返回了错误（5xx，如 MinerU 不可用）。"""

    def __init__(self, message: str) -> None:
        super().__init__(message)


class ExtractionService:
    """每次请求先打 primary（8000）；连接失败时回落 alt（7900）。

    失败回落状态记 `_alt_dead_until` 时间戳（30s 内不再试 alt），
    任何 alt 成功都会清掉该状态。
    """

    def __init__(
        self,
        primary_base_url: str = PRIMARY_BASE_URL,
        alt_base_url: str = ALT_BASE_URL,
    ) -> None:
        self._primary = primary_base_url.rstrip("/")
        self._alt = alt_base_url.rstrip("/")
        self._alt_dead_until: float = 0.0
        self._client: httpx.AsyncClient | None = None

    @property
    def client(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(
                timeout=httpx.Timeout(READ_TIMEOUT, connect=3.0)
            )
        return self._client

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

    # ---------- 回落机制 ----------

    def _alt_usable(self) -> bool:
        return time.monotonic() >= self._alt_dead_until

    def _mark_alt_dead(self) -> None:
        self._alt_dead_until = time.monotonic() + ALT_FAILURE_TTL

    def _mark_alt_alive(self) -> None:
        self._alt_dead_until = 0.0

    async def _request(
        self,
        path: str,
        *,
        method: str = "GET",
        read_timeout: float = READ_TIMEOUT,
        **kwargs: Any,
    ) -> httpx.Response:
        """primary → alt 回落。连接失败回落 alt；5xx 时也回落 alt，
        两个上游都 5xx 则抛 ExtractionUpstreamError（含最后一个 5xx 响应）。
        """
        client = self.client
        primary_failed: Exception | None = None
        primary_status: int | None = None

        try:
            resp = await client.request(
                method,
                f"{self._primary}{_UPSTREAM_PREFIX}{path}",
                timeout=httpx.Timeout(read_timeout, connect=3.0),
                **kwargs,
            )
            if resp.status_code < 500:
                return resp
            primary_status = resp.status_code
            primary_failed = RuntimeError(f"primary 返回 {resp.status_code}")
        except (httpx.HTTPError, httpx.TimeoutException) as exc:
            primary_failed = exc

        # ---- 回落 alt ----
        if not self._alt_usable():
            raise ExtractionUnavailable(
                "抽取服务（8000）不可用，且 7900 最近不可用，稍后自动重试",
                cause=primary_failed,
            )
        try:
            resp = await client.request(
                method,
                f"{self._alt}{_UPSTREAM_PREFIX}{path}",
                timeout=httpx.Timeout(read_timeout, connect=3.0),
                **kwargs,
            )
            if resp.status_code < 500:
                self._mark_alt_alive()
                return resp
            self._mark_alt_dead()
            raise ExtractionUpstreamError(
                f"两个抽取上游均返回错误：8000 状态 {primary_status}，"
                f"7900 状态 {resp.status_code}"
            )
        except (httpx.HTTPError, httpx.TimeoutException) as exc:
            self._mark_alt_dead()
            raise ExtractionUnavailable(
                "抽取服务不可用（8000 与 7900 均不可达）",
                cause=exc,
            ) from exc

    # ---------- 归一化 ----------

    @staticmethod
    def normalize_job(item: dict[str, Any]) -> dict[str, Any]:
        """列表 item 归一化（API.md §10 列表形态）。

        title = source_filename 优先 → paper_id → job_id。
        """
        job_id = item.get("job_id")
        title = item.get("source_filename") or item.get("paper_id") or job_id
        stage = item.get("stage")
        return {
            "job_id": job_id,
            "title": title,
            "status": item.get("status"),
            "stage": stage,
            "substage": item.get("substage"),
            "progress_pct": derive_progress_pct(stage),
            "message": item.get("message"),
            "paper_id": item.get("paper_id"),
            "family": item.get("family"),
            "source_filename": item.get("source_filename"),
            "stage_count": item.get("stage_count"),
            "species_count": item.get("species_count"),
            "figure_table_caption_count": item.get("figure_table_caption_count"),
            "figure_table_stage_count": item.get("figure_table_stage_count"),
            "updated_at": item.get("updated_at"),
            "config_url": item.get("config_url"),
            "methodology_plan_url": item.get("methodology_plan_url"),
        }

    # ---------- 业务操作 ----------

    async def list_jobs(self, limit: int = 20) -> dict[str, Any]:
        """GET /extract-jobs（可带 ?limit=，默认 20）。"""
        resp = await self._request(f"/extract-jobs?limit={limit}")
        payload = resp.json()
        return {
            "jobs": [self.normalize_job(item) for item in payload.get("jobs", [])],
            "count": len(payload.get("jobs", [])),
        }

    async def get_job(self, job_id: str) -> dict[str, Any]:
        """GET /extract-jobs/{id}：扁平详情归一化（API.md §10）。

        plan：completed 时请求 /plan；未完成时上游 409 → 置 null 不报错。
        """
        resp = await self._request(f"/extract-jobs/{job_id}")
        item = resp.json()

        status = item.get("status")
        stage = item.get("stage")
        job_id_val = item.get("job_id") or job_id
        title = (
            item.get("source_filename")
            or item.get("paper_id")
            or job_id_val
        )

        plan: dict[str, Any] | None = None
        if status == "completed":
            try:
                plan_resp = await self._request(f"/extract-jobs/{job_id}/plan")
            except ExtractionUnavailable as exc:
                logger.warning(
                    "extraction: job %s 的 plan 请求失败（上游不可用）：%s", job_id, exc
                )
            else:
                if plan_resp.status_code == 409:
                    # 上游认为未完成（与扁平详情不一致），契约要求置 null 不报错
                    plan = None
                else:
                    plan = plan_resp.json()

        return {
            "job_id": job_id_val,
            "title": title,
            "status": status,
            "stage": stage,
            "progress_pct": derive_progress_pct(stage),
            "message": item.get("message"),
            "paper_id": item.get("paper_id"),
            "family": item.get("family"),
            "stats": {
                "stage_count": item.get("stage_count"),
                "species_count": item.get("species_count"),
                "figure_table_caption_count": item.get("figure_table_caption_count"),
                "figure_table_stage_count": item.get("figure_table_stage_count"),
            },
            "detail": item,
            "plan": plan,
        }

    async def _passthrough(self, path: str) -> dict[str, Any]:
        """processing-status / progress / summary / run-log：JSON 透传。

        上游 404 视为空数据（rows 为空属正常现象，前端显示空态）。
        """
        resp = await self._request(path)
        return resp.json()

    async def upload(
        self,
        *,
        filename: str,
        stream: AsyncIterator[bytes] | bytes,
        content_type: str | None,
        paper_id: str = "uploaded_paper",
        family: str = "",
        output_name: str = "",
    ) -> dict[str, Any]:
        """POST /extract-upload：multipart 转发（内存有界）。

        httpx 的 multipart 编码只接受同步 file-like/bytes，异步流会先
        spool 进 SpooledTemporaryFile（≤5MB 驻内存，超出落磁盘），
        因此不会整文件读进内存。
        返回上游形态 + message:"已提交抽取任务"。
        """
        created = not isinstance(stream, (bytes, bytearray))
        if created:
            spooled: Any = tempfile.SpooledTemporaryFile(max_size=_SPOOL_MAX_MEMORY)
        else:
            spooled = stream
        try:
            if created:
                async for chunk in stream:
                    await anyio.to_thread.run_sync(spooled.write, chunk)
                spooled.seek(0)
            resp = await self._request(
                _UPLOAD_ROUTE,
                method="POST",
                read_timeout=UPLOAD_READ_TIMEOUT,
                data={
                    "paper_id": paper_id,
                    "family": family,
                    "output_name": output_name,
                },
                files={"file": (filename, spooled, content_type)},
            )
        finally:
            if created:
                spooled.close()
        payload = resp.json()
        payload["message"] = _SUCCESS_UPLOAD_MESSAGE
        return payload


# 模块级单例：路由层共享（app 关闭时 close）。
service = ExtractionService()


async def close_extraction() -> None:
    """app shutdown 钩子。"""
    await service.close()