"""Paper tools (论文): the library, the document, and the write path.

The write path is the reason this package exists. A model asked to "write the
paper into the system" arrives with a generated document, not with paragraph
positions, so the tools accept either and the server does the translating:

* :func:`paper_write` takes one request per paragraph, each addressed by
  ``heading`` ("1. Introduction") or by ``position``;
* :func:`paper_write_text` takes the whole generated document as one string and
  finds the sections itself, so a twenty-paragraph paper is one tool call.

Both end in :func:`app.crud.paper.replace_paragraph` — the same function the
paragraph editor calls — which is why a paper written through MCP and a paper
written by hand are indistinguishable afterwards.
"""

from __future__ import annotations

import json
from typing import Annotated, Any, Literal

from mcp.server.mcpserver import MCPServer
from mcp.server.mcpserver.exceptions import ToolError
from pydantic import Field

from app.crud import paper as paper_crud
from app.mcp import support as s
from app.mcp.specs import ParagraphSpec
from app.mcp.tools.registry import mcp_tool
from app.models import Paper
from app.schemas.paper import (
    PaperCreate,
    PaperUpdate,
    ParagraphUpdate,
    SentenceCreate,
)


def register(server: MCPServer) -> None:
    """Register every paper tool on ``server``."""

    @mcp_tool(server, title="列出论文")
    def paper_list(
        keyword: Annotated[
            str | None, Field(description="按标题、作者或关键词模糊搜索")
        ] = None,
        status: Annotated[
            Literal["draft", "writing", "done"] | None,
            Field(description="按状态过滤"),
        ] = None,
        template_id: Annotated[
            int | None, Field(description="按所用模板过滤")
        ] = None,
        page: Annotated[int, Field(description="页码，从 1 开始")] = 1,
        page_size: Annotated[int, Field(description="每页条数，最大 200")] = 20,
    ) -> str:
        """列出/搜索论文（论文列表，不含正文）。返回 id、标题、模板、进度。"""
        page, page_size = s.page_args(page, page_size)
        with s.session() as db:
            items, total = paper_crud.list_papers(
                db,
                keyword=keyword,
                status=status,
                template_id=template_id,
                page=page,
                page_size=page_size,
            )
            return s.dumps(
                s.page_payload(
                    items=[
                        {
                            "id": item.id,
                            "title": item.title,
                            "template_id": item.template_id,
                            "template": item.template_name,
                            "author": item.author,
                            "status": item.status,
                            "keywords": item.keywords,
                            "paragraphs": item.paragraph_count,
                            "template_paragraphs": item.template_paragraph_count,
                            "sentences": item.sentence_count,
                        }
                        for item in items
                    ],
                    total=total,
                    page=page,
                    page_size=page_size,
                )
            )

    @mcp_tool(server, title="取一篇论文")
    def paper_get(
        paper_id: Annotated[int, Field(description="论文 id")],
    ) -> str:
        """取一篇论文的元信息与写作进度，不含正文。"""
        with s.session() as db:
            return s.dumps(s.paper_row(s.paper_or_fail(db, paper_id)))

    @mcp_tool(server, title="新建论文")
    def paper_create(
        title: Annotated[str, Field(description="论文标题，不能为空")],
        template_id: Annotated[
            int | None, Field(description="所用模板 id；给出后论文立刻拥有整篇结构")
        ] = None,
        author: Annotated[str | None, Field(description="作者")] = None,
        status: Annotated[
            Literal["draft", "writing", "done"] | None,
            Field(description="状态，默认 draft"),
        ] = None,
        keywords: Annotated[
            str | None, Field(description="关键词，逗号分隔；会去重并存成规范形式")
        ] = None,
        target_journal: Annotated[
            str | None, Field(description="投稿目标期刊")
        ] = None,
    ) -> str:
        """新建一篇论文。只写元信息，正文为空；结构来自模板，随时可换。"""
        payload = PaperCreate(
            title=title,
            template_id=template_id,
            author=author,
            status=status or "draft",
            keywords=keywords,
            target_journal=target_journal,
        )
        with s.session() as db:
            _assert_template_exists(db, payload.template_id)
            return s.dumps(s.paper_row(paper_crud.create(db, payload)))

    @mcp_tool(server, title="修改论文")
    def paper_update(
        paper_id: Annotated[int, Field(description="论文 id")],
        title: Annotated[str | None, Field(description="新标题；不传则不变")] = None,
        template_id: Annotated[
            int | None,
            Field(description="切换到另一个模板；句子留在原位置，标题随之改变"),
        ] = None,
        clear_template: Annotated[
            bool, Field(description="true 表示把模板置空（正文全部落到“未设定”）")
        ] = False,
        author: Annotated[
            str | None, Field(description="作者；空字符串表示清空")
        ] = None,
        status: Annotated[
            Literal["draft", "writing", "done"] | None, Field(description="状态")
        ] = None,
        keywords: Annotated[
            str | None, Field(description="关键词；空字符串表示清空")
        ] = None,
        target_journal: Annotated[
            str | None, Field(description="目标期刊；空字符串表示清空")
        ] = None,
    ) -> str:
        """修改论文元信息，或切换模板（换模板会重塑全文结构，正文不丢）。"""
        values: dict[str, Any] = {}
        if title is not None:
            values["title"] = title
        if clear_template:
            values["template_id"] = None
        elif template_id is not None:
            values["template_id"] = template_id
        if author is not None:
            values["author"] = author
        if status is not None:
            values["status"] = status
        if keywords is not None:
            values["keywords"] = keywords
        if target_journal is not None:
            values["target_journal"] = target_journal

        if not values:
            s.fail("没有给出任何要修改的字段")

        payload = PaperUpdate(**values)
        # The schemas normalise (trimmed title, canonical keywords, blank ->
        # None); only the keys actually sent are applied, so a field left out
        # is not mistaken for an explicit null.
        applied = payload.model_dump(exclude_unset=True)

        with s.session() as db:
            paper = s.paper_or_fail(db, paper_id)
            if "template_id" in applied:
                _assert_template_exists(db, applied["template_id"])

            template_changed = (
                "template_id" in applied
                and applied["template_id"] != paper.template_id
            )
            updated = paper_crud.update(
                db, paper, values=applied, template_changed=template_changed
            )
            return s.dumps({**s.paper_row(updated), "template_changed": template_changed})

    @mcp_tool(server, title="删除论文")
    def paper_delete(
        paper_ids: Annotated[int | list[int], Field(description="论文 id，或 id 数组")],
        confirm: Annotated[
            bool,
            Field(description="必须显式传 true；删除会连同句子和引用一起消失，不可恢复"),
        ] = False,
    ) -> str:
        """删除论文（连句子和引用一起）。必须 confirm=true 才会执行。"""
        ids = s.id_list(paper_ids)
        if not ids:
            s.fail("必须给出至少一个论文 id")
        if not confirm:
            s.fail(
                f"删除论文 {ids} 会连同其全部句子和引用一起消失，且无法恢复。"
                "确认无误后请带 confirm=true 重新调用"
            )
        with s.session() as db:
            return s.dumps({"deleted": paper_crud.delete_many(db, ids), "ids": ids})

    @mcp_tool(server, title="论文骨架")
    def paper_outline(
        paper_id: Annotated[int, Field(description="论文 id")],
    ) -> str:
        """取论文的段落骨架：每段的 position、标题、已写句子数。写正文前先看这个。"""
        with s.session() as db:
            paper = s.paper_or_fail(db, paper_id)
            document = paper_crud.build_document(db, paper)
            written = {
                paragraph.paper_template_filed_sort: len(paragraph.sentences)
                for paragraph in document.paragraphs
            }
            return s.dumps(
                {
                    "paper": s.paper_row(paper),
                    "paragraphs": [
                        {**item, "sentences": written.get(item["position"], 0)}
                        for item in s.heading_index(paper)
                    ],
                    "warnings": document.warnings,
                }
            )

    @mcp_tool(server, title="读论文全文")
    def paper_document(
        paper_id: Annotated[int, Field(description="论文 id")],
        format: Annotated[
            Literal["text", "json", "both"],
            Field(description="text=可直接阅读的全文（默认），json=带 id 的结构，both=两者都要"),
        ] = "text",
        include_empty: Annotated[
            bool, Field(description="是否输出还没有内容的段落")
        ] = False,
        citations: Annotated[
            bool, Field(description="是否带上引用标记与参考文献列表")
        ] = True,
    ) -> str:
        """读整篇论文：按模板顺序输出全文（含引用编号和参考文献）。"""
        with s.session() as db:
            paper = s.paper_or_fail(db, paper_id)
            document = paper_crud.build_document(db, paper)

            payload: dict[str, Any] = {"paper": s.paper_row(paper)}
            if format in ("text", "both"):
                payload["text"] = s.document_text(
                    document, include_empty=include_empty, citations=citations
                )
            if format in ("json", "both"):
                payload["paragraphs"] = [
                    {
                        **s.paragraph_row(paragraph, preview=0),
                        "sentences": [
                            s.sentence_row(row, separator=row.separator_before)
                            for row in paragraph.sentences
                        ],
                    }
                    for paragraph in document.paragraphs
                    if include_empty or paragraph.sentences
                ]
                payload["citations"] = [
                    {
                        "index": item.index,
                        "reference_id": item.reference_id,
                        "quote": item.quote,
                        "sentence_id": item.sentence_id,
                    }
                    for item in document.citations
                ]
            if document.warnings:
                payload["warnings"] = document.warnings
            return s.dumps(payload)

    @mcp_tool(server, title="批量写段落")
    def paper_write(
        paper_id: Annotated[int, Field(description="论文 id")],
        paragraphs: Annotated[
            list[ParagraphSpec],
            Field(
                description=(
                    "要写的段落数组。每段用 heading（如 “1. Introduction”）或 position 指定，"
                    "正文放 sentences（字符串数组或带 citations 的对象数组），"
                    "或放 text + split 让服务端切句"
                )
            ),
        ],
        mode: Annotated[
            Literal["replace", "append"],
            Field(description="replace=整段覆盖（默认），append=追加到该段末尾"),
        ] = "replace",
    ) -> str:
        """把生成好的段落写入论文：按标题或位置定位，一次可写多段。AI 写完直接调用这个。"""
        if not paragraphs:
            s.fail("paragraphs 不能为空")

        with s.session() as db:
            paper = s.paper_or_fail(db, paper_id)
            written: list[dict[str, Any]] = []

            for spec in paragraphs:
                position = s.resolve_position(
                    paper, position=spec.position, heading=spec.heading
                )
                inputs = s.paragraph_inputs(spec)
                heading = _heading_at(paper, position)

                if mode == "append":
                    for item in inputs:
                        paper_crud.append_sentence(
                            db,
                            paper,
                            SentenceCreate(
                                paper_template_filed_sort=position,
                                content=item.content,
                                citations=item.citations,
                                sort=item.sort,
                            ),
                        )
                else:
                    paper_crud.replace_paragraph(
                        db,
                        paper,
                        position,
                        ParagraphUpdate(sentences=inputs),
                    )

                written.append(
                    {
                        "position": position,
                        "heading": heading,
                        "mode": mode,
                        "sentences": len(inputs),
                        "chars": sum(len(item.content) for item in inputs),
                    }
                )

            return s.dumps(
                {
                    "paper_id": paper.id,
                    "written": written,
                    "total_sentences": sum(item["sentences"] for item in written),
                }
            )

    @mcp_tool(server, title="整篇写入")
    def paper_write_text(
        paper_id: Annotated[int, Field(description="论文 id")],
        text: Annotated[
            str,
            Field(
                description=(
                    "整篇论文的文本。用 Markdown 标题（# 0 Abstract / ## 1. Introduction）"
                    "或与模板同名的标题行分段，标题下的内容写进对应段落"
                )
            ),
        ],
        split: Annotated[
            Literal["line", "sentence", "paragraph"],
            Field(description="段落内怎么切句：line 一行一句（默认）/ sentence 按句号 / paragraph 整段一句"),
        ] = "line",
        mode: Annotated[
            Literal["replace", "append"],
            Field(description="replace=整段覆盖（默认），append=追加到该段末尾"),
        ] = "replace",
        strict: Annotated[
            bool,
            Field(description="true=有标题匹配不上就整篇不写；false=能写的先写，其余在 unmatched 里返回"),
        ] = False,
    ) -> str:
        """把一整篇生成好的论文按标题切段后写入论文。适合“AI 写完直接灌进去”。"""
        sections = _split_sections(text)
        if not sections:
            s.fail("没有从 text 里找到任何与模板同名的标题，请改用 paper_write 按 position 写")

        paragraphs: list[ParagraphSpec] = []
        unmatched: list[str] = []

        with s.session() as db:
            paper = s.paper_or_fail(db, paper_id)
            for heading, body in sections:
                try:
                    s.resolve_position(paper, heading=heading)
                except ToolError as error:
                    # ToolError carries the candidate list the model needs to
                    # correct the heading, so it is reported rather than raised
                    # when the caller asked for a lenient run.
                    unmatched.append(f"{heading}：{error}")
                    if strict:
                        s.fail(f"标题“{heading}”无法定位到模板段落；{error}")
                    continue
                paragraphs.append(
                    ParagraphSpec(heading=heading, text=body, split=split)
                )

        if not paragraphs:
            s.fail("没有任何标题能定位到模板段落：" + "；".join(unmatched))

        payload = json.loads(paper_write(paper_id=paper_id, paragraphs=paragraphs, mode=mode))
        if unmatched:
            payload["unmatched"] = unmatched
        return s.dumps(payload)


# --- helpers -----------------------------------------------------------------


def _assert_template_exists(db: Any, template_id: int | None) -> None:
    """Refuse a paper pointing at a template that is not there.

    Checked in Python rather than by a foreign key, exactly as the REST route
    does: TiDB parses but does not enforce ``FOREIGN KEY``, so an unchecked
    write would leave a paper with no structure and no way to notice.
    """
    if template_id is None:
        return
    s.template_or_fail(db, template_id)


def _heading_at(paper: Paper, position: int) -> str | None:
    """What the paper's template calls this position, if it names it at all."""
    for item in s.heading_index(paper):
        if item["position"] == position:
            return str(item["heading"])
    return None


#: A Markdown ATX heading, the most likely shape of generated output.
_MARKDOWN_HEADING = "#"


def _split_sections(text: str) -> list[tuple[str, str]]:
    """Cut a generated document into ``(heading, body)`` pairs.

    Two heading shapes are recognised, in this order:

    1. a Markdown heading line — ``# 0 Abstract`` / ``## 1. Introduction``;
    2. a plain line that is *exactly* one of the template's headings after
       normalisation, which is what a model produces when it copies the outline
       out of :func:`paper_outline` and writes prose under it.

    A body line is never mistaken for a heading: rule 2 needs a short line with
    no sentence-ending punctuation, and it still has to match a heading the
    caller's paper actually defines. Anything before the first heading is
    dropped — there is no paragraph to put it in — which is why the tool
    refuses outright when no heading matched at all.
    """
    sections: list[tuple[str, str]] = []
    heading: str | None = None
    body: list[str] = []

    def close() -> None:
        if heading is not None:
            sections.append((heading, "\n".join(body).strip()))

    for line in text.splitlines():
        stripped = line.strip()
        candidate = stripped.lstrip(_MARKDOWN_HEADING).strip() if stripped.startswith(_MARKDOWN_HEADING) else None

        if candidate is not None:
            if not candidate:
                continue
            close()
            heading, body = candidate, []
            continue

        if _looks_like_plain_heading(stripped):
            close()
            heading, body = stripped, []
            continue

        if heading is not None:
            body.append(line)

    close()
    # A heading with an empty body is a section the caller left blank; keeping
    # it would overwrite a written paragraph with nothing.
    return [(name, block) for name, block in sections if block]


def _looks_like_plain_heading(line: str) -> bool:
    """Whether a plain line could be a heading rather than prose.

    Deliberately conservative. The caller's paper decides the final answer in
    :func:`app.mcp.support.resolve_position`; this only avoids handing it every
    sentence in the document.
    """
    if not line or len(line) > 60:
        return False
    if line[-1] in "。！？.!?,，、;；:：":
        return False
    # A heading has no more than a handful of words before the first full stop;
    # a prose line at this length that ends without punctuation is rare enough
    # that the template match is what settles it.
    return True