"""多章链图单测（chain-workflow 设计 §3 / §10）——节点纯逻辑 + 全图 MemorySaver 流程。

全部注入 mock 网关 + fake session_factory + fake repos + MemorySaver，无真 LLM、无真
Postgres。验：① 无冲突全跑到 END 且 written 满；② 有冲突 → `__interrupt__`；③ resume
（`Command(resume=decisions)`）→ accept → END；④ `decide` 纯逻辑；⑤ accept 推进章号。
asyncio_mode=auto。
"""

from __future__ import annotations

import uuid
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any

from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command
from pydantic import BaseModel
from ww_agents import Conflict, ContinuityReview
from ww_core.orchestrator.chain import (
    accept_chapter,
    build_chain_graph,
    decide,
    has_conflicts,
    initial_chain_state,
    write_chapter,
)
from ww_llm_gateway.types import LlmRequest, LlmResponse, ServedBy, Usage

PROJECT = uuid.UUID("00000000-0000-0000-0000-000000000001")
USER = uuid.UUID("00000000-0000-0000-0000-0000000000aa")


# ---- 测试替身 ----


def _usage() -> Usage:
    return Usage(
        provider="fake",
        model="fake",
        input_tokens=1,
        output_tokens=1,
        cost_minor=0,
        currency="USD",
    )


class FakeChainGateway:
    """链网关 mock：write（无 output_schema）回固定正文；审稿按 `output_schema` 回 parsed。

    `conflicts`：注给 continuity 审的冲突清单（空 = 无冲突章）。其余审 schema 回各自空实例。
    """

    def __init__(
        self,
        *,
        draft_text: str,
        conflicts: list[Conflict],
        by_schema: dict[type, BaseModel],
    ) -> None:
        self._draft = draft_text
        self._continuity = ContinuityReview(conflicts=conflicts)
        self._by_schema = by_schema
        self.write_calls = 0
        self.review_calls = 0
        self.write_inputs: list[str] = []

    async def run(self, req: LlmRequest) -> LlmResponse:
        schema: type[BaseModel] | None = req.output_schema
        if schema is None:
            self.write_calls += 1
            self.write_inputs.append(req.input if isinstance(req.input, str) else "")
            # 每次写章产不同正文（便于断言下一章续写注入了上一章正文）。
            draft = f"{self._draft}#{self.write_calls}"
            return LlmResponse(
                text=draft,
                parsed=None,
                usage=_usage(),
                served_by=ServedBy(provider="fake", model="fake"),
            )
        self.review_calls += 1
        parsed: BaseModel = (
            self._continuity if schema is ContinuityReview else self._by_schema[schema]
        )
        return LlmResponse(
            text=parsed.model_dump_json(),
            parsed=parsed,
            usage=_usage(),
            served_by=ServedBy(provider="fake", model="fake"),
        )


class FakeDraftView:
    def __init__(self, content: str) -> None:
        self.content = content


class FakeChapterRepo:
    """内存章草稿 repo：save_draft 存正文，get_draft 回读，latest_accepted 读已验收正文。"""

    def __init__(self, store: dict[int, str], accepted: dict[int, str] | None = None) -> None:
        self._store = store
        self._accepted = accepted if accepted is not None else {}

    async def save_draft(
        self, project_id: uuid.UUID, chapter_no: int, *, text: str, volume: int = 1
    ) -> Any:
        self._store[chapter_no] = text
        return FakeDraftView(text)

    async def get_draft(self, project_id: uuid.UUID, chapter_no: int) -> Any:
        content = self._store.get(chapter_no)
        return FakeDraftView(content) if content is not None else None

    async def latest_accepted(self, project_id: uuid.UUID, chapter_no: int) -> Any:
        content = self._accepted.get(chapter_no)
        return FakeDraftView(content) if content is not None else None


class FakeReviewRecord:
    def __init__(self, chapter_no: int, conflicts: list[dict[str, Any]]) -> None:
        self.id = uuid.uuid4()
        self.chapter_no = chapter_no
        self.conflicts = conflicts


class FakeReviewRepo:
    """内存审稿 repo：record 追加留痕，list_for_chapter 回新→旧。"""

    def __init__(self, records: dict[int, list[FakeReviewRecord]]) -> None:
        self._records = records

    async def record(
        self,
        project_id: uuid.UUID,
        chapter_no: int,
        *,
        chapter_version: int | None = None,
        conflicts: list[dict[str, Any]],
        foreshadow_sug: list[dict[str, Any]] | None = None,
        style: dict[str, Any] | None = None,
        pace: dict[str, Any] | None = None,
        characterization: dict[str, Any] | None = None,
        health_score: int | None = None,
    ) -> Any:
        rec = FakeReviewRecord(chapter_no, list(conflicts))
        self._records.setdefault(chapter_no, []).insert(0, rec)
        return rec

    async def list_for_chapter(self, project_id: uuid.UUID, chapter_no: int) -> list[Any]:
        return list(self._records.get(chapter_no, []))


class FakeSession:
    def __init__(self) -> None:
        self.commits = 0

    async def commit(self) -> None:
        self.commits += 1


class FakeContext:
    def __init__(self) -> None:
        self.stable_core = "## 世界观硬规则\n灵气可凝丹"
        self.volatile = "## 写作指令\n写第 N 章"


@asynccontextmanager
async def _session_factory() -> AsyncIterator[FakeSession]:
    """fake 短事务 session（每次新建；满足 SessionFactory）。"""
    yield FakeSession()


def _empty_by_schema() -> dict[type, BaseModel]:
    """非 continuity 审（pace/style/foreshadow）的空 parsed——按各 spec 的 output_schema。"""
    from ww_agents import foreshadow_spec, pace_spec, style_drift_spec

    out: dict[type, BaseModel] = {}
    for spec in (foreshadow_spec, pace_spec, style_drift_spec):
        schema = spec.output_schema
        assert schema is not None
        out[schema] = schema()
    return out


def _make_harness(*, conflicts: list[Conflict]) -> dict[str, Any]:
    """造一套共享内存态 + 工厂闭包 + accept 间谍，供图/节点测试复用。"""
    draft_store: dict[int, str] = {}
    accepted_store: dict[int, str] = {}
    review_store: dict[int, list[FakeReviewRecord]] = {}
    accepted: list[dict[str, Any]] = []

    gateway = FakeChainGateway(
        draft_text="第 N 章正文。", conflicts=conflicts, by_schema=_empty_by_schema()
    )

    async def gateway_builder(session: Any) -> Any:
        """按节点 session 建网关的替身（不变量 #1）：忽略 session，恒返同一 mock 网关。"""
        return gateway

    def memory_repos_factory(session: Any) -> Any:
        return object()  # assemble 被 fake 替换，不实际用 repos

    def chapter_repo_factory(session: Any) -> Any:
        return FakeChapterRepo(draft_store, accepted_store)

    def review_repo_factory(session: Any) -> Any:
        return FakeReviewRepo(review_store)

    async def assemble(repos: Any, project_id: uuid.UUID, chapter_no: int) -> Any:
        return FakeContext()

    async def accept_op(
        *,
        project_id: uuid.UUID,
        chapter_no: int,
        user_id: uuid.UUID,
        decisions: list[Any],
    ) -> None:
        accepted.append({"chapter_no": chapter_no, "decisions": list(decisions)})
        # 仿真验收落库：把本章草稿晋升为 accepted（供下一章续写读前文）。
        if chapter_no in draft_store:
            accepted_store[chapter_no] = draft_store[chapter_no]

    return {
        "gateway": gateway,
        "gateway_builder": gateway_builder,
        "session_factory": _session_factory,
        "memory_repos_factory": memory_repos_factory,
        "chapter_repo_factory": chapter_repo_factory,
        "review_repo_factory": review_repo_factory,
        "assemble": assemble,
        "accept_op": accept_op,
        "accepted": accepted,
        "draft_store": draft_store,
        "accepted_store": accepted_store,
        "review_store": review_store,
    }


def _build(h: dict[str, Any], checkpointer: Any) -> Any:
    return build_chain_graph(
        h["gateway_builder"],
        session_factory=h["session_factory"],
        memory_repos_factory=h["memory_repos_factory"],
        chapter_repo_factory=h["chapter_repo_factory"],
        review_repo_factory=h["review_repo_factory"],
        assemble=h["assemble"],
        accept_op=h["accept_op"],
        checkpointer=checkpointer,
    )


# ---- decide 纯逻辑 ----


async def test_decide_no_conflicts_sets_false() -> None:
    review_store: dict[int, list[FakeReviewRecord]] = {1: [FakeReviewRecord(1, [])]}

    def review_repo_factory(session: Any) -> Any:
        return FakeReviewRepo(review_store)

    out = await decide(
        {"project_id": PROJECT, "user_id": USER, "current_chapter_no": 1},
        session_factory=_session_factory,
        review_repo_factory=review_repo_factory,
    )

    assert out == {"has_conflicts": False}


async def test_decide_with_conflicts_sets_true() -> None:
    review_store: dict[int, list[FakeReviewRecord]] = {
        1: [FakeReviewRecord(1, [{"kind": "x", "detail": "矛盾"}])]
    }

    def review_repo_factory(session: Any) -> Any:
        return FakeReviewRepo(review_store)

    out = await decide(
        {"project_id": PROJECT, "user_id": USER, "current_chapter_no": 1},
        session_factory=_session_factory,
        review_repo_factory=review_repo_factory,
    )

    assert out == {"has_conflicts": True}


def test_has_conflicts_reads_flag() -> None:
    assert has_conflicts({"has_conflicts": True}) is True
    assert has_conflicts({"has_conflicts": False}) is False
    assert has_conflicts({}) is False


# ---- accept 节点：推进章号 + 追加 written + 透传裁决 ----


async def test_accept_chapter_advances_and_records_decisions() -> None:
    accepted: list[dict[str, Any]] = []

    async def accept_op(
        *,
        project_id: uuid.UUID,
        chapter_no: int,
        user_id: uuid.UUID,
        decisions: list[Any],
    ) -> None:
        accepted.append({"chapter_no": chapter_no, "decisions": list(decisions)})

    out = await accept_chapter(
        {
            "project_id": PROJECT,
            "user_id": USER,
            "current_chapter_no": 3,
            "written": [1, 2],
            "last_chapter_no": 5,
        },
        decisions=[{"conflict_index": 0, "verdict": "ignore"}],
        accept_op=accept_op,
    )

    assert out["current_chapter_no"] == 4
    assert out["written"] == [1, 2, 3]
    assert out["has_conflicts"] is False
    assert accepted == [
        {"chapter_no": 3, "decisions": [{"conflict_index": 0, "verdict": "ignore"}]}
    ]


# ---- write 节点：收集版（非流式）落稿 ----


async def test_write_chapter_saves_collected_draft() -> None:
    h = _make_harness(conflicts=[])
    state = initial_chain_state(project_id=PROJECT, user_id=USER, start_chapter_no=1, count=1)

    out = await write_chapter(
        state,
        gateway_builder=h["gateway_builder"],
        session_factory=h["session_factory"],
        memory_repos_factory=h["memory_repos_factory"],
        chapter_repo_factory=h["chapter_repo_factory"],
        assemble=h["assemble"],
    )

    assert out == {}  # state 不变（正文在 DB）
    assert h["draft_store"][1] == "第 N 章正文。#1"
    # 收集版：write 走非流式 run
    assert h["gateway"].write_calls == 1


# ---- 全图（MemorySaver）：无冲突全自动跑满 ----


async def test_chain_runs_all_chapters_when_no_conflicts() -> None:
    h = _make_harness(conflicts=[])
    graph = _build(h, MemorySaver())
    config: RunnableConfig = {"configurable": {"thread_id": "chain-noconflict"}}
    initial = initial_chain_state(project_id=PROJECT, user_id=USER, start_chapter_no=1, count=2)

    final = await graph.ainvoke(initial, config=config)

    assert "__interrupt__" not in final
    assert final["written"] == [1, 2]
    assert final["current_chapter_no"] == 3  # last(2) + 1，越界 → END
    assert [a["chapter_no"] for a in h["accepted"]] == [1, 2]
    assert h["accepted"][0]["decisions"] == []  # 无冲突章 decisions=[]


# ---- 全图：有冲突 → interrupt → resume → accept → END ----


async def test_chain_interrupts_on_conflict_then_resumes() -> None:
    conflict = Conflict(type="性格漂移", where="第3段", suggestion="统一主角姓名为张三")
    h = _make_harness(conflicts=[conflict])
    saver = MemorySaver()  # 同进程同 saver 跨两次 invoke → interrupt 可恢复
    graph = _build(h, saver)
    config: RunnableConfig = {"configurable": {"thread_id": "chain-conflict"}}
    initial = initial_chain_state(project_id=PROJECT, user_id=USER, start_chapter_no=1, count=1)

    # 初始 run：第 1 章四审报冲突 → interrupt 暂停。
    interrupted = await graph.ainvoke(initial, config=config)

    assert "__interrupt__" in interrupted
    assert h["accepted"] == []  # 尚未验收

    # resume：作者裁决该冲突 → 续跑 accept → END。
    decisions = [{"conflict_index": 0, "verdict": "ignore", "note": "笔误，后续统一"}]
    final = await graph.ainvoke(Command(resume=decisions), config=config)

    assert "__interrupt__" not in final
    assert final["written"] == [1]
    assert [a["chapter_no"] for a in h["accepted"]] == [1]
    assert h["accepted"][0]["decisions"] == decisions  # 裁决透传给 accept_op


# ---- F2 续写式链 continue_volume ----


async def test_write_chapter_draft_volume_omits_prior_text() -> None:
    """draft_volume（默认）写章不注入前文——仍走 assemble 的 volatile（无回归）。"""
    h = _make_harness(conflicts=[])
    h["accepted_store"][0] = "上一章正文末尾片段。"  # 即使有前文也不应注入
    state = initial_chain_state(
        project_id=PROJECT, user_id=USER, start_chapter_no=1, count=1
    )  # 默认 chain_key=draft_volume

    await write_chapter(
        state,
        gateway_builder=h["gateway_builder"],
        session_factory=h["session_factory"],
        memory_repos_factory=h["memory_repos_factory"],
        chapter_repo_factory=h["chapter_repo_factory"],
        assemble=h["assemble"],
    )

    write_input = h["gateway"].write_inputs[0]
    assert write_input == "## 写作指令\n写第 N 章"  # 原 assemble volatile，未改写
    assert "上一章正文末尾片段" not in write_input


async def test_write_chapter_continue_volume_injects_prior_accepted_text() -> None:
    """continue_volume 写第 2 章时，注入第 1 章 accepted 正文（经 build_continuation_context）。"""
    h = _make_harness(conflicts=[])
    h["accepted_store"][1] = "第一章已验收的正文末尾。"
    state = initial_chain_state(
        project_id=PROJECT,
        user_id=USER,
        start_chapter_no=2,
        count=1,
        chain_key="continue_volume",
    )

    await write_chapter(
        state,
        gateway_builder=h["gateway_builder"],
        session_factory=h["session_factory"],
        memory_repos_factory=h["memory_repos_factory"],
        chapter_repo_factory=h["chapter_repo_factory"],
        assemble=h["assemble"],
    )

    write_input = h["gateway"].write_inputs[0]
    assert "第一章已验收的正文末尾。" in write_input
    assert "前文正文" in write_input  # build_continuation_context 的小节标题


async def test_continue_volume_chain_second_chapter_sees_first_chapter_text() -> None:
    """全图 E2E：continue_volume 两章——第 2 章请求上下文含第 1 章正文（不变量 #1/#5）。"""
    h = _make_harness(conflicts=[])
    graph = _build(h, MemorySaver())
    config: RunnableConfig = {"configurable": {"thread_id": "chain-continue"}}
    initial = initial_chain_state(
        project_id=PROJECT,
        user_id=USER,
        start_chapter_no=1,
        count=2,
        chain_key="continue_volume",
    )

    final = await graph.ainvoke(initial, config=config)

    assert "__interrupt__" not in final
    assert final["written"] == [1, 2]
    # 第 1 章写出的正文（accept_op 晋升为 accepted）应出现在第 2 章的写章请求里。
    first_chapter_text = h["accepted_store"][1]
    second_write_input = h["gateway"].write_inputs[1]
    assert first_chapter_text in second_write_input
    # 第 1 章请求无前文（前一章 0 不存在）→ 降级占位，不报错。
    first_write_input = h["gateway"].write_inputs[0]
    assert "前文正文" in first_write_input
