"""QA 闭环：生成 → 校验 → 仅重生成失败章 → 复校验（Phase 5）。"""
from __future__ import annotations

from pathlib import Path

from genesis.inference.engine import InferenceEngine
from genesis.inference.prompt_registry import PromptRegistry
from genesis.qa.guardrails import DEFAULT_MAX_QA_ROUNDS, QALoopController
from genesis.qa.report import QAReport
from genesis.qa.validator import QAValidator
from genesis.writer.context_builder import build_contexts
from genesis.writer.docx_injector import Block, DocxInjector
from genesis.writer.models import ChapterContent
from genesis.writer.orchestrator import _section_id_of
from genesis.writer.renderer import render_chapter_blocks
from genesis.writer.writer_agent import WriterAgent
from genesis.writer.writer_state import WriterState


class QALoop:
    def __init__(self, max_rounds: int = DEFAULT_MAX_QA_ROUNDS) -> None:
        self.controller = QALoopController(max_rounds=max_rounds)

    def _build(self, structured_source, samples_dir, engine, prompt_registry, template_path, output_path, session_id, only_ids=None, prev=None):
        ctxs = build_contexts(structured_source, samples_dir)
        state = WriterState([c.chapter_id for c in ctxs])
        agent = WriterAgent(session_id=session_id, engine=engine, prompt_registry=prompt_registry, state=state)
        contents_map = dict(prev) if prev else {}
        order = [c.chapter_id for c in ctxs]
        sections: dict[str, list[Block]] = {}
        for ctx in ctxs:
            if only_ids is not None and ctx.chapter_id not in only_ids and ctx.chapter_id in contents_map:
                content = contents_map[ctx.chapter_id]
            else:
                content = agent.generate_chapter(ctx)
                contents_map[ctx.chapter_id] = content
            blocks = render_chapter_blocks(content)
            sec_id = _section_id_of(ctx.template_marker.section_placeholder)
            if sec_id:
                sections[sec_id] = blocks
        tpl = template_path or getattr(structured_source.template, "file_name", None)
        if not tpl:
            raise ValueError("template_path 必须提供")
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        doc = DocxInjector(tpl).inject(sections, meta={})
        doc.save(output_path)
        return [contents_map[cid] for cid in order]

    def run(self, structured_source, output_path, session_id="writer", samples_dir="samples", engine=None, prompt_registry=None, template_path=None) -> QAReport:
        engine = engine or InferenceEngine()
        prompt_registry = prompt_registry or PromptRegistry()
        validator = QAValidator()
        contents = self._build(structured_source, samples_dir, engine, prompt_registry, template_path, output_path, session_id)
        report = validator.validate_doc(contents, structured_source)
        while self.controller.can_continue() and report.failed_chapters:
            self.controller.advance()
            contents = self._build(
                structured_source,
                samples_dir,
                engine,
                prompt_registry,
                template_path,
                output_path,
                session_id,
                only_ids=set(report.failed_chapters),
                prev={c.chapter_id: c for c in contents},
            )
            report = validator.validate_doc(contents, structured_source)
        return report