"""Writer 章节数据 token 预算（OV5 DataGate 接入）。

背景：design.md §6.8 ① 要求 `DataGate.load(structured_source, selector=该章数据)`，
但此前 `_format_chapter_data` 直接渲染、未过 DataGate，且 writer_agent 的
`_chunk_source` 只做「可用性验证」不生效（死代码）。本文件验证：
  - 章节数据超出 token 预算时按「最大表优先」省略，并留下省略标记（不抛错、不阻断生成）
  - 预算可由 GenerationContext.data_token_budget 注入
  - 死代码 _chunk_source 已移除
"""
from genesis.data_models import (
    CellValue,
    ExcelTable,
    ParsedTemplate,
    ChapterMarker,
    Provenance,
    SheetType,
    StructuredSource,
)
from genesis.orchestrator.datagate import DataGate, DataSelector
from genesis.writer.models import ChapterSpec, GenerationContext
from genesis.writer.writer_agent import WriterAgent


def _cell(v):
    return CellValue(value=v, provenance=Provenance("要件.xlsx", "s1", 1, "A", "列"))


def _table(name, sheet_type, headers, rows):
    return ExcelTable(
        name=name,
        detected_type=sheet_type,
        extraction_method="openpyxl",
        headers=headers,
        rows=[{h: _cell(v) for h, v in zip(headers, row)} for row in rows],
    )


def _source(tables):
    template = ParsedTemplate(
        file_name="t.docx",
        sections=[ChapterMarker(type="heading", name="x", level=1)],
        placeholders={},
        styles={"used": []},
    )
    return StructuredSource(
        tables=tables, template=template, rule_docs=[], image_analyses=[],
        existing_system=None, comments=[],
    )


def _ctx(chapter_id, source, budget=None):
    kwargs = {} if budget is None else {"data_token_budget": budget}
    return GenerationContext(
        chapter_id=chapter_id, title="機能一覧",
        template_marker=ChapterSpec(
            chapter_id=chapter_id, title="機能一覧",
            section_placeholder=f"section:{chapter_id}",
        ),
        structured_source=source,
        write_rules=[], design_rules=[], template_styles=set(),
        **kwargs,
    )


def _small_table():
    return _table("機能一覧", SheetType.FUNCTION, ["機能ID", "機能名"], [["F001", "止損"]])


def _big_table():
    rows = [[f"F{i:03d}", "非常に長い説明文" * 20] for i in range(200)]
    return _table("大量機能", SheetType.FUNCTION, ["機能ID", "説明"], rows)


# ---------- 预算常量 ----------

def test_default_data_token_budget_exposed():
    from genesis.writer.models import MAX_CHAPTER_DATA_TOKENS
    assert MAX_CHAPTER_DATA_TOKENS == 8_000


# ---------- 超预算：省略 + 标记（不抛错） ----------

def test_tiny_budget_omits_all_tables_with_marker():
    src = _source([_small_table()])
    data = _ctx("function_list", src, budget=1).to_vars()["data"]
    # 预算=1 → 无表可留 → 仅省略标记，且不抛 DataGateError
    assert data == "" or "省略" in data
    assert "F001" not in data


def test_budget_drops_largest_table_and_keeps_smaller():
    small, big = _small_table(), _big_table()
    src = _source([small, big])
    # 测量用 gate 需足够大（大表单独即超默认 8000 预算，否则测量本身会抛错）
    gate = DataGate(max_total_tokens=10 ** 9)
    small_tokens = gate.load(src, DataSelector(table_ids=["機能一覧"])).token_estimate
    big_tokens = gate.load(src, DataSelector(table_ids=["大量機能"])).token_estimate
    assert big_tokens > small_tokens

    # 预算仅够小表 → 大表被省略，小表内容保留，标记点名被省略的表
    data = _ctx("function_list", src, budget=small_tokens + 1).to_vars()["data"]
    assert "F001" in data and "止損" in data
    assert "大量機能" not in data or "省略" in data
    assert "省略" in data and "大量機能" in data


# ---------- 预算内：原样渲染 ----------

def test_within_budget_renders_without_marker():
    src = _source([_small_table()])
    data = _ctx("function_list", src).to_vars()["data"]
    assert "F001" in data and "止損" in data
    assert "省略" not in data


# ---------- 死代码已移除 ----------

def test_writer_agent_has_no_dead_chunk_source():
    assert not hasattr(WriterAgent, "_chunk_source")