src/services/drafting.py

import asyncio
import hashlib
import io
import json
import re
import threading
import unicodedata
from collections import OrderedDict
from collections.abc import AsyncIterator
from concurrent.futures import ThreadPoolExecutor
from typing import Any

from docx import Document
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml.ns import qn
from docx.shared import Pt
from langchain_core.messages import HumanMessage, SystemMessage

from src.models.schemas import (
    DraftCheckIssue,
    DraftExportRequest,
    DraftPlanningQuestion,
    DraftPlanningResult,
    DraftRequest,
    DraftResponse,
    DraftTemplateSource,
    RetrieveCitation,
)
from src.services.llm import get_llm
from src.services.retrieval_access import (
    RetrievalAccessCoordinator,
    WorkspaceRequest,
    source_access_context,
)
from src.services.retrieval_common import (
    DEFAULT_MAX_CHUNKS_PER_DOCUMENT,
    _dict_to_source,
    _merge_sources,
    reciprocal_rank_fusion_lists,
    simplify_legal_query,
)

# Văn bản hành chính theo Nghị định 30/2020/NĐ-CP dùng phông Times New Roman.
DOCX_FONT_NAME = "Times New Roman"

# Bound the pre-streaming phase so a slow/hung retrieval backend (Chroma,
# embeddings, BM25) cannot block draft generation indefinitely. On timeout we
# degrade gracefully — generate from whatever sources we already have.
DRAFT_RETRIEVAL_QUERY_TIMEOUT_S = 20.0
DRAFT_PLANNING_TIMEOUT_S = 30.0
# Draft căn cứ should be DISTINCT documents — one chunk per document is enough to cite a
# legal basis, so we cap at 1 and over-fetch candidates so the final top_k stays diverse.
DRAFT_MAX_CHUNKS_PER_DOCUMENT = 1
_DRAFT_CANDIDATE_MULTIPLIER = 3
# Bound the thread pool that fans out the per-sub-query retrieval calls.
_DRAFT_RETRIEVAL_MAX_WORKERS = 5

# In-process LRU cache for the grounding check. The check is a 20-40s LLM call that is
# deterministic in its inputs, so an identical (requirement, generated html, citations,
# template) re-check — e.g. a re-render or retry — can be served instantly. Bounded to
# avoid unbounded growth; cleared on process restart.
_DRAFT_CHECK_CACHE_MAX = 128
_draft_check_cache: "OrderedDict[str, tuple[str, int, list[DraftCheckIssue]]]" = OrderedDict()
_draft_check_cache_lock = threading.Lock()
# Max wait for the next streamed token. If the LLM stalls (provider buffering,
# network drop) we stop instead of hanging the connection forever.
DRAFT_STREAM_IDLE_TIMEOUT_S = 60.0

OFFICIAL_TEMPLATE_SOURCE = DraftTemplateSource(
    title="Nghị định 30/2020/NĐ-CP",
    url="https://vanban.chinhphu.vn/?docid=199378&pageid=27160",
    note=(
        "Thể thức, kỹ thuật trình bày theo Phụ lục I; bảng chữ viết tắt và mẫu "
        "trình bày văn bản hành chính theo Phụ lục III."
    ),
)


TEMPLATE_GUIDES = {
    "cong_van": {
        "name": "CÔNG VĂN",
        "summary": "Văn bản hành chính dùng để trao đổi, đề nghị, hướng dẫn hoặc trả lời công việc.",
        "required_blocks": [
            "Tên cơ quan, tổ chức chủ quản và tên cơ quan, tổ chức ban hành.",
            "Số, ký hiệu văn bản theo bảng chữ viết tắt tại Phụ lục III.",
            "Địa danh và thời gian ban hành văn bản.",
            "Tên loại hoặc trích yếu nội dung bắt đầu bằng 'V/v ...' nếu phù hợp.",
            "Dòng 'Kính gửi:' đúng đối tượng nhận.",
            "Nội dung nêu căn cứ, bối cảnh, yêu cầu/đề nghị và thời hạn thực hiện nếu có.",
            "Nơi nhận và khối chữ ký của người có thẩm quyền.",
        ],
    },
    "to_trinh": {
        "name": "TỜ TRÌNH",
        "summary": "Văn bản trình cấp có thẩm quyền xem xét, phê duyệt hoặc quyết định một vấn đề.",
        "required_blocks": [
            "Tên cơ quan, tổ chức chủ quản và tên cơ quan, tổ chức ban hành.",
            "Số, ký hiệu văn bản theo Phụ lục III.",
            "Địa danh và thời gian ban hành văn bản.",
            "Tên loại 'TỜ TRÌNH' và trích yếu nội dung.",
            "Dòng 'Kính gửi:' nêu đúng cơ quan/người có thẩm quyền xem xét.",
            "Nội dung gồm sự cần thiết, căn cứ pháp lý/thực tiễn, nội dung trình, đề xuất/kiến nghị.",
            "Nơi nhận và khối chữ ký của người có thẩm quyền.",
        ],
    },
    "thong_bao": {
        "name": "THÔNG BÁO",
        "summary": "Văn bản truyền đạt thông tin, kết luận, nội dung triển khai hoặc yêu cầu thực hiện.",
        "required_blocks": [
            "Tên cơ quan, tổ chức chủ quản và tên cơ quan, tổ chức ban hành.",
            "Số, ký hiệu văn bản theo Phụ lục III.",
            "Địa danh và thời gian ban hành văn bản.",
            "Tên loại 'THÔNG BÁO' và trích yếu nội dung.",
            "Nội dung thông báo nêu rõ sự việc, đối tượng áp dụng, thời gian, địa điểm, yêu cầu thực hiện nếu có.",
            "Nơi nhận và khối chữ ký của người có thẩm quyền.",
        ],
    },
    "quyet_dinh_noi_bo": {
        "name": "QUYẾT ĐỊNH",
        "summary": "Văn bản quyết định nội bộ, phải có phần căn cứ và điều khoản thi hành.",
        "required_blocks": [
            "Tên cơ quan, tổ chức chủ quản và tên cơ quan, tổ chức ban hành.",
            "Số, ký hiệu văn bản theo Phụ lục III.",
            "Địa danh và thời gian ban hành văn bản.",
            "Tên loại 'QUYẾT ĐỊNH' và dòng 'Về việc ...'.",
            "Các căn cứ ban hành, mỗi căn cứ một đoạn, dùng văn bản pháp lý đã truy xuất khi có.",
            "Cụm 'QUYẾT ĐỊNH:' trước phần điều khoản.",
            "Các điều khoản: Điều 1 nội dung quyết định; Điều 2 tổ chức thực hiện; Điều cuối hiệu lực/trách nhiệm thi hành.",
            "Nơi nhận và khối chữ ký của người có thẩm quyền.",
        ],
    },
}


def _template_guide(template_key: str) -> str:
    template = TEMPLATE_GUIDES.get(template_key, TEMPLATE_GUIDES["cong_van"])
    required_blocks = "\n".join(f"- {block}" for block in template["required_blocks"])
    return f"""Loại văn bản: {template["name"]}
Mục đích: {template["summary"]}
Nguồn mẫu bắt buộc: {OFFICIAL_TEMPLATE_SOURCE.title}, {OFFICIAL_TEMPLATE_SOURCE.note}
Checklist bố cục phải có:
{required_blocks}"""


def plan_draft(request: DraftRequest) -> DraftPlanningResult:
    guide = _template_guide(request.template_key.value)
    system_prompt = f"""Bạn là agent lập kế hoạch soạn thảo văn bản pháp luật/hành chính Việt Nam. Phân tích yêu cầu TRƯỚC khi truy xuất VBPL; không bịa căn cứ, không viết bản thảo ở bước này.
{guide}

CHECKLIST XÁC MINH BẮT BUỘC trước khi cho ready (mỗi điểm còn thiếu là một câu hỏi xác minh):
1. Thẩm quyền ký / cấp phê duyệt.
2. Đối tượng áp dụng (toàn đơn vị hay bộ phận cụ thể).
3. Phạm vi nội dung cần thông báo/trình/quyết định/xin ý kiến.
4. Thời điểm hiệu lực hoặc thời hạn thực hiện.
5. Nơi nhận và đơn vị chịu trách nhiệm.
6. Căn cứ cần dùng (VBPL công khai, quy chế nội bộ, tài liệu riêng).
7. Kết quả mong muốn.

Hỏi từng điểm thiếu, mỗi câu ngắn và độc lập, kèm đúng 3 phương án sát câu hỏi (không thêm "Khác" vì UI đã có ô tự nhập).
Không trả ready chỉ vì có thể tô vàng phần thiếu. Nếu yêu cầu chưa xác nhận đủ các điểm trên thì phải đặt ít nhất một câu hỏi xác minh; chỉ trả ready khi đủ dữ kiện hoặc người dùng đã trả lời/skip xác minh.

Trả về DUY NHẤT JSON hợp lệ với schema:
{{
  "status": "ready" | "needs_clarification" | "needs_breakdown" | "needs_grill",
  "reasoning_summary": "tóm tắt ngắn lý do",
  "needs_retrieval": true | false,
  "questions": ["câu hỏi nếu còn thiếu dữ kiện"],
  "question_options": [{{"question": "câu hỏi", "options": ["A", "B", "C"]}}],
  "breakdown": ["vấn đề cần tách để tra cứu"],
  "retrieval_queries": ["query cụ thể để retrieve VBPL"],
  "suggested_prompt": "yêu cầu đã làm rõ/tối ưu nếu có"
}}

Chọn status: needs_clarification khi thiếu dữ kiện trọng yếu (đối tượng, thời gian, thẩm quyền, nội dung, căn cứ); needs_breakdown khi nhiều vấn đề pháp lý độc lập (tạo breakdown + retrieval_queries); needs_grill khi mơ hồ/rủi ro cao/dễ sai thẩm quyền; ready khi đủ dữ kiện (tạo 2-5 retrieval_queries cụ thể).
needs_retrieval=false chỉ khi chỉ cần thể thức/mẫu chung hoặc chỉnh câu chữ; =true nếu có kết luận pháp lý, căn cứ, quyền/nghĩa vụ, thời hạn, chế tài, hiệu lực hoặc viện dẫn VBPL.
"""
    human_prompt = f"""Loại văn bản: {request.template_key.value}
Đơn vị ban hành: {request.issuing_unit}
Nơi nhận: {request.recipient}
Yêu cầu người dùng: {request.draft_requirement}
"""
    # Planner emits structured JSON — temperature=0 is faster and far less likely
    # to produce malformed output that falls back to the default plan. The light model
    # is enough for this rule-following classification and cuts planning latency.
    response = get_llm(temperature=0, light=True).invoke(
        [
            SystemMessage(content=system_prompt),
            HumanMessage(content=human_prompt),
        ]
    )
    return _parse_planning_result(str(response.content), request)


def _parse_planning_result(content: str, request: DraftRequest) -> DraftPlanningResult:
    text = content.strip()
    if text.startswith("```json"):
        text = text[7:]
    if text.startswith("```"):
        text = text[3:]
    if text.endswith("```"):
        text = text[:-3]
    text = text.strip()

    try:
        data = json.loads(text)
    except json.JSONDecodeError:
        data = {
            "status": "ready",
            "reasoning_summary": "Không đọc được JSON planning, dùng yêu cầu gốc để truy xuất.",
            "needs_retrieval": True,
            "questions": [],
            "breakdown": [],
            "retrieval_queries": [request.draft_requirement],
            "suggested_prompt": request.draft_requirement,
        }

    status = data.get("status", "ready")
    if status not in {"ready", "needs_clarification", "needs_breakdown", "needs_grill"}:
        status = "ready"
    queries = _clean_string_list(data.get("retrieval_queries"))
    questions = _clean_string_list(data.get("questions"))
    if status == "ready" and not queries:
        queries = [request.draft_requirement]

    return DraftPlanningResult(
        status=status,
        reasoning_summary=str(data.get("reasoning_summary", "")).strip(),
        needs_retrieval=bool(data.get("needs_retrieval", True)),
        questions=questions,
        question_options=_normalize_question_options(data.get("question_options"), questions),
        breakdown=_clean_string_list(data.get("breakdown")),
        retrieval_queries=queries[:5],
        suggested_prompt=str(data.get("suggested_prompt", "")).strip(),
    )


def _normalize_question_options(value: Any, questions: list[str]) -> list[DraftPlanningQuestion]:
    raw_items = value if isinstance(value, list) else []
    normalized: list[DraftPlanningQuestion] = []

    for index, question in enumerate(questions):
        raw_item = raw_items[index] if index < len(raw_items) and isinstance(raw_items[index], dict) else {}
        raw_options = _clean_string_list(raw_item.get("options"))
        options = raw_options[:3]
        fallback_options = _fallback_options_for_question(question)
        for fallback in fallback_options:
            if len(options) >= 3:
                break
            if fallback not in options:
                options.append(fallback)
        normalized.append(DraftPlanningQuestion(question=question, options=options[:3]))

    return normalized


def _fallback_options_for_question(question: str) -> list[str]:
    normalized = _normalize_vietnamese_text(question)
    if any(keyword in normalized for keyword in ["hieu luc", "thoi diem", "ngay nao", "khi nao", "thoi gian"]):
        return [
            "Có hiệu lực từ ngày ký",
            "Có hiệu lực từ một ngày cụ thể",
            "Áp dụng cho một giai đoạn cụ thể",
        ]
    if any(keyword in normalized for keyword in ["tham quyen", "nguoi ky", "ai ky", "phe duyet", "ban hanh"]):
        return [
            "Thủ trưởng đơn vị ký ban hành",
            "Trưởng bộ phận chuyên môn ký",
            "Trình cấp có thẩm quyền phê duyệt trước",
        ]
    if any(keyword in normalized for keyword in ["can cu", "co so phap ly", "van ban phap ly", "quy dinh nao"]):
        return [
            "Dùng căn cứ VBPL công khai",
            "Dùng thêm quy chế nội bộ",
            "Kết hợp VBPL công khai và tài liệu nội bộ",
        ]
    if any(keyword in normalized for keyword in ["kinh gui", "noi nhan", "gui ai", "nguoi nhan"]):
        return [
            "Gửi đúng nơi nhận đã nhập",
            "Gửi các đơn vị trực thuộc liên quan",
            "Gửi cơ quan cấp trên hoặc cơ quan phối hợp",
        ]
    if any(keyword in normalized for keyword in ["pham vi", "ap dung", "doi tuong", "don vi nao", "cho ai"]):
        return [
            "Áp dụng cho toàn bộ đơn vị",
            "Áp dụng cho một phòng ban cụ thể",
            "Áp dụng cho đơn vị trực thuộc liên quan",
        ]
    return [
        "Giữ theo thông tin người dùng đã nhập",
        "Làm rõ bằng câu trả lời tự nhập",
        "Chưa xác định, yêu cầu planner hỏi lại cụ thể hơn",
    ]


def _normalize_vietnamese_text(value: str) -> str:
    text = unicodedata.normalize("NFD", value.lower())
    without_marks = "".join(char for char in text if unicodedata.category(char) != "Mn")
    return without_marks.replace("đ", "d")


def _clean_string_list(value: Any) -> list[str]:
    if not isinstance(value, list):
        return []
    cleaned: list[str] = []
    for item in value:
        text = str(item or "").strip()
        if text and text not in cleaned:
            cleaned.append(text)
    return cleaned


def generate_draft(
    request: DraftRequest,
    coordinator: RetrievalAccessCoordinator,
    current_user: Any,
    guest_session_id: str | None,
) -> DraftResponse:
    workspace = request.workspace or WorkspaceRequest(type="public_only")
    planning = request.planning or plan_draft(request)
    if planning.status != "ready" and not request.force_generate_with_highlights:
        return DraftResponse(
            html="",
            citations=[],
            warnings=[
                "Cần làm rõ yêu cầu trước khi truy xuất VBPL và soạn thảo."
            ],
            template_source=OFFICIAL_TEMPLATE_SOURCE,
            planning=planning,
        )

    if planning.needs_retrieval:
        sources, warning = _retrieve_planned_sources(
            request=request,
            coordinator=coordinator,
            workspace=workspace,
            current_user=current_user,
            guest_session_id=guest_session_id,
            planning=planning,
        )
    else:
        sources, warning = [], ""

    system_msg, human_msg = _build_draft_messages(request, planning, sources)
    llm = get_llm()
    response = llm.invoke([system_msg, human_msg])
    html_content = _strip_html_fence(response.content.strip())

    citations = _build_citations(sources)
    warnings_list = [warning] if warning else []

    return DraftResponse(
        html=html_content,
        citations=citations,
        warnings=warnings_list,
        template_source=OFFICIAL_TEMPLATE_SOURCE,
        planning=planning,
    )


def _build_draft_messages(
    request: "DraftRequest",
    planning: "DraftPlanningResult",
    sources: list[dict[str, Any]],
) -> tuple[Any, Any]:
    """Return (SystemMessage, HumanMessage) for draft generation."""
    context_str = "\n\n".join(
        [f"[{i+1}] {s.get('metadata', {}).get('title', 'Unknown')}: {s.get('content', '')[:500]}" for i, s in enumerate(sources)]
    )
    guide = _template_guide(request.template_key.value)
    clarification_prompt = _clarification_highlight_prompt(planning, request.force_generate_with_highlights)
    retrieval_prompt = _retrieval_usage_prompt(planning)

    system_prompt = f"""Bạn là một chuyên viên pháp lý cấp cao của nhà nước Việt Nam, am hiểu sâu sắc và có nhiệm vụ soạn thảo văn bản hành chính tuân thủ TUYỆT ĐỐI theo quy định về thể thức và kỹ thuật trình bày văn bản hành chính tại Nghị định 30/2020/NĐ-CP.
Nguồn mẫu chính thức: {OFFICIAL_TEMPLATE_SOURCE.title} - {OFFICIAL_TEMPLATE_SOURCE.url}
Phải bám Phụ lục I về thể thức, kỹ thuật trình bày văn bản hành chính và Phụ lục III về bảng chữ viết tắt tên loại, mẫu trình bày văn bản hành chính.
{guide}

YÊU CẦU ĐỊNH DẠNG NGHIÊM NGẶT (TRẢ VỀ DUY NHẤT HTML):
Bạn PHẢI sử dụng các thẻ HTML với thuộc tính style để căn lề. KHÔNG ĐƯỢC dùng Markdown. KHÔNG bao bọc kết quả trong ```html.
Bạn PHẢI dùng cấu trúc HTML sau cho phần đầu văn bản:
<p style="text-align: left"><b>[TÊN CƠ QUAN CHỦ QUẢN]</b><br/><b>{request.issuing_unit}</b><br/>Số: .../[Ký hiệu]</p>
<p style="text-align: center"><b>CỘNG HÒA XÃ HỘI CHỦ NGHĨA VIỆT NAM</b><br/><b>Độc lập - Tự do - Hạnh phúc</b></p>
<p style="text-align: right"><i>..., ngày ... tháng ... năm ...</i></p>

Phần tên loại văn bản phải là thẻ <h2> căn giữa: <h2 style="text-align: center"><b>[TÊN LOẠI VĂN BẢN]</b></h2>
Nội dung chính dùng thẻ <p style="text-align: justify">.
Phần cuối văn bản phải có Nơi nhận và Chữ ký người có thẩm quyền.
Không được tự bịa căn cứ pháp lý. Chỉ dùng căn cứ từ ngữ cảnh tài liệu pháp lý hoặc từ yêu cầu người dùng; nếu thiếu căn cứ, phải ghi cảnh báo nguồn.
{retrieval_prompt}
{clarification_prompt}

Ngữ cảnh tài liệu pháp lý để tham khảo và áp dụng:
{context_str}

Cảnh báo: Nếu ngữ cảnh tài liệu không đủ thông tin pháp lý để soạn thảo, hãy chèn một đoạn cảnh báo nhỏ ở cuối văn bản bằng thẻ <p><i>Cảnh báo: Không tìm thấy căn cứ pháp lý rõ ràng.</i></p>.
"""

    human_prompt = f"""Đơn vị ban hành: {request.issuing_unit}
Nơi nhận: {request.recipient}
Yêu cầu nội dung: {request.draft_requirement}
Hãy viết toàn bộ văn bản hoàn chỉnh dựa trên yêu cầu trên.
"""
    return SystemMessage(content=system_prompt), HumanMessage(content=human_prompt)


def _build_citations(sources: list[dict[str, Any]]) -> list["RetrieveCitation"]:
    return [
        RetrieveCitation(
            title=s.get("metadata", {}).get("title", ""),
            excerpt=s.get("content", "")[:200] + "...",
            citation_label=f"[{i+1}]",
            url=s.get("metadata", {}).get("url", ""),
            source_uri=s.get("metadata", {}).get("source_uri", ""),
        )
        for i, s in enumerate(sources)
    ]


def _strip_html_fence(text: str) -> str:
    if text.startswith("```html"):
        text = text[7:]
    elif text.startswith("```"):
        text = text[3:]
    if text.endswith("```"):
        text = text[:-3]
    return text.strip()


async def _retrieve_planned_sources_async(
    *,
    request: "DraftRequest",
    coordinator: "RetrievalAccessCoordinator",
    workspace: "WorkspaceRequest",
    current_user: Any,
    guest_session_id: str | None,
    planning: "DraftPlanningResult",
) -> tuple[list[dict[str, Any]], str]:
    """Run simplified retrieval queries in parallel; RRF-fuse and cap at top_k."""
    queries = _planned_retrieval_queries(planning, request)

    candidate_k = _candidate_top_k(request.top_k)
    access_context = _draft_source_access_context(
        coordinator=coordinator,
        workspace=workspace,
        include_private_library=request.include_private_library,
        current_user=current_user,
        guest_session_id=guest_session_id,
    )

    async def _run_one(query: str) -> tuple[list[dict[str, Any]], str]:
        return await asyncio.wait_for(
            asyncio.to_thread(
                coordinator.retrieve,
                query=query,
                retrieval_mode=request.retrieval_mode,
                workspace=workspace,
                include_public_vbpl=request.include_public_vbpl,
                include_private_library=request.include_private_library,
                top_k=candidate_k,
                user=current_user,
                guest_session_id=guest_session_id,
                access_context=access_context,
            ),
            timeout=DRAFT_RETRIEVAL_QUERY_TIMEOUT_S,
        )

    # return_exceptions=True so one failed/timed-out query never kills the
    # whole draft — we collect what succeeds and surface the rest as a warning.
    results = await asyncio.gather(*[_run_one(q) for q in queries], return_exceptions=True)

    warnings: list[str] = []
    failed_queries = 0
    ranked_lists: list[list[Any]] = []
    for result in results:
        if isinstance(result, BaseException):
            failed_queries += 1
            continue
        retrieved, warning = result
        if warning:
            warnings.append(warning)
        ranked_lists.append([_dict_to_source(source) for source in retrieved])

    sources = _fuse_ranked_sources(
        ranked_lists,
        top_k=request.top_k,
        max_chunks_per_document=DRAFT_MAX_CHUNKS_PER_DOCUMENT,
    )

    if failed_queries:
        if sources:
            warnings.append(
                f"Một số truy vấn nguồn ({failed_queries}/{len(queries)}) bị lỗi hoặc quá thời gian; "
                "soạn thảo dựa trên các nguồn còn lại."
            )
        else:
            warnings.append(
                "Không truy xuất được nguồn pháp lý (lỗi hoặc quá thời gian); "
                "soạn thảo theo thể thức và yêu cầu người dùng, kiểm tra kỹ căn cứ pháp lý."
            )

    return sources, " ".join(_dedupe_text(warnings))


async def generate_draft_stream(
    request: "DraftRequest",
    coordinator: "RetrievalAccessCoordinator",
    current_user: Any,
    guest_session_id: str | None,
) -> AsyncIterator[dict[str, Any]]:
    """Async generator yielding SSE-ready frames: meta, token*, done (or error)."""
    workspace = request.workspace or WorkspaceRequest(type="public_only")
    if request.planning is not None:
        planning = request.planning
    else:
        try:
            planning = await asyncio.wait_for(
                asyncio.to_thread(plan_draft, request),
                timeout=DRAFT_PLANNING_TIMEOUT_S,
            )
        except TimeoutError:
            # Planner unavailable/slow — fall back to a ready plan so the user
            # still gets a draft instead of an indefinite spinner.
            planning = DraftPlanningResult(
                status="ready",
                reasoning_summary="Bỏ qua bước lập kế hoạch do quá thời gian; soạn thảo trực tiếp từ yêu cầu.",
                needs_retrieval=True,
                questions=[],
                question_options=[],
                breakdown=[],
                retrieval_queries=[request.draft_requirement],
                suggested_prompt=request.draft_requirement,
            )

    if planning.status != "ready" and not request.force_generate_with_highlights:
        yield {
            "event": "meta",
            "data": {
                "html": "",
                "citations": [],
                "warnings": ["Cần làm rõ yêu cầu trước khi truy xuất VBPL và soạn thảo."],
                "template_source": OFFICIAL_TEMPLATE_SOURCE.model_dump(),
                "planning": planning.model_dump(),
            },
        }
        yield {"event": "done", "data": {}}
        return

    # Planning finished — tell the client right away so it can advance the
    # progress UI instead of sitting on "Đang phân tích yêu cầu" until the
    # first token (which only comes after retrieval).
    yield {
        "event": "status",
        "data": {
            "phase": "retrieving" if planning.needs_retrieval else "generating",
            "needs_retrieval": planning.needs_retrieval,
            "planning": planning.model_dump(),
        },
    }

    if planning.needs_retrieval:
        try:
            sources, warning = await _retrieve_planned_sources_async(
                request=request,
                coordinator=coordinator,
                workspace=workspace,
                current_user=current_user,
                guest_session_id=guest_session_id,
                planning=planning,
            )
        except Exception as exc:
            # Retrieval is best-effort: never let it abort generation. The draft
            # can still be written from the template + user requirement.
            sources = []
            warning = (
                "Không truy xuất được nguồn pháp lý (lỗi hệ thống truy xuất); "
                f"soạn thảo theo thể thức và yêu cầu người dùng. Chi tiết: {exc}"
            )
    else:
        sources, warning = [], ""

    citations = _build_citations(sources)
    warnings_list = [warning] if warning else []

    yield {
        "event": "meta",
        "data": {
            "citations": [c.model_dump() for c in citations],
            "warnings": warnings_list,
            "template_source": OFFICIAL_TEMPLATE_SOURCE.model_dump(),
            "planning": planning.model_dump(),
        },
    }

    system_msg, human_msg = _build_draft_messages(request, planning, sources)
    llm = get_llm(streaming=True)

    accumulated = ""
    fence_stripped = False
    stream_iter = llm.astream([system_msg, human_msg]).__aiter__()
    while True:
        try:
            chunk = await asyncio.wait_for(
                stream_iter.__anext__(), timeout=DRAFT_STREAM_IDLE_TIMEOUT_S
            )
        except StopAsyncIteration:
            break
        except TimeoutError:
            if accumulated:
                # Got partial content before stalling — finalize what we have.
                break
            yield {
                "event": "error",
                "data": {"message": "Mô hình không phản hồi kịp thời, vui lòng thử lại."},
            }
            return
        raw = chunk.content if hasattr(chunk, "content") else str(chunk)
        if not raw:
            continue
        prev_len = len(accumulated)
        accumulated += raw
        if not fence_stripped:
            if accumulated.startswith("```html") and len(accumulated) > 7:
                accumulated = accumulated[7:].lstrip("\n")
                fence_stripped = True
                prev_len = 0
            elif accumulated.startswith("```") and len(accumulated) > 3:
                accumulated = accumulated[3:].lstrip("\n")
                fence_stripped = True
                prev_len = 0
        clean_delta = accumulated[prev_len:]
        if clean_delta:
            yield {"event": "token", "data": {"delta": clean_delta}}

    final_html = _strip_html_fence(accumulated)
    yield {"event": "done", "data": {"html": final_html}}


def _planned_retrieval_queries(
    planning: "DraftPlanningResult", request: "DraftRequest"
) -> list[str]:
    """Simplify the planner's retrieval queries the same way the chat RAG flow does.

    Strips instruction boilerplate (and preserves legal numbers/diacritics) so the query
    handed to retrieval carries the discriminative signal. Capped at 5 sub-queries.
    """
    raw = planning.retrieval_queries or [planning.suggested_prompt or request.draft_requirement]
    return _dedupe_text([simplify_legal_query(query) for query in raw[:5]])


def _candidate_top_k(top_k: int) -> int:
    """Over-fetch per sub-query so that after collapsing to one chunk per document we
    still have enough DISTINCT documents to fill top_k."""
    return min(20, max(top_k, top_k * _DRAFT_CANDIDATE_MULTIPLIER))


def _draft_source_access_context(
    *,
    coordinator: RetrievalAccessCoordinator,
    workspace: WorkspaceRequest,
    include_private_library: bool,
    current_user: Any,
    guest_session_id: str | None,
):
    if not include_private_library or workspace.type == "public_only":
        return None
    return source_access_context(coordinator.db, current_user, guest_session_id)


def _fuse_ranked_sources(
    ranked_lists: list[list[Any]],
    *,
    top_k: int,
    max_chunks_per_document: int = DEFAULT_MAX_CHUNKS_PER_DOCUMENT,
) -> list[dict[str, Any]]:
    """Fuse per-sub-query ranked source lists with Reciprocal Rank Fusion.

    Plain concat + truncate-at-top_k lets the first sub-query fill every slot and starves
    a document that ranks #1 in a later sub-query. RRF keeps each sub-query represented;
    the document cap then keeps the final căn cứ list to distinct documents.
    """
    fused = reciprocal_rank_fusion_lists(ranked_lists)
    merged = _merge_sources(
        fused,
        [],
        limit=top_k,
        max_chunks_per_document=max_chunks_per_document,
    )
    return [source.to_dict() for source in merged]


def _retrieve_planned_sources(
    *,
    request: DraftRequest,
    coordinator: RetrievalAccessCoordinator,
    workspace: WorkspaceRequest,
    current_user: Any,
    guest_session_id: str | None,
    planning: DraftPlanningResult,
) -> tuple[list[dict[str, Any]], str]:
    queries = _planned_retrieval_queries(planning, request)
    if not queries:
        return [], ""

    candidate_k = _candidate_top_k(request.top_k)
    access_context = _draft_source_access_context(
        coordinator=coordinator,
        workspace=workspace,
        include_private_library=request.include_private_library,
        current_user=current_user,
        guest_session_id=guest_session_id,
    )

    def _run_one(query: str) -> tuple[list[dict[str, Any]], str]:
        return coordinator.retrieve(
            query=query,
            retrieval_mode=request.retrieval_mode,
            workspace=workspace,
            include_public_vbpl=request.include_public_vbpl,
            include_private_library=request.include_private_library,
            top_k=candidate_k,
            user=current_user,
            guest_session_id=guest_session_id,
            access_context=access_context,
        )

    # Fan the sub-queries out concurrently — each does its own embedding + Chroma + BM25
    # round-trip, so sequential execution made latency the sum of all sub-queries.
    # ThreadPoolExecutor.map preserves input order in the results.
    with ThreadPoolExecutor(
        max_workers=min(len(queries), _DRAFT_RETRIEVAL_MAX_WORKERS)
    ) as executor:
        results = list(executor.map(_run_one, queries))

    warnings: list[str] = []
    ranked_lists: list[list[Any]] = []
    for retrieved, warning in results:
        if warning:
            warnings.append(warning)
        ranked_lists.append([_dict_to_source(source) for source in retrieved])

    sources = _fuse_ranked_sources(
        ranked_lists,
        top_k=request.top_k,
        max_chunks_per_document=DRAFT_MAX_CHUNKS_PER_DOCUMENT,
    )
    return sources, " ".join(_dedupe_text(warnings))


def _clarification_highlight_prompt(planning: DraftPlanningResult, force_highlights: bool) -> str:
    if not force_highlights or planning.status == "ready":
        return ""

    missing_items = planning.questions or planning.breakdown or ["thông tin cần người dùng làm rõ"]
    missing_html = "\n".join(f"- {item}" for item in missing_items)
    return f"""
YÊU CẦU KHI CÒN THIẾU THÔNG TIN:
Người dùng đã yêu cầu cứ soạn bản nháp dù còn điểm cần làm rõ. Bạn PHẢI soạn tiếp như bình thường, nhưng KHÔNG tự bịa các phần thiếu.
Mỗi vùng thiếu thông tin phải được tô vàng bằng đúng dạng HTML:
<span data-clarification="true" style="background-color: #FEF3C7; padding: 0 2px;">[Cần làm rõ: mô tả ngắn phần thiếu]</span>
Các điểm hiện cần làm rõ:
{missing_html}
"""


def _retrieval_usage_prompt(planning: DraftPlanningResult) -> str:
    if planning.needs_retrieval:
        return ""
    return """
PLANNER ĐÃ XÁC ĐỊNH KHÔNG CẦN RETRIEVE VBPL:
Chỉ soạn theo thể thức/mẫu hành chính và thông tin người dùng cung cấp.
Không viện dẫn điều, khoản, số hiệu văn bản pháp luật, thời hạn pháp lý hoặc căn cứ pháp lý cụ thể nếu không có trong yêu cầu người dùng.
"""


def _dedupe_text(items: list[str]) -> list[str]:
    seen = set()
    deduped = []
    for item in items:
        text = str(item or "").strip()
        if not text or text in seen:
            continue
        seen.add(text)
        deduped.append(text)
    return deduped


_FORMAT_RULES: dict[str, list[tuple[str, str, str]]] = {
    # (regex_pattern, error_severity, description_vi)
    "common": [
        (
            r"CỘNG\s+HÒA\s+XÃ\s+HỘI\s+CHỦ\s+NGHĨA\s+VIỆT\s+NAM",
            "error",
            "Thiếu dòng quốc hiệu 'CỘNG HÒA XÃ HỘI CHỦ NGHĨA VIỆT NAM'",
        ),
        (
            r"Độc\s+lập\s*[-–]\s*Tự\s+do\s*[-–]\s*Hạnh\s+phúc",
            "error",
            "Thiếu dòng tiêu ngữ 'Độc lập - Tự do - Hạnh phúc'",
        ),
        (
            r"text-align:\s*right",
            "error",
            "Thiếu dòng địa danh và thời gian ban hành (căn phải)",
        ),
        (
            r"Nơi\s+nhận",
            "error",
            "Thiếu mục 'Nơi nhận' ở phần cuối văn bản",
        ),
        (
            r"<h[1-6][^>]*>",
            "error",
            "Thiếu tiêu đề tên loại văn bản (thẻ h2 căn giữa)",
        ),
    ],
    "cong_van": [
        (r"[Kk]ính\s+gửi", "error", "Thiếu dòng 'Kính gửi:' xác định nơi nhận chính"),
        (r"V/v\b|[Vv]ề\s+việc", "warning", "Nên có trích yếu nội dung 'V/v ...' hoặc 'Về việc ...'"),
    ],
    "to_trinh": [
        (r"[Kk]ính\s+gửi", "error", "Thiếu dòng 'Kính gửi:' trình lên cơ quan có thẩm quyền"),
        (r"[Ss]ự\s+cần\s+thiết|[Cc]ăn\s+cứ", "warning", "Nên có phần trình bày sự cần thiết hoặc căn cứ"),
        (r"[Kk]ính?\s+trình|[Đđ]ề\s+nghị", "warning", "Nên có câu kính trình hoặc đề nghị phê duyệt"),
    ],
    "thong_bao": [
        (r"[Tt]hông\s+báo", "error", "Tiêu đề phải là 'THÔNG BÁO'"),
    ],
    "quyet_dinh_noi_bo": [
        (r"[Cc]ăn\s+cứ", "error", "Thiếu phần căn cứ ban hành quyết định"),
        (r"QUYẾT\s+ĐỊNH\s*:", "error", "Thiếu cụm 'QUYẾT ĐỊNH:' phân tách phần điều khoản"),
        (r"[Đđ]iều\s+1\b", "error", "Thiếu Điều 1 — nội dung quyết định"),
        (r"[Đđ]iều\s+2\b", "warning", "Nên có Điều 2 — tổ chức thực hiện"),
    ],
}

_SPACING_PATTERNS = [
    (
        r"(<p[^>]*>\s*(?:<br\s*/?>)?\s*</p>\s*){3,}",
        "warning",
        "Có đoạn trắng liên tiếp (≥3 thẻ <p> rỗng) — kiểm tra khoảng cách giữa các phần",
    ),
    (
        r"<p[^>]*style=\"[^\"]*text-align:\s*justify[^\"]*\"[^>]*>\s*(?:<br\s*/?>)?\s*</p>",
        "warning",
        "Có thẻ <p justify> rỗng — có thể thừa khoảng trắng trong phần nội dung",
    ),
]


def _check_draft_format(html: str, template_key: str) -> list[DraftCheckIssue]:
    issues: list[DraftCheckIssue] = []
    idx = 0

    rules = _FORMAT_RULES.get("common", []) + _FORMAT_RULES.get(template_key, [])
    for pattern, severity, description in rules:
        if not re.search(pattern, html, re.IGNORECASE | re.DOTALL):
            issues.append(DraftCheckIssue(
                id=f"fmt-{idx}",
                type="format",
                severity=severity,  # type: ignore[arg-type]
                description=description,
            ))
            idx += 1

    for pattern, severity, description in _SPACING_PATTERNS:
        if re.search(pattern, html, re.IGNORECASE | re.DOTALL):
            issues.append(DraftCheckIssue(
                id=f"spacing-{idx}",
                type="format",
                severity=severity,  # type: ignore[arg-type]
                description=description,
            ))
            idx += 1

    return issues


def _draft_check_cache_key(
    html: str, requirement: str, citations: list[dict], template_key: str
) -> str:
    payload = json.dumps(
        {
            "html": html,
            "requirement": requirement,
            "citations": citations,
            "template_key": template_key,
        },
        sort_keys=True,
        ensure_ascii=False,
        default=str,
    )
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def check_draft_grounding(
    html: str,
    requirement: str,
    citations: list[dict],
    template_key: str = "cong_van",
) -> tuple[str, int, list[DraftCheckIssue]]:
    cache_key = _draft_check_cache_key(html, requirement, citations, template_key)
    with _draft_check_cache_lock:
        cached = _draft_check_cache.get(cache_key)
        if cached is not None:
            _draft_check_cache.move_to_end(cache_key)
            return cached

    format_issues = _check_draft_format(html, template_key)

    source_block = "\n".join(
        f"[{i+1}] {c.get('title', '')}: {c.get('excerpt', '')}"
        for i, c in enumerate(citations)
    ) or "(Không có nguồn truy xuất)"

    # Grounding check is a constrained annotation task — the light model handles it and
    # roughly halves the call latency.
    llm = get_llm(light=True)
    messages = [
        SystemMessage(content=(
            "Bạn là chuyên gia kiểm tra độ tin cậy văn bản pháp luật Việt Nam.\n"
            "Nhiệm vụ: rà soát bản thảo HTML, tìm những đoạn văn có thể không được hỗ trợ "
            "bởi yêu cầu của người dùng hoặc các nguồn đã truy xuất.\n"
            "Những đoạn cần đánh dấu bao gồm:\n"
            "- Số hiệu điều khoản, văn bản pháp lý cụ thể không xuất hiện trong nguồn.\n"
            "- Thời hạn, ngày tháng pháp lý không có căn cứ.\n"
            "- Tên cơ quan, thẩm quyền ký không được người dùng xác nhận.\n"
            "- Kết luận pháp lý không thể truy xuất từ nguồn hoặc yêu cầu.\n"
            "Bọc MỖI đoạn không chắc chắn trong:\n"
            '<span data-clarification="true" style="background-color: #FEF3C7; padding: 0 2px;">[đoạn gốc]</span>\n'
            "KHÔNG đánh dấu những phần rõ ràng được nêu trong yêu cầu hoặc có trong nguồn.\n"
            "Trả về DUY NHẤT HTML đã chú thích, không giải thích, không bọc trong ```html."
        )),
        HumanMessage(content=(
            f"Yêu cầu người dùng:\n{requirement}\n\n"
            f"Nguồn đã truy xuất:\n{source_block}\n\n"
            f"Bản thảo HTML:\n{html}"
        )),
    ]
    result = str(llm.invoke(messages).content).strip()
    for prefix in ("```html", "```"):
        if result.startswith(prefix):
            result = result[len(prefix):]
    if result.endswith("```"):
        result = result[:-3]
    result = result.strip()

    uncertain_count = len(re.findall(r'data-clarification="true"', result))

    grounding_issues: list[DraftCheckIssue] = []
    if uncertain_count > 0:
        grounding_issues.append(DraftCheckIssue(
            id="grounding-0",
            type="grounding",
            severity="warning",
            description=f"Có {uncertain_count} đoạn được tô vàng vì chưa có căn cứ pháp lý rõ ràng — kiểm tra và bổ sung hoặc xóa bỏ.",
        ))

    result_tuple = (result, uncertain_count, format_issues + grounding_issues)
    with _draft_check_cache_lock:
        _draft_check_cache[cache_key] = result_tuple
        _draft_check_cache.move_to_end(cache_key)
        while len(_draft_check_cache) > _DRAFT_CHECK_CACHE_MAX:
            _draft_check_cache.popitem(last=False)
    return result_tuple


def patch_draft_selection(selected_html: str, instruction: str) -> str:
    llm = get_llm()
    messages = [
        SystemMessage(content=(
            "Bạn là chuyên viên soạn thảo văn bản pháp luật/hành chính Việt Nam.\n"
            "Người dùng cung cấp một đoạn HTML từ bản nháp và yêu cầu sửa đổi cụ thể.\n"
            "Trả về DUY NHẤT đoạn HTML đã được sửa. Không thêm giải thích, không bọc trong ```html.\n"
            "Chỉ chỉnh sửa đúng theo yêu cầu; không thay đổi phần không liên quan."
        )),
        HumanMessage(content=(
            f"Đoạn gốc (HTML):\n{selected_html}\n\n"
            f"Yêu cầu sửa đổi:\n{instruction}\n\n"
            "Trả về đoạn HTML đã sửa:"
        )),
    ]
    result = str(llm.invoke(messages).content).strip()
    for prefix in ("```html", "```"):
        if result.startswith(prefix):
            result = result[len(prefix):]
    if result.endswith("```"):
        result = result[:-3]
    return result.strip()


def _apply_font(run) -> None:
    """Force Times New Roman on a run, including the East Asian slot so Vietnamese
    diacritics render in the same face."""
    run.font.name = DOCX_FONT_NAME
    rpr = run._element.get_or_add_rPr()
    rfonts = rpr.find(qn("w:rFonts"))
    if rfonts is None:
        rfonts = rpr.makeelement(qn("w:rFonts"), {})
        rpr.append(rfonts)
    for attr in ("w:ascii", "w:hAnsi", "w:cs", "w:eastAsia"):
        rfonts.set(qn(attr), DOCX_FONT_NAME)


def export_docx(request: DraftExportRequest) -> io.BytesIO:
    doc = Document()

    # Set Times New Roman as the document-wide default (Normal style); headings
    # and individual runs are also normalized explicitly below.
    normal_font = doc.styles["Normal"].font
    normal_font.name = DOCX_FONT_NAME
    normal_rpr = doc.styles["Normal"].element.get_or_add_rPr()
    normal_rfonts = normal_rpr.find(qn("w:rFonts"))
    if normal_rfonts is None:
        normal_rfonts = normal_rpr.makeelement(qn("w:rFonts"), {})
        normal_rpr.append(normal_rfonts)
    for attr in ("w:ascii", "w:hAnsi", "w:cs", "w:eastAsia"):
        normal_rfonts.set(qn(attr), DOCX_FONT_NAME)

    doc.add_heading(request.title, 0)

    # Basic regex-based HTML parsing for DOCX
    # This is a rudimentary parser sufficient for the MVP's limited HTML tag support.

    # Normalize html slightly
    html = request.html.replace("\n", " ")

    # Extract blocks (p, h1-h6, ul, ol)
    # We will split the HTML by block elements
    block_pattern = re.compile(r'<(p|h[1-6]|ul|ol)[^>]*>(.*?)</\1>', re.IGNORECASE)

    for match in block_pattern.finditer(html):
        tag = match.group(1).lower()
        content = match.group(2)

        if tag.startswith('h'):
            level = int(tag[1])
            paragraph = doc.add_heading('', level=level)
        elif tag == 'p':
            paragraph = doc.add_paragraph()
            # alignment checking could be done here if classes were preserved,
            # but we'll stick to basic runs for now.
            if 'text-align="center"' in match.group(0) or 'text-align: center' in match.group(0):
                 paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
            elif 'text-align="right"' in match.group(0) or 'text-align: right' in match.group(0):
                 paragraph.alignment = WD_ALIGN_PARAGRAPH.RIGHT
            elif 'text-align="justify"' in match.group(0) or 'text-align: justify' in match.group(0):
                 paragraph.alignment = WD_ALIGN_PARAGRAPH.JUSTIFY
        elif tag in ('ul', 'ol'):
            # parse list items
            li_pattern = re.compile(r'<li[^>]*>(.*?)</li>', re.IGNORECASE)
            for li_match in li_pattern.finditer(content):
                li_content = li_match.group(1)
                style = 'List Bullet' if tag == 'ul' else 'List Number'
                paragraph = doc.add_paragraph(style=style)
                _parse_runs(paragraph, li_content)
            continue

        _parse_runs(paragraph, content)

    if request.warnings:
        doc.add_heading("Cảnh báo", level=1)
        for w in request.warnings:
            doc.add_paragraph(w, style='List Bullet')

    if request.citations:
        doc.add_heading("Tài liệu tham khảo", level=1)
        for c in request.citations:
            p = doc.add_paragraph(style='List Bullet')
            p.add_run(f"{c.citation_label} {c.title}").bold = True
            if c.excerpt:
                p.add_run(f" - {c.excerpt}")

    # Normalize every run to Times New Roman (covers title, headings, body,
    # warnings and citations regardless of their originating style).
    for paragraph in doc.paragraphs:
        for run in paragraph.runs:
            _apply_font(run)

    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer


def _parse_runs(paragraph, html_content):
    # Splits inline tags: <b>, <i>, <u>, <strong>, <em>
    # This regex attempts to find tags and plain text.
    # We will use a simple scanner.

    # remove unsupported inline tags except b, i, u, strong, em, br, span
    clean_html = re.sub(r'</?(div|a|img)[^>]*>', '', html_content)

    tokens = re.split(r'(</?(?:b|i|u|strong|em|br|span)[^>]*>)', clean_html, flags=re.IGNORECASE)

    bold = False
    italic = False
    underline = False
    font_size_stack: list[float | None] = []

    for token in tokens:
        token_lower = token.lower()
        if token_lower in ('<b>', '<strong>'):
            bold = True
        elif token_lower in ('</b>', '</strong>'):
            bold = False
        elif token_lower in ('<i>', '<em>'):
            italic = True
        elif token_lower in ('</i>', '</em>'):
            italic = False
        elif token_lower == '<u>':
            underline = True
        elif token_lower == '</u>':
            underline = False
        elif token_lower.startswith('<span'):
            font_size_stack.append(_extract_font_size_pt(token))
        elif token_lower == '</span>':
            if font_size_stack:
                font_size_stack.pop()
        elif token_lower in ('<br>', '<br/>', '<br />'):
            paragraph.add_run("\n")
        elif not token.startswith('<'):
            # Text node
            if token:
                # unescape basic HTML entities
                text = token.replace('&nbsp;', ' ').replace('&lt;', '<').replace('&gt;', '>').replace('&amp;', '&')
                run = paragraph.add_run(text)
                run.bold = bold
                run.italic = italic
                run.underline = underline
                active_font_size = next((size for size in reversed(font_size_stack) if size), None)
                if active_font_size:
                    run.font.size = Pt(active_font_size)


def _extract_font_size_pt(tag: str) -> float | None:
    match = re.search(r'font-size\s*:\s*([0-9]+(?:\.[0-9]+)?)(pt|px)?', tag, re.IGNORECASE)
    if not match:
        return None

    value = float(match.group(1))
    unit = (match.group(2) or "pt").lower()
    if unit == "px":
        return value * 0.75
    return value