"""文笔层：把已过闸门的编年稿大量扩写成有现场的长文。

与 volume.write_episode 的分工：
- volume 从史料素材写出「骨架正确」的一集（事实优先，1500 字内）；
- expand 拿那一稿逐事件展开（场景优先，目标 8000 字）。

**为什么要分块**：一次调用来写 8000 字，模型会跑偏、提前收尾、漏掉后半段事件。
所以按「年份场景」把原稿切块，每块目标约 2600 字，逐块展开后用上一块结尾承接，
再把各块拼接。每块都只喂它自己那一段史料，避免不同块写出重复内容。

这是 A/B 层：产物写成 `<key>.literary.md`，永不覆盖原稿。
"""
from __future__ import annotations

import math
import re
from collections import Counter
from dataclasses import replace
from pathlib import Path

from dfannals import deslop, volume
from dfannals import episodes as ep_mod
from dfannals.chronicle import cjk_length
from dfannals.config import (
    EXPAND_CEILING,
    EXPAND_CHUNK_TARGET,
    EXPAND_FLOOR,
    EXPAND_TARGET_CHARS,
    LITERARY_MODEL,
    MAX_TOKENS_EXPAND,
    PROMPT_DIR,
)
from dfannals.factcheck import Suspicion
from dfannals.legends import World
from dfannals.llm import chat
from dfannals.normalize import Normalized, normalize
from dfannals.threads import Thread

SPEC_FILE = PROMPT_DIR / "literary-expander.md"
TAIL_CHARS = 300          # 上一块结尾取多少字做承接
YEAR_RE = re.compile(r"^\s*\d{3,4}\s*年")
ANY_YEAR_RE = re.compile(r"(\d{3,4})\s*年")


def scene_groups(draft: str) -> list[list[str]]:
    """按「年份开头的段落」切场景组；紧随其后的非年份段落归入同一组。"""
    groups: list[list[str]] = []
    for para in (p for p in draft.split("\n") if p.strip()):
        if YEAR_RE.match(para) or not groups:
            groups.append([para])
        else:
            groups[-1].append(para)
    return groups


def plan_chunks(draft: str, per_chunk_chars: int) -> list[str]:
    """把原稿装进若干块，尽量在场景边界切割。"""
    groups = scene_groups(draft)
    chunks: list[list[str]] = []
    current: list[str] = []
    size = 0
    for group in groups:
        glen = sum(cjk_length(line) for line in group)
        if current and size + glen > per_chunk_chars:
            chunks.append(current)
            current, size = [], 0
        current += group
        size += glen
    if current:
        chunks.append(current)
    return ["\n".join(c).strip() for c in chunks if "".join(c).strip()]


def _chunk_episode(episode: ep_mod.Episode, chunk_text: str) -> ep_mod.Episode:
    """把这一块覆盖的史料事件挑出来，作为该块的素材。"""
    years = [int(m.group(1)) for m in ANY_YEAR_RE.finditer(chunk_text)]
    if not years:
        return episode
    low, high = min(years), max(years)
    events = [e for e in episode.events if low <= e.year <= high] or episode.events
    return replace(episode, events=events, start_year=low, end_year=high)


def _title_form(value: str) -> str:
    """YETI → Yeti、EAGLE_MAN → Eagle_man：正文里模型就是这么写的。"""
    return "_".join(part[:1].upper() + part[1:].lower() for part in value.split("_"))


def allowed_names(world: World, thread: Thread, episode: ep_mod.Episode) -> list[str]:
    """本集史料里真实出现过的专名。

    这份名单既发给模型（事前防幻觉），也用于修复可疑专名（事后）。
    实测：不给名单时模型造出了 "Farewell" 这个史料里不存在的地名。
    """
    names: set[str] = set()
    heros, others = ep_mod.cast_of(world, thread, episode)
    for fid in [*heros, *others, *thread.members]:
        name = world.figure_name(fid)
        if name:
            names.add(name)
        figure = world.figures.get(fid)
        if figure and figure.race:
            names.add(figure.race)
            names.add(_title_form(figure.race))
    for event in episode.events:
        for ref in world.render_refs(event):
            value = ref.partition("=")[2].strip()
            if value:
                names.add(value)
    return sorted(n for n in names if n)


def chunk_messages(
    world: World,
    thread: Thread,
    episode: ep_mod.Episode,
    total_episodes: int,
    chunk_text: str,
    chunk_target: int,
    index: int,
    count: int,
    prev_tail: str = "",
    elided: Counter | None = None,
) -> list[dict]:
    spec = SPEC_FILE.read_text(encoding="utf-8")
    material = volume.render_material(world, thread, episode, total_episodes, "", elided)
    user = [
        material,
        "",
        f"### 这一块原稿（第 {index}/{count} 块，只写这一块的年代范围）",
        "",
        chunk_text,
        "",
        "### 这一块的 AI 味诊断（这些必须改掉）",
        "",
        deslop.digest(chunk_text),
        "",
    ]
    if prev_tail:
        user += ["### 上一块写到哪儿（用于承接状态，不要复述）", "", prev_tail, ""]
    names = allowed_names(world, thread, episode)
    if names:
        user += [
            "### 允许使用的专名（史料里真实存在，逗号分隔；除此之外的人名地名一律不许写）",
            "",
            "、".join(names),
            "",
            "（上面提到的种族名可以当普通名词用，不要写成新角色。）",
            "",
        ]
    user += [
        (
            f"### 要求：把这一块逐事件展开到 {chunk_target} 字左右，只输出正文。"
            "不要写标题，不要重复上一块已经写过的内容，也不要跳到下一块的年代。"
        ),
    ]
    return [{"role": "system", "content": spec}, {"role": "user", "content": "\n".join(user)}]


def _fit_length(
    messages: list[dict],
    text: str,
    *,
    floor: int,
    ceiling: int,
    target: int,
    gateway=None,
    model: str = LITERARY_MODEL,
    max_repairs: int = 1,
) -> str:
    """把长度拉进区间：太短继续展开，太长压缩（都不许动史实）。"""
    for _ in range(max_repairs):
        length = cjk_length(text)
        if floor <= length <= ceiling:
            break
        ask = (
            f"只有 {length} 字，还差很多。继续展开：把这一段里还没落到现场的事件逐件写成场景，"
            f"补对白、动作、视角人物此刻的感知与算计，写到 {target} 字左右。"
            "不许新增史料里没有的事件或专名，也不要用形容词灌水。"
            if length < floor
            else f"有 {length} 字，超了。删掉重复交代和没有功能的描写，压到 {ceiling} 字以内，"
                 "冲突与转折的现场感要保留。"
        )
        text = chat(messages + [{"role": "assistant", "content": text},
                                {"role": "user", "content": ask}],
                    gateway=gateway, model=model, max_tokens=MAX_TOKENS_EXPAND).content
    return text.strip()


def build_messages(
    world: World,
    thread: Thread,
    episode: ep_mod.Episode,
    total_episodes: int,
    draft: str,
    prev_tail: str = "",
    elided: Counter | None = None,
    target: int = EXPAND_TARGET_CHARS,
) -> list[dict]:
    """整集一次写完时用（目标不长时）：素材 + 原稿 + 诊断 + 字数目标。"""
    spec = SPEC_FILE.read_text(encoding="utf-8")
    material = volume.render_material(world, thread, episode, total_episodes, prev_tail, elided)
    user = [
        material,
        "",
        "### 原稿（事实与事件顺序按它，写法不要学它）",
        "",
        draft.strip(),
        "",
        "### 原稿的 AI 味诊断（这些必须改掉）",
        "",
        deslop.digest(draft),
        "",
        (
            f"### 现在开始：按原稿的事件顺序逐件展开，写到 {target} 字左右"
            f"（{EXPAND_FLOOR}–{EXPAND_CEILING} 字以内），只输出正文。"
        ),
    ]
    return [{"role": "system", "content": spec}, {"role": "user", "content": "\n".join(user)}]


def expand_episode(
    world: World,
    thread: Thread,
    episode: ep_mod.Episode,
    total_episodes: int,
    draft: str,
    prev_tail: str = "",
    *,
    gateway=None,
    model: str = LITERARY_MODEL,
    target: int = EXPAND_TARGET_CHARS,
    elided: Counter | None = None,
    max_repairs: int = 1,
) -> Normalized:
    """大量扩写：分块展开 → 拼接 → 规范化标点。

    返回 Normalized，notes 里带上每块的实际字数，便于发现哪一块没写够。
    """
    # 标题不进正文块：否则模型会把 "# 守门人" 当正文写进去。最后再拼回开头。
    lines = draft.split("\n")
    headings = [ln for ln in lines if ln.strip().startswith("#")]
    body = "\n".join(ln for ln in lines if not ln.strip().startswith("#")).strip()

    length = cjk_length(body) or 1
    blocks = max(1, math.ceil(target / EXPAND_CHUNK_TARGET))
    chunks = plan_chunks(body, max(1, round(length / blocks)))
    chunk_target = max(700, target // len(chunks))

    parts: list[str] = []
    notes: list[str] = []
    for index, chunk in enumerate(chunks, 1):
        prev_tail_text = parts[-1][-TAIL_CHARS:] if parts else prev_tail[-TAIL_CHARS:]
        sub_episode = _chunk_episode(episode, chunk)
        messages = chunk_messages(world, thread, sub_episode, total_episodes, chunk,
                                  chunk_target, index, len(chunks), prev_tail_text, elided)
        text = chat(messages, gateway=gateway, model=model,
                    max_tokens=MAX_TOKENS_EXPAND).content
        text = _fit_length(messages, text,
                           floor=int(chunk_target * 0.75), ceiling=int(chunk_target * 1.35),
                           target=chunk_target, gateway=gateway, model=model,
                           max_repairs=max_repairs)
        parts.append(text)
        notes.append(f"第 {index}/{len(chunks)} 块：{cjk_length(text)} 字（目标 {chunk_target}）")

    stitched = "\n\n".join(parts)
    total_now = cjk_length(stitched)

    # 拼接后总量仍不足，再整篇催一次（只此一次，避免无限加长）
    if total_now < EXPAND_FLOOR:
        messages = build_messages(world, thread, episode, total_episodes, body, prev_tail,
                                  elided, target)
        stitched = _fit_length(messages, stitched, floor=EXPAND_FLOOR, ceiling=EXPAND_CEILING,
                               target=target, gateway=gateway, model=model, max_repairs=1)

    result = normalize(stitched.strip())
    result.notes = notes + result.notes
    if headings:
        result.text = "\n\n".join([*headings, result.text])
    return result


def literary_path(directory: Path, episode: ep_mod.Episode) -> Path:
    """扩写稿的落盘位置：与原稿并列，便于 A/B 对照。"""
    return directory / f"{episode.key}.literary.md"


NAME_FIX_SPEC = """你在做一件很窄的事：正文里出现了史料中不存在的专名（模型凭空造的人名、地名、称号）。

把包含这些专名的句子改掉：
- 优先换成「允许使用的专名」里真实存在、语义接近的名字；
- 没有合适的名字，就把句子改成不指名（例如把"把 Farewell 挡回去"改成"把下一个要拦的人挡回去"）。

不许增删情节，不许改动没被点到的句子，不许把段落合并或拆分。
只输出改好的段落，每段以 [序号] 开头，序号与段数必须和输入一致。不写任何解释。
"""


_NUMBERED_RE = re.compile(r"^\s*\[(\d+)\]\s*(.+?)\s*$")


def _parse_numbered(raw: str) -> dict[int, str]:
    """把模型返回的 `[n] 段落` 解析成 {序号: 段落}，容忍跨行。"""
    out: dict[int, str] = {}
    current: int | None = None
    buffer: list[str] = []
    for line in raw.split("\n"):
        match = _NUMBERED_RE.match(line)
        if match:
            if current is not None:
                out[current] = "\n".join(buffer).strip()
            current = int(match.group(1))
            buffer = [match.group(2)]
        elif current is not None and line.strip():
            buffer.append(line.strip())
    if current is not None:
        out[current] = "\n".join(buffer).strip()
    return {k: v for k, v in out.items() if v}


def repair_names(
    world: World,
    thread: Thread,
    episode: ep_mod.Episode,
    text: str,
    suspects: list[Suspicion],
    *,
    gateway=None,
    model: str = LITERARY_MODEL,
) -> Normalized:
    """只重写含可疑专名的那几段，其他段落一字不动。

    整篇重写代价高（几千字）且模型会顺手改别的地方；段级手术安全得多——
    段数对不上就不写回，宁可不改也不把全文换掉。
    """
    bad = {s.name for s in suspects}
    paragraphs = text.split("\n")
    targets = [i for i, para in enumerate(paragraphs) if any(name in para for name in bad)]
    if not targets:
        return normalize(text)

    numbered = "\n".join(f"[{n}] {paragraphs[i]}" for n, i in enumerate(targets, 1))
    messages = [
        {"role": "system", "content": NAME_FIX_SPEC},
        {"role": "user", "content": "\n".join([
            "### 允许使用的专名", "", "、".join(allowed_names(world, thread, episode)), "",
            "### 史料里找不到的专名", "",
            *[f"- {s.name}（出现 {s.count} 次）｜例：…{s.context}…" for s in suspects], "",
            f"### 需要修的段落（共 {len(targets)} 段）", "", numbered, "",
            f"### 要求：只输出改好的 {len(targets)} 段，每段以 [序号] 开头。",
        ])},
    ]
    reply = chat(messages, gateway=gateway, model=model, max_tokens=MAX_TOKENS_EXPAND).content
    fixed = _parse_numbered(reply)

    applied = 0
    for n, index in enumerate(targets, 1):
        if n in fixed:
            paragraphs[index] = fixed[n]
            applied += 1

    result = normalize("\n".join(paragraphs))
    result.notes = [f"可疑专名修复：改写 {applied}/{len(targets)} 段"] + result.notes
    return result


def _hits(text: str) -> tuple[int, float, int]:
    """返回（硬伤数、每千字硬伤密度、字数）。"""
    hard = [h for h in deslop.check(text) if h.level == "hard"]
    count = sum(h.count for h in hard)
    length = cjk_length(text)
    return count, count * 1000 / (length or 1), length


def compare(draft: str, literary: str) -> str:
    """A/B 对照：字数与 AI 味硬伤密度。"""
    rows = ["A/B 对照："]
    for label, text in (("原稿", draft), ("扩写稿", literary)):
        count, density, length = _hits(text)
        rows.append(f"  {label:<8}{length:>6} 字｜硬伤 {count:>4} 处（{density:.1f}/千字）")
    before = _hits(draft)[2]
    if before:
        rows.append(f"  篇幅倍数：{_hits(literary)[2] / before:.1f}×")
    return "\n".join(rows)


def full_report(draft: str, literary: str) -> str:
    """给 CLI 用：A/B 对照 + 两稿各自的诊断。"""
    return "\n".join([
        compare(draft, literary),
        "",
        "— 原稿诊断 —",
        deslop.report(draft),
        "",
        "— 扩写稿诊断 —",
        deslop.report(literary),
    ])