"""read_document: 让纯文本主模型"读"扫描件 / 图片型 PDF。

markitdown 只能抽 PDF 的文本层,扫描件(纯图页)抽出来是空的 —— 这条死路由本 tool
补上:整份 PDF base64 内联喂豆包 seed-2.0-lite 的文档理解(方舟按页栅格化后视觉
OCR),产出 markdown 全文。格式 / 上限 / 成本均经 scripts/probe_ark_doc.py 实测:
file 内容块 + `data:application/pdf;base64,` 前缀;单页栅格化 3600 万像素上限;
输入约 1300 token/页(100 页约 ¥0.09)。

与 look_at_image 同一 model / key / 记账通道(usage_events kind="vision"),
配置在 config/media/doubao.yaml 的 vision 段(max_pdf_mb / max_pdf_pages)。
"""
from __future__ import annotations

from pathlib import Path
from typing import Any, Optional
from uuid import UUID

from core.ark_client import ArkConfig
from core.storage.usage import record_vision_usage

from .base import Tool
from .output import compact_tool_output
from .image_ref import _CONTAINER_ROOT, load_pdf_as_data_url, resolve_in_root
from .media_common import ark_chat_with_retry, extract_chat_answer, record_usage_safe

_DEFAULT_QUESTION = (
    "这是一份多页 PDF 文档。请逐页把其中的文字完整 OCR 成 markdown:"
    "每页以「== 第N页 ==」开头;表格转成 markdown 表格;保留标题层级与段落换行;"
    "公式尽量用 LaTeX;不要总结、不要遗漏、不要自行补充原文没有的内容。"
)

# 保存到文件时,tool 返回值里带的正文预览长度(全文在文件里,预览只为让模型确认质量)
_PREVIEW_CHARS = 1500


def _count_pdf_pages(pdf: Path) -> Optional[int]:
    """软探页数(pdfminer 随 markitdown[pdf] 已在依赖里);解析不了返 None 不拦路。"""
    try:
        from pdfminer.pdfdocument import PDFDocument
        from pdfminer.pdfpage import PDFPage
        from pdfminer.pdfparser import PDFParser

        with open(pdf, "rb") as f:
            return sum(1 for _ in PDFPage.create_pages(PDFDocument(PDFParser(f))))
    except Exception:
        return None


class ReadDocumentTool(Tool):
    name = "read_document"
    description = (
        "Read a PDF (including SCANNED/image-only PDFs that markitdown can't extract) using "
        "Doubao Seed 2.0 Lite document understanding — OCRs every page into markdown. "
        "Use when markitdown output for a PDF is empty/near-empty (scanned document), or to "
        "ask a specific question about a PDF's content. Pass the PDF path; optionally "
        "`question` (default: full per-page OCR to markdown) and `save_md` (relative path to "
        "write the full text, e.g. 'source/xxx.md' — recommended for multi-page docs so the "
        "full text lands in a file instead of flooding context). Costs roughly 0.001-0.002 "
        "CNY per page; docs over the page limit must be split first."
    )
    parameters = {
        "type": "object",
        "properties": {
            "document": {
                "type": "string",
                "description": (
                    "PDF 相对路径(task_dir 内,如 'source/xxx.pdf',或用户消息里"
                    "`[用户上传的文件]` 行给的路径)。"
                ),
            },
            "question": {
                "type": "string",
                "description": (
                    "想从文档里知道什么(可选)。如「第3章的检测指标是什么」。"
                    "不传则默认逐页完整 OCR 成 markdown。"
                ),
            },
            "save_md": {
                "type": "string",
                "description": (
                    "把全文写到这个相对路径(可选,如 'source/xxx.md')。多页 OCR 建议必传:"
                    "全文落文件,tool 只返回开头预览,不撑爆上下文。"
                ),
            },
        },
        "required": ["document"],
    }

    def __init__(
        self,
        *,
        ark_cfg: ArkConfig,
        vision_variant_cfg: dict,
        variant_key: str,
        working_dir: Path,
        task_id: UUID,
        user_id: UUID,
        base_dir: Optional[Path] = None,
        user_root: Optional[Path] = None,
    ) -> None:
        super().__init__(base_dir, user_root=user_root)
        self.ark_cfg = ark_cfg
        self.cfg = vision_variant_cfg
        self.variant_key = variant_key
        self.working_dir = Path(working_dir)
        self.task_id = task_id
        self.user_id = user_id

    def execute(
        self,
        document: str,
        question: Optional[str] = None,
        save_md: Optional[str] = None,
    ) -> str:
        if not (document or "").strip():
            return "[Error] document(PDF 路径)不能为空"

        cfg = self.cfg
        max_bytes = int(float(cfg.get("max_pdf_mb", 30)) * 1024 * 1024)
        data_url, disp, err = load_pdf_as_data_url(
            document.strip(),
            working_dir=self.working_dir,
            user_root=self.user_root,
            display_fn=self._display,
            max_bytes=max_bytes,
        )
        if err:
            return err

        # 页数软闸:约 1300 token/页,超过上限会撞模型上下文窗口 → 白付一次失败调用。
        # pdfminer 解析不了(加密/损坏)不拦,让 API 报错兜底。
        max_pages = int(cfg.get("max_pdf_pages", 100))
        resolved = resolve_in_root(document.strip(), self.working_dir, self.user_root)
        n_pages = _count_pdf_pages(resolved) if resolved is not None else None
        if n_pages is not None and n_pages > max_pages:
            return (
                f"[Error] 文档 {n_pages} 页超过单次 {max_pages} 页上限(上下文约束)。"
                f"先把 PDF 拆成不超过 {max_pages} 页的分卷再逐卷调用。"
            )

        q = (question or "").strip() or _DEFAULT_QUESTION
        model_id = cfg["model_id"]
        timeout_s = float(cfg.get("doc_request_timeout_s", 600))
        endpoint = cfg.get("endpoint", "/chat/completions")

        body: dict[str, Any] = {
            "model": model_id,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": q},
                        {
                            "type": "file",
                            "file": {"filename": Path(disp).name, "file_data": data_url},
                        },
                    ],
                }
            ],
        }

        # 超时透明重试(细节见 media_common.ark_chat_with_retry)
        resp, api_err = ark_chat_with_retry(
            self.ark_cfg, endpoint, body,
            timeout_s=timeout_s,
            retries=int(cfg.get("timeout_retries", 1)),
            tool_name="read_document",
        )
        if api_err:
            return api_err
        assert resp is not None

        answer, truncated = extract_chat_answer(resp)
        if not answer:
            return (
                "[Error] 文档理解响应缺内容(模型未返回文本)。"
                "可能 PDF 损坏 / 页面超单页像素上限,稍后重试或先重存该 PDF。"
            )

        usage = resp.get("usage") or {}
        tin = int(usage.get("prompt_tokens", 0) or 0)
        tout = int(usage.get("completion_tokens", 0) or 0)

        cost = record_usage_safe(
            "read_document", record_vision_usage,
            task_id=self.task_id,
            user_id=self.user_id,
            model_profile=f"doubao.{self.variant_key}",
            prompt_tokens=tin,
            completion_tokens=tout,
            input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)),
            output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)),
            extra_units={"document": disp},
        )
        cost_cny = float(cost or 0)

        banner = (
            f"[read_document] model={model_id} · document={disp}"
            + (f" · pages={n_pages}" if n_pages else "")
            + f" · tokens={tin}+{tout} · cost=¥{cost_cny:.4f}"
        )
        trunc_note = (
            "\n[注意] 输出被模型 token 上限截断,末尾页可能缺失 —— "
            "用 question 按页段分次问(如「只 OCR 第 51-100 页」)补齐。"
        ) if truncated else ""

        if save_md and save_md.strip():
            saved_disp, save_err = self._save_md(save_md.strip(), answer)
            if save_err:
                # 保存失败不吞掉 OCR 结果:降级为直接返回(截断保护)
                return f"{banner}\n{save_err}(全文改为直接返回)\n\n{compact_tool_output(answer)}{trunc_note}"
            preview = answer[:_PREVIEW_CHARS]
            more = f"\n...(预览截断,全文 {len(answer)} 字符在文件里)" if len(answer) > _PREVIEW_CHARS else ""
            return f"{banner}\nsaved: {saved_disp}\n\n{preview}{more}{trunc_note}"

        return f"{banner}\n\n{compact_tool_output(answer)}{trunc_note}"

    def _save_md(self, rel: str, text: str) -> tuple[str, str]:
        """把全文写到 task 内。返回 (display_path, error)。

        与读取侧 resolve_in_root 同款三形态(相对 / 宿主绝对 / 容器 `/workspace/...`),
        但目标是写入 —— 不要求文件已存在,只做 user_root 边界校验。
        """
        p = Path(rel)
        is_container = rel == _CONTAINER_ROOT or rel.startswith(_CONTAINER_ROOT + "/")
        if self.user_root is not None and is_container:
            target = self.user_root / rel[len(_CONTAINER_ROOT):].lstrip("/")
        elif p.is_absolute():
            target = p
        else:
            target = self.working_dir / p
        target = target.resolve()
        root = (self.user_root or self.working_dir).resolve()
        try:
            target.relative_to(root)
        except ValueError:
            return "", f"[Error] save_md 越界: {rel!r}(须在 task 目录内)"
        try:
            target.parent.mkdir(parents=True, exist_ok=True)
            target.write_text(text, encoding="utf-8")
        except OSError as e:
            return "", f"[Error] 写入 {rel} 失败: {type(e).__name__}: {e}"
        return self._display(target), ""