import pytest
from types import SimpleNamespace

from genesis.writer.writer_agent import (
    WriterAgent,
    WRITER_PROMPT_TEMPLATE,
    CHAPTER_OUTPUT_SCHEMA,
)
from genesis.writer.models import GenerationContext, ChapterSpec
from genesis.writer.writer_state import WriterState
from genesis.writer.exceptions import WriterGenerationError


class FakeEngine:
    def __init__(self):
        self.calls = 0

    def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
        self.calls += 1
        return SimpleNamespace(
            data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
            status="ok",
        )


class FakePromptRegistry:
    @staticmethod
    def get_or_create(name, template):
        return SimpleNamespace(name=name, version="1", template=template)


def _ctx(cid, title):
    return GenerationContext(
        chapter_id=cid,
        title=title,
        template_marker=ChapterSpec(chapter_id=cid, title=title, section_placeholder=None),
        structured_source="source text",
        write_rules=["W1"],
        design_rules=["D1"],
        template_styles={"Heading1"},
        prior_state=None,
    )


def test_generate_chapter_success():
    agent = WriterAgent(
        session_id="s",
        engine=FakeEngine(),
        prompt_registry=FakePromptRegistry(),
        state=WriterState(["db_design"]),
    )
    content = agent.generate_chapter(_ctx("db_design", "DB 設計"))
    assert content.chapter_id == "db_design"
    assert content.blocks and content.blocks[0].type == "paragraph"


def test_prompt_template_enforces_title_language():
    # 语言约束：正文需与章节标题语言一致（模板为日文时输出日文）
    assert "语言约束" in WRITER_PROMPT_TEMPLATE
    assert "{{title}}" in WRITER_PROMPT_TEMPLATE


def test_prompt_template_has_impact_context_var():
    # 影响调查结果作为生成主上下文：模板必须包含 impact 变量（无影响书时渲染为空串）
    assert "影响调查上下文" in WRITER_PROMPT_TEMPLATE
    assert "{{impact}}" in WRITER_PROMPT_TEMPLATE


def test_prompt_template_enforces_topic_and_chapter_data():
    # 主题约束：正文必须围绕本章标题主题，依据章节数据，禁止套用系统整体架构
    assert "主题约束" in WRITER_PROMPT_TEMPLATE
    assert "参考资料（本章对应数据）" in WRITER_PROMPT_TEMPLATE
    assert "{{data}}" in WRITER_PROMPT_TEMPLATE
    # 旧的全量 source 变量不再使用
    assert "{{source}}" not in WRITER_PROMPT_TEMPLATE


def test_prompt_template_has_sub_headings_var():
    # 子节结构（design.md §6.5：H2/H3 归入本章）：模板必须提供小节变量并指导按小节组织
    assert "{{sub_headings}}" in WRITER_PROMPT_TEMPLATE
    assert "小节" in WRITER_PROMPT_TEMPLATE


def test_prompt_template_forbids_heading_blocks_when_no_sub_headings():
    # 无子节的章：LLM 不得自造 heading 块（真实试运行暴露：帳票一覧/バッチ一覧
    # 章内出现与章同名的自造 H2）
    assert "不得输出任何 type=heading 的内容块" in WRITER_PROMPT_TEMPLATE


class HeadingEngine(FakeEngine):
    """返回含 heading 块的输出（模拟 LLM 自造标题行为）。"""

    def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
        return SimpleNamespace(
            data={
                "title": variables["title"],
                "blocks": [
                    {"type": "heading", "level": 2, "text": variables["title"]},
                    {"type": "paragraph", "text": "正文内容"},
                ],
            },
            status="ok",
        )


def test_generate_chapter_drops_headings_when_no_sub_headings():
    """无子节结构的章：LLM 自造的 heading 块被程序化剔除（模板结构为准）。"""
    agent = WriterAgent(
        session_id="s", engine=HeadingEngine(),
        prompt_registry=FakePromptRegistry(),
        state=WriterState(["batch_list"]),
    )
    ctx = _ctx("batch_list", "7. バッチ一覧")
    content = agent.generate_chapter(ctx)
    assert all(b.type != "heading" for b in content.blocks)
    assert [b.text for b in content.blocks] == ["正文内容"]


def test_generate_chapter_keeps_headings_when_sub_headings_exist():
    """有子节结构的章：heading 块保留（按小节组织的内容）。"""
    agent = WriterAgent(
        session_id="s", engine=HeadingEngine(),
        prompt_registry=FakePromptRegistry(),
        state=WriterState(["function_list"]),
    )
    ctx = _ctx("function_list", "2. 機能一覧")
    ctx.template_marker.sub_headings = ["2.1 機能一覧表"]
    content = agent.generate_chapter(ctx)
    assert any(b.type == "heading" for b in content.blocks)


class RecordingEngine(FakeEngine):
    """记录 variables 以断言 prompt 渲染输入。"""

    def __init__(self):
        super().__init__()
        self.last_vars = None

    def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
        self.last_vars = dict(variables)
        return super().chat_structured(
            session_id=session_id, prompt=prompt,
            variables=variables, schema=schema, retry_count=retry_count,
        )


class RealStylePromptRegistry:
    """模拟真实 PromptRegistry：无 get_or_create，get() 返回模板字符串（非 Prompt）。"""

    def __init__(self):
        self._tpl = {}

    def register(self, name, version, template):
        self._tpl[(name, version)] = template

    def get(self, name, version=None):
        return self._tpl[(name, version)]


class RenderingEngine:
    """模拟真实引擎渲染行为：仅当收到 Prompt 对象时才用 jinja2 渲染（engine._render_prompt 契约）。"""

    def __init__(self):
        self.last_rendered = None

    def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
        from jinja2 import Template
        from genesis.inference.types import Prompt as P
        tpl = prompt.template if isinstance(prompt, P) else prompt  # str 原样返回 → 不渲染
        self.last_rendered = Template(tpl).render(**variables) if isinstance(prompt, P) else tpl
        return SimpleNamespace(
            data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
            status="ok",
        )


def test_prompt_is_rendered_with_real_style_registry():
    """回归：真实 PromptRegistry.get 返回模板字符串 → 引擎按 str 原样发送，
    {{data}}/{{title}} 等占位符从未被替换（真实试运行暴露）。
    WriterAgent 必须保证交给引擎的是可渲染的 Prompt 对象。"""
    agent = WriterAgent(
        session_id="s",
        engine=RenderingEngine(),
        prompt_registry=RealStylePromptRegistry(),
        state=WriterState(["db_design"]),
    )
    ctx = _ctx("db_design", "DB 設計")
    agent.generate_chapter(ctx)
    assert "{{" not in agent.engine.last_rendered          # 无残留占位符
    assert "DB 設計" in agent.engine.last_rendered         # title 已渲染
    assert "W1" in agent.engine.last_rendered              # write_rules 已渲染


def test_generate_chapter_passes_chapter_scoped_data_var():
    from genesis.data_models import (
        CellValue, ExcelTable, ParsedTemplate, ChapterMarker,
        Provenance, SheetType, StructuredSource,
    )

    def cell(v):
        return CellValue(value=v, provenance=Provenance("f.xlsx", "s", 1, "A", "列"))

    src = StructuredSource(
        tables=[
            ExcelTable("DB定義", SheetType.DATABASE, "openpyxl",
                       ["テーブルID", "テーブル名"],
                       [{"テーブルID": cell("T001"), "テーブル名": cell("trade_order")}]),
            ExcelTable("機能一覧", SheetType.FUNCTION, "openpyxl",
                       ["機能ID", "機能名"],
                       [{"機能ID": cell("F001"), "機能名": cell("止损风控")}]),
        ],
        template=ParsedTemplate("t.docx",
                                [ChapterMarker(type="heading", name="x", level=1)],
                                {}, {"used": []}),
        rule_docs=[], image_analyses=[], existing_system=None, comments=[],
    )
    ctx = GenerationContext(
        chapter_id="db_design", title="DB 設計",
        template_marker=ChapterSpec(chapter_id="db_design", title="DB 設計", section_placeholder=None),
        structured_source=src, write_rules=["W1"], design_rules=["D1"],
        template_styles={"Heading1"},
    )
    engine = RecordingEngine()
    agent = WriterAgent(session_id="s", engine=engine,
                        prompt_registry=FakePromptRegistry(),
                        state=WriterState(["db_design"]))
    agent.generate_chapter(ctx)
    data = engine.last_vars["data"]
    assert "T001" in data and "trade_order" in data   # 本章对应数据
    assert "止损风控" not in data                      # 其他章数据被排除


def test_generate_chapter_retries():
    class Boom(FakeEngine):
        def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
            raise RuntimeError("boom")

    agent = WriterAgent(
        session_id="s",
        engine=Boom(),
        prompt_registry=FakePromptRegistry(),
        state=WriterState(["db_design"]),
        max_retries=2,
    )
    with pytest.raises(WriterGenerationError):
        agent.generate_chapter(_ctx("db_design", "DB 設計"))


class AsyncFakeEngine:
    async def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
        return SimpleNamespace(
            data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
            status="ok",
        )


def test_generate_chapter_with_async_engine():
    agent = WriterAgent(
        session_id="s",
        engine=AsyncFakeEngine(),
        prompt_registry=FakePromptRegistry(),
        state=WriterState(["db_design"]),
    )
    content = agent.generate_chapter(_ctx("db_design", "DB 設計"))
    assert content.chapter_id == "db_design"
    assert content.blocks and content.blocks[0].type == "paragraph"


def test_generate_chapter_propagates_engine_error_detail():
    class FailingEngine:
        def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
            return SimpleNamespace(
                data={},
                status="failed",
                error="LLM HTTP 401: invalid key",
                error_code="LLM_NETWORK_ERROR",
            )

    agent = WriterAgent(
        session_id="s",
        engine=FailingEngine(),
        prompt_registry=FakePromptRegistry(),
        state=WriterState(["db_design"]),
        max_retries=1,
    )
    with pytest.raises(WriterGenerationError) as exc:
        agent.generate_chapter(_ctx("db_design", "DB 設計"))
    assert "LLM HTTP 401" in str(exc.value)
    assert "LLM_NETWORK_ERROR" in str(exc.value)