"""End-to-end smoke test for the MCP server, spoken as a real MCP client.

Run it from ``backend/``::

    .venv/bin/python scripts/smoke_mcp.py

It spawns ``scripts/mcp_server.py`` exactly the way Claude Code, Codex and the
Harness do — a child process, JSON-RPC over its stdin/stdout — and then walks
the whole writing loop through tool calls: create a paper on a template, read
the outline, write paragraphs by *heading*, write a whole generated document in
one call, edit one sentence without rewriting its paragraph, search across
papers, switch the template, and delete everything it created.

Speaking the real protocol is the point. Calling the tool functions directly
would test the domain layer, which ``scripts/smoke_papers.py`` already covers;
what is untested until a client connects is the part that only exists on the
wire — argument schemas, tool discovery, error results, and the stdio stream
staying clean enough to carry the protocol.

Exits non-zero on the first failed expectation and cleans up after itself.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

import anyio

BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

from mcp import types  # noqa: E402
from mcp.client.session import ClientSession  # noqa: E402
from mcp.client.stdio import StdioServerParameters, stdio_client  # noqa: E402
from sqlalchemy import or_, select  # noqa: E402

from app.db.session import SessionLocal  # noqa: E402
from app.models import Paper, Template  # noqa: E402

SERVER_SCRIPT = BACKEND_DIR / "scripts" / "mcp_server.py"

_checks = 0
_created_papers: list[int] = []
_created_templates: list[int] = []


def check(label: str, condition: bool, detail: str = "") -> None:
    """Assert one expectation and report it."""
    global _checks
    _checks += 1
    if not condition:
        raise AssertionError(f"{label}: {detail}")
    print(f"  ok  {label}")


def payload(result: types.CallToolResult) -> Any:
    """The JSON a tool returned, refusing an error result.

    A tool that failed returns ``isError`` with the message as its content, not
    an exception — so a smoke test that only looked at the content would read a
    refusal as data. That is the failure mode this helper exists to prevent.
    """
    text = "".join(
        block.text for block in result.content if isinstance(block, types.TextContent)
    )
    if result.is_error:
        raise AssertionError(f"tool returned an error: {text}")
    return json.loads(text) if text else None


async def call(session: ClientSession, tool: str, **arguments: Any) -> Any:
    """Call one tool and return its JSON payload.

    The parameter is ``tool``, not ``name``: several tools here take an
    argument called ``name`` (``template_create``, ``field_create``) and a
    helper whose own parameter shadowed it would fail with a TypeError that
    looks nothing like the tool call it came from.
    """
    return payload(await session.call_tool(tool, arguments))


async def expect_error(session: ClientSession, tool: str, **arguments: Any) -> str:
    """Call one tool expecting a refusal, and return the message."""
    result = await session.call_tool(tool, arguments)
    if not result.is_error:
        raise AssertionError(f"{tool} should have refused, but returned {result.content}")
    return "".join(
        block.text for block in result.content if isinstance(block, types.TextContent)
    )


def pick_template() -> tuple[int, list[dict[str, Any]]]:
    """The template with the most fields, and its outline as the tools see it."""
    with SessionLocal() as db:
        template = db.scalars(
            select(Template).order_by(Template.id.asc())
        ).all()
        best = max(template, key=lambda item: len(item.items))
        return best.id, [
            {"position": item.sort, "name": item.field.name} for item in best.items
        ]


@asynccontextmanager
async def connect(url: str | None, token: str | None) -> AsyncIterator[ClientSession]:
    """Open a session over stdio (default) or against a running HTTP server.

    Both paths yield the same :class:`ClientSession`, so every check below runs
    unchanged on either transport — which is the point of testing the protocol
    rather than the functions: the tools must behave identically whether a
    client spawned them or connected to them.
    """
    if url:
        # headers live on the HTTP client, not on the transport call.
        import httpx2
        from mcp.client.streamable_http import streamable_http_client

        headers = {"Authorization": f"Bearer {token}"} if token else {}
        async with httpx2.AsyncClient(headers=headers) as http_client:
            async with streamable_http_client(url, http_client=http_client) as (
                read,
                write,
            ):
                async with ClientSession(read, write) as session:
                    yield session
        return

    parameters = StdioServerParameters(
        command=sys.executable,
        args=[str(SERVER_SCRIPT)],
        cwd=str(BACKEND_DIR),
    )
    async with stdio_client(parameters) as (read, write):
        async with ClientSession(read, write) as session:
            yield session


async def run(url: str | None = None, token: str | None = None) -> int:
    template_id, outline = pick_template()
    check("a template with fields exists", len(outline) >= 3, f"got {len(outline)}")
    first = outline[0]["name"]
    second = outline[1]["name"]

    async with connect(url, token) as session:
        handshake = await session.initialize()
        check(
            "server announces itself",
            handshake.server_info.name == "paper-doc",
            str(handshake.server_info),
        )

        print("\n1. tool discovery")
        listed = await session.list_tools()
        names = {tool.name for tool in listed.tools}
        check("29 tools are published", len(names) == 29, f"got {len(names)}")
        check(
            "every group is reachable",
            {"paper_write", "paragraph_write", "sentence_add", "template_get", "field_list"}
            <= names,
        )
        check(
            "no tool returns a structured envelope",
            all(tool.output_schema is None for tool in listed.tools),
        )
        write_tool = next(tool for tool in listed.tools if tool.name == "paper_write")
        check(
            "paper_write describes its paragraph array",
            "paragraphs" in write_tool.input_schema["properties"],
        )

        print("\n2. create a paper on that template")
        paper = await call(
            session,
            "paper_create",
            title="SMOKE MCP 论文",
            template_id=template_id,
            author="smoke",
            keywords="测试, smoke, 测试",
        )
        _created_papers.append(paper["id"])
        check("created", paper["title"] == "SMOKE MCP 论文")
        check("keywords are deduplicated", paper["keywords"] == "测试, smoke", paper["keywords"])
        check(
            "the whole structure exists with no content",
            paper["paragraphs"] == 0 and paper["template_paragraphs"] == len(outline),
            json.dumps(paper, ensure_ascii=False),
        )

        print("\n3. read the outline before writing")
        skeleton = await call(session, "paper_outline", paper_id=paper["id"])
        check(
            "the outline carries every position",
            len(skeleton["paragraphs"]) == len(outline),
            json.dumps(skeleton, ensure_ascii=False),
        )
        check("nothing has been written", all(item["sentences"] == 0 for item in skeleton["paragraphs"]))

        print("\n4. write paragraphs by heading")
        written = await call(
            session,
            "paper_write",
            paper_id=paper["id"],
            paragraphs=[
                {
                    "heading": first,
                    "sentences": ["第一句。", "第二句。"],
                },
                {
                    "heading": second.lstrip("0123456789. "),
                    "text": "这一整段只有一行，按句号切成两句。第二句在这里。",
                    "split": "sentence",
                },
            ],
        )
        check("two paragraphs written", len(written["written"]) == 2, json.dumps(written, ensure_ascii=False))
        check("four sentences total", written["total_sentences"] == 4, str(written["total_sentences"]))
        check(
            "both headings resolved to template positions",
            all(item["heading"] for item in written["written"]),
            json.dumps(written, ensure_ascii=False),
        )

        print("\n5. a heading the template does not have is refused, with candidates")
        message = await expect_error(
            session,
            "paragraph_write",
            paper_id=paper["id"],
            heading="不存在的标题",
            sentences=["x"],
        )
        check("the refusal lists the real headings", first in message, message)

        print("\n6. read one paragraph back")
        paragraph = await call(
            session, "paragraph_get", paper_id=paper["id"], heading=first
        )
        check("two sentences stored", len(paragraph["sentences"]) == 2, json.dumps(paragraph, ensure_ascii=False))
        check(
            "CJK seams take no separator",
            all(row["separator_before"] == "" for row in paragraph["sentences"]),
            json.dumps(paragraph["sentences"], ensure_ascii=False),
        )

        print("\n7. append a sentence with a citation")
        added = await call(
            session,
            "sentence_add",
            paper_id=paper["id"],
            heading=first,
            content="A cited sentence.",
            citations=[{"quote": "引用的原话", "reference_id": 7}],
        )
        check("appended", added["sentence"]["content"] == "A cited sentence.")
        check("the citation is stored", added["sentence"]["citations"][0]["quote"] == "引用的原话")

        after = await call(session, "paragraph_get", paper_id=paper["id"], heading=first)
        check(
            "the seam before it takes a space",
            after["sentences"][-1]["separator_before"] == " ",
            json.dumps(after["sentences"][-1], ensure_ascii=False),
        )

        print("\n8. edit one sentence without touching its paragraph")
        edited = await call(
            session,
            "sentence_update",
            paper_id=paper["id"],
            sentence_id=after["sentences"][0]["id"],
            content="改过的第一句。",
        )
        check("edited", edited["sentence"]["content"] == "改过的第一句。")
        still = await call(session, "paragraph_get", paper_id=paper["id"], heading=first)
        check("the other sentences kept their ids", still["sentences"][1]["id"] == after["sentences"][1]["id"])

        print("\n9. delete one sentence")
        before = await call(session, "sentence_list", paper_id=paper["id"])
        await call(
            session,
            "sentence_delete",
            paper_id=paper["id"],
            sentence_id=still["sentences"][-1]["id"],
        )
        remaining = await call(session, "sentence_list", paper_id=paper["id"])
        check(
            "exactly one sentence fewer",
            remaining["total"] == before["total"] - 1,
            f"{before['total']} -> {remaining['total']}",
        )

        print("\n10. write a whole generated document in one call")
        generated = "\n".join(
            [f"# {item['name']}" + "\n" + f"{item['name']}的正文写在标题下面。这句话是第二句。"
             for item in outline]
        )
        whole = await call(
            session, "paper_write_text", paper_id=paper["id"], text=generated, split="sentence"
        )
        check(
            "every heading in the document was written",
            len(whole["written"]) == len(outline),
            json.dumps(whole, ensure_ascii=False),
        )
        check("nothing was left unmatched", "unmatched" not in whole, json.dumps(whole, ensure_ascii=False))

        print("\n11. read the paper as a document")
        document = await call(session, "paper_document", paper_id=paper["id"])
        check("the text carries every heading", all(item["name"] in document["text"] for item in outline))
        check("and every paragraph", document["text"].count("正文写在标题下面") == len(outline))
        structured = await call(
            session, "paper_document", paper_id=paper["id"], format="json"
        )
        check(
            "the json form carries sentence ids",
            all(
                row.get("id") is not None
                for paragraph in structured["paragraphs"]
                for row in paragraph["sentences"]
            ),
        )

        print("\n12. search across papers")
        found = await call(session, "sentence_search", keyword="正文写在标题下面", paper_id=paper["id"])
        check("the search finds what was written", found["total"] >= 1, str(found["total"]))

        print("\n13. move a paragraph, then clear the one it moved onto")
        home = outline[0]
        target = outline[1]
        relocation = await call(
            session,
            "paragraph_move",
            paper_id=paper["id"],
            heading=home["name"],
            target_heading=target["name"],
        )
        check(
            "the move reports both ends",
            relocation["from"]["position"] == home["position"]
            and relocation["to"]["position"] == target["position"],
            json.dumps(relocation, ensure_ascii=False),
        )
        emptied = await call(session, "paragraph_get", paper_id=paper["id"], position=home["position"])
        check("the source paragraph is empty now", emptied["sentences"] == [], json.dumps(emptied, ensure_ascii=False))
        merged = await call(session, "paragraph_get", paper_id=paper["id"], position=target["position"])
        check(
            "the target holds both paragraphs' sentences",
            len(merged["sentences"]) == 4,
            json.dumps(merged["sentences"], ensure_ascii=False),
        )

        cleared = await call(
            session, "paragraph_delete", paper_id=paper["id"], position=target["position"]
        )
        check("clearing reports how much it removed", cleared["deleted_sentences"] == 4, json.dumps(cleared, ensure_ascii=False))
        after_clear = await call(session, "paragraph_get", paper_id=paper["id"], position=target["position"])
        check("and the paragraph is empty", after_clear["sentences"] == [], json.dumps(after_clear, ensure_ascii=False))

        print("\n14. the field library round-trips")
        field = await call(session, "field_create", name="SMOKE MCP 字段", level=2, font_size=9.5)
        check("created with its typography", field["level"] == 2 and field["font_size"] == 9.5, json.dumps(field, ensure_ascii=False))
        renamed = await call(session, "field_update", field_id=field["id"], name="SMOKE MCP 字段（改）")
        check("renamed", renamed["name"] == "SMOKE MCP 字段（改）", json.dumps(renamed, ensure_ascii=False))
        dropped = await call(session, "field_delete", field_ids=field["id"])
        check("deleted while no template places it", dropped["deleted"] == 1, json.dumps(dropped, ensure_ascii=False))
        check("and it is reported as found, not missing", dropped["missing"] == [], json.dumps(dropped, ensure_ascii=False))

        print("\n15. switch the template and keep the content")
        spare = await call(
            session,
            "template_create",
            name="SMOKE MCP 模板",
            abstract="临时模板",
            field_ids=[_field_id(outline[0]["name"])],
        )
        _created_templates.append(spare["id"])
        check("a template can be built from field ids", spare["fields"][0]["position"] == 10, json.dumps(spare, ensure_ascii=False))

        rebuilt = await call(
            session,
            "template_update",
            template_id=spare["id"],
            name="SMOKE MCP 模板",
            field_ids=[_field_id(outline[0]["name"]), _field_id(outline[1]["name"])],
        )
        check(
            "the outline is replaced wholesale and renumbered",
            [item["position"] for item in rebuilt["fields"]] == [10, 20],
            json.dumps(rebuilt, ensure_ascii=False),
        )

        sentences_before = (await call(session, "sentence_list", paper_id=paper["id"]))["total"]
        switched = await call(
            session, "paper_update", paper_id=paper["id"], template_id=spare["id"]
        )
        check("the switch is reported", switched["template_changed"] is True)
        switched_document = await call(
            session, "paper_document", paper_id=paper["id"], format="json"
        )
        check(
            "content survives the switch, under 未设定",
            sum(len(row["sentences"]) for row in switched_document["paragraphs"])
            == sentences_before,
            json.dumps(switched_document, ensure_ascii=False)[:400],
        )
        check(
            "positions the new template does not define are marked unmatched",
            any(row["matched"] is False for row in switched_document["paragraphs"]),
            json.dumps(switched_document, ensure_ascii=False)[:400],
        )

        print("\n16. a template in use cannot be deleted")
        message = await expect_error(session, "template_delete", template_ids=spare["id"])
        check("the refusal names the blocker", "SMOKE MCP 模板" in message, message)

        print("\n17. deleting a paper needs an explicit confirmation")
        message = await expect_error(session, "paper_delete", paper_ids=paper["id"])
        check("the refusal explains itself", "confirm=true" in message, message)

        print("\n18. delete the paper and the template")
        deleted = await call(session, "paper_delete", paper_ids=paper["id"], confirm=True)
        check("deleted", deleted["deleted"] == 1, json.dumps(deleted, ensure_ascii=False))
        _created_papers.remove(paper["id"])
        gone = await call(session, "template_delete", template_ids=spare["id"])
        check("the template is free to go now", gone["deleted"] == 1, json.dumps(gone, ensure_ascii=False))
        _created_templates.remove(spare["id"])

        print(f"\nall {_checks} checks passed")
        return 0


def _field_id(name: str) -> int:
    """The library field behind a heading name, for the template-create step."""
    from app.models import TemplateFieldLibrary

    with SessionLocal() as db:
        row = db.scalars(
            select(TemplateFieldLibrary).where(TemplateFieldLibrary.name == name)
        ).first()
        if row is None:
            raise AssertionError(f"field {name!r} not in the library")
        return row.id


def cleanup() -> None:
    """Remove anything this run created, even after a failed expectation."""
    if _created_papers:
        with SessionLocal() as db:
            for paper in db.scalars(select(Paper).where(Paper.id.in_(_created_papers))).all():
                db.delete(paper)
            db.commit()
            print(f"cleaned up papers {_created_papers}")
    if _created_templates:
        with SessionLocal() as db:
            for template in db.scalars(
                select(Template).where(Template.id.in_(_created_templates))
            ).all():
                db.delete(template)
            db.commit()
            print(f"cleaned up templates {_created_templates}")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    """``--url`` switches the suite from spawning a server to dialling one."""
    parser = argparse.ArgumentParser(description="paper-doc MCP 冒烟测试")
    parser.add_argument(
        "--url",
        default=os.getenv("MCP_SMOKE_URL"),
        help="连到已在运行的 streamable-http 服务，例如 http://127.0.0.1:8931/mcp；"
        "不传则自己拉起 stdio 子进程",
    )
    parser.add_argument(
        "--token",
        default=os.getenv("MCP_HTTP_TOKEN"),
        help="HTTP 传输的 Bearer 口令",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)

    # The smoke test writes real rows, so it removes its own first: a previous
    # run that died before cleanup would otherwise leave a duplicate title.
    with SessionLocal() as db:
        stale = db.scalars(
            select(Paper).where(or_(Paper.title == "SMOKE MCP 论文"))
        ).all()
        for paper in stale:
            db.delete(paper)
        stale_templates = db.scalars(
            select(Template).where(Template.name == "SMOKE MCP 模板")
        ).all()
        for template in stale_templates:
            db.delete(template)
        db.commit()

    try:
        return anyio.run(run, args.url, args.token)
    except AssertionError as error:
        print(f"\nFAILED: {error}", file=sys.stderr)
        return 1
    finally:
        cleanup()


if __name__ == "__main__":
    sys.exit(main())