"""Chapter brief aggregation and Chinese chapter assembly."""

from __future__ import annotations

import json
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Callable

from scripts.runtime.roles import RoleDefinition, RuntimeProfile
from scripts.runtime.skills import SkillRegistry
from scripts.runtime.tasks import load_task_cards, validate_packet
from scripts.runtime.workers import ChatClient


def validate_chapter_brief(brief: dict) -> None:
    required = {
        "chapter_id",
        "chapter_title",
        "packet_ids",
        "core_claims",
        "evidence_items",
        "counter_evidence",
        "source_ids",
        "open_questions",
        "assembly_notes",
    }
    missing = sorted(required - set(brief))
    if missing:
        raise ValueError(f"chapter brief missing fields: {missing}")
    if not brief["chapter_id"]:
        raise ValueError("chapter_id required")
    if not brief["packet_ids"]:
        raise ValueError("chapter brief requires at least one packet")
    if not brief["core_claims"]:
        raise ValueError("chapter brief requires core_claims")
    if not brief["evidence_items"]:
        raise ValueError("chapter brief requires evidence_items")
    if not brief["counter_evidence"]:
        raise ValueError("chapter brief requires counter_evidence")


def validate_chapter_markdown_citations(markdown: str, brief: dict) -> None:
    validate_chapter_brief(brief)
    cited = set(re.findall(r"\[(src_[A-Za-z0-9_-]+)\]", markdown))
    allowed = set(brief.get("source_ids") or [])
    unknown = sorted(cited - allowed)
    if unknown:
        raise ValueError(f"unknown citation ids in {brief['chapter_id']}: {unknown}")


def _chapter_title_from_id(chapter_id: str) -> str:
    try:
        index = int(chapter_id.replace("ch", ""))
        return f"第{index}章"
    except ValueError:
        return chapter_id


def build_chapter_briefs(project_root: Path) -> list[dict]:
    cards = load_task_cards(project_root / "phase2" / "task_cards.json")
    grouped: dict[str, list[tuple[str, dict]]] = {}
    for card in cards:
        packet_path = project_root / card.output_packet
        if not packet_path.exists():
            continue
        packet = json.loads(packet_path.read_text(encoding="utf-8"))
        validate_packet(packet)
        for chapter_id in card.chapter_ids:
            grouped.setdefault(chapter_id, []).append((card.task_id, packet))

    briefs: list[dict] = []
    out_dir = project_root / "phase2" / "chapter_briefs"
    out_dir.mkdir(parents=True, exist_ok=True)
    for chapter_id in sorted(grouped):
        packet_pairs = sorted(grouped[chapter_id], key=lambda item: item[0])
        packet_ids = [item[0] for item in packet_pairs]
        packets = [item[1] for item in packet_pairs]
        source_ids = sorted({sid for packet in packets for sid in packet.get("source_ids", [])})
        brief = {
            "chapter_id": chapter_id,
            "chapter_title": _chapter_title_from_id(chapter_id),
            "packet_ids": packet_ids,
            "core_claims": [claim for packet in packets for claim in packet.get("claims", [])],
            "evidence_items": [item for packet in packets for item in packet.get("evidence_items", [])],
            "counter_evidence": [item for packet in packets for item in packet.get("counter_evidence", [])],
            "source_ids": source_ids,
            "open_questions": [q for packet in packets for q in packet.get("open_questions", [])],
            "assembly_notes": [
                "用中文写正式章节，英文仅保留在必要的来源标题、原文摘录、DOI/URL 中。",
                "避免碎片化：不要按 packet 逐段堆砌，要先提炼本章主线，再组织证据。",
                "每个事实、数字和关键判断都必须保留 [src_xxx] 引用。",
                "必须纳入 counter_evidence，并说明它如何影响结论置信度。",
            ],
        }
        validate_chapter_brief(brief)
        (out_dir / f"{chapter_id}.json").write_text(
            json.dumps(brief, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        briefs.append(brief)
    return briefs


def build_chapter_user_prompt(brief: dict) -> str:
    return (
        "请根据以下 chapter brief 写一章正式中文 Markdown 正文。\n"
        "目标是形成一个完整章节，而不是 packet 摘要。避免碎片化，按金字塔结构组织：章首先给结论，再用证据支撑。\n"
        "要求：标题必须是观点型判断；每个数字和事实保留 [src_xxx]；纳入反方证据；不要出现调度元数据。\n"
        "禁止写空泛咨询腔。每个二级小节都必须至少落下 2 个具体审计发现、法规要求、SOP/记录/参数/现场观察或整改证据；不要只写原则。\n"
        "正文末尾必须增加“证据落点与待补证据”小节，用表格列出：关键判断、已使用证据 source_id、已落地整改动作、仍缺证据。若证据不足，直接标注需回炉 Phase 2，不要用泛泛表述补齐。\n"
        "只输出 Markdown，不要输出解释。\n\n"
        f"{json.dumps(brief, ensure_ascii=False, indent=2)}"
    )


class ChapterAssemblyWorker:
    def __init__(
        self,
        *,
        role: RoleDefinition,
        client: ChatClient,
        skill_registry: SkillRegistry | None = None,
    ) -> None:
        self.role = role
        self.client = client
        self.skill_registry = skill_registry or SkillRegistry()

    def _system_prompt(self) -> str:
        skill_texts = []
        for name in self.role.skills:
            try:
                skill_texts.append(f"# Skill: {name}\n\n{self.skill_registry.read(name)}")
            except FileNotFoundError:
                skill_texts.append(f"# Skill: {name}\n\n[missing skill: {name}]")
        return (
            "你是 Deep Research v0.20 的中文章节组装 worker。\n"
            "你的职责是把结构化证据包收束成连贯章节，解决并发研究造成的碎片化。\n"
            "不得编造来源，不得删除关键反方证据。\n\n"
            + "\n\n".join(skill_texts)
        )

    def write_chapter(self, *, project_root: Path, brief: dict) -> Path:
        validate_chapter_brief(brief)
        markdown = self.client.chat_complete(
            model=self.role.model,
            system=self._system_prompt(),
            user=build_chapter_user_prompt(brief),
            temperature=self.role.temperature,
            max_tokens=self.role.max_tokens,
            tag=f"chapter:{brief['chapter_id']}",
        )
        validate_chapter_markdown_citations(markdown, brief)
        out = project_root / "phase2" / "drafts" / f"{brief['chapter_id']}.md"
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(markdown.rstrip() + "\n", encoding="utf-8")
        return out


def _write_chapter_error(project_root: Path, brief: dict, error: Exception) -> None:
    path = project_root / "phase2" / "chapter_errors" / f"{brief.get('chapter_id', 'unknown')}.json"
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "chapter_id": brief.get("chapter_id"),
        "status": "failed",
        "error": str(error),
    }
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def run_chapter_assembly_workers(
    *,
    project_root: Path,
    briefs: list[dict],
    runtime: RuntimeProfile,
    client_factory: Callable[[RoleDefinition], ChatClient],
    workers: int,
) -> int:
    role = runtime.role_for_task("chapter_assembly")
    max_workers = max(1, min(workers, role.max_concurrency))

    def run_one(brief: dict) -> tuple[dict, Path | None, Exception | None]:
        try:
            worker = ChapterAssemblyWorker(role=role, client=client_factory(role))
            return brief, worker.write_chapter(project_root=project_root, brief=brief), None
        except Exception as error:
            return brief, None, error

    written = 0
    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = [pool.submit(run_one, brief) for brief in briefs]
        for future in as_completed(futures):
            brief, path, error = future.result()
            if error is not None:
                _write_chapter_error(project_root, brief, error)
                continue
            if path is None:
                _write_chapter_error(project_root, brief, RuntimeError("chapter worker returned no output path"))
                continue
            written += 1
    return written