backend/chatbot/structured_legal.py

"""
Structured legal answer helpers using LangChain output parsers.
"""

from __future__ import annotations

import json
import logging
import textwrap
from functools import lru_cache
from typing import List, Optional, Sequence

from langchain.output_parsers import PydanticOutputParser
from langchain.schema import OutputParserException
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


class LegalCitation(BaseModel):
    """Single citation item pointing back to a legal document."""

    document_title: str = Field(..., description="Tên văn bản pháp luật.")
    section_code: str = Field(..., description="Mã điều/khoản được trích dẫn.")
    page_range: Optional[str] = Field(
        None, description="Trang hoặc khoảng trang trong tài liệu."
    )
    summary: str = Field(
        ...,
        description="1-2 câu mô tả nội dung chính của trích dẫn, phải liên quan trực tiếp câu hỏi.",
    )
    snippet: str = Field(
        ..., description="Trích đoạn ngắn gọn (≤500 ký tự) lấy từ tài liệu gốc."
    )


class LegalAnswer(BaseModel):
    """Structured answer returned by the LLM."""

    summary: str = Field(
        ...,
        description="Đoạn mở đầu tóm tắt kết luận chính, phải nhắc văn bản áp dụng (ví dụ Quyết định 69/QĐ-TW).",
    )
    details: List[str] = Field(
        ...,
        description="Tối thiểu 2 gạch đầu dòng mô tả từng hình thức/điều khoản. Mỗi gạch đầu dòng phải nhắc mã điều hoặc tên văn bản.",
    )
    citations: List[LegalCitation] = Field(
        ...,
        description="Danh sách trích dẫn; phải có ít nhất 1 phần tử tương ứng với các tài liệu đã cung cấp.",
    )


@lru_cache(maxsize=1)
def get_legal_output_parser() -> PydanticOutputParser:
    """Return cached parser to enforce structured output."""

    return PydanticOutputParser(pydantic_object=LegalAnswer)


def build_structured_legal_prompt(
    query: str,
    documents: Sequence,
    parser: PydanticOutputParser,
    prefill_summary: Optional[str] = None,
    retry_hint: Optional[str] = None,
) -> str:
    """Construct prompt instructing the LLM to return structured JSON."""

    doc_blocks = []
    for idx, doc in enumerate(documents[:5], 1):
        document = getattr(doc, "document", None)
        title = getattr(document, "title", "") or "Không rõ tên văn bản"
        code = getattr(document, "code", "") or "N/A"
        section_code = getattr(doc, "section_code", "") or "Không rõ điều"
        section_title = getattr(doc, "section_title", "") or ""
        page_range = _format_page_range(doc)
        content = getattr(doc, "content", "") or ""
        snippet = (content[:800] + "...") if len(content) > 800 else content

        block = textwrap.dedent(
            f"""
            TÀI LIỆU #{idx}
            Văn bản: {title} (Mã: {code})
            Điều/khoản: {section_code} - {section_title}
            Trang: {page_range or 'Không rõ'}
            Trích đoạn:
            {snippet}
            """
        ).strip()
        doc_blocks.append(block)

    docs_text = "\n\n".join(doc_blocks)
    reference_lines = []
    title_section_pairs = []
    for doc in documents[:5]:
        document = getattr(doc, "document", None)
        title = getattr(document, "title", "") or "Không rõ tên văn bản"
        section_code = getattr(doc, "section_code", "") or "Không rõ điều"
        reference_lines.append(f"- {title} | {section_code}")
        title_section_pairs.append((title, section_code))
    reference_text = "\n".join(reference_lines)
    prefill_block = ""
    if prefill_summary:
        prefill_block = textwrap.dedent(
            f"""
            Bản tóm tắt tiếng Việt đã có sẵn (hãy dùng lại, diễn đạt ngắn gọn hơn, KHÔNG thêm thông tin mới):
            {prefill_summary.strip()}
            """
        ).strip()
    format_instructions = parser.get_format_instructions()
    retry_hint_block = ""
    if retry_hint:
        retry_hint_block = textwrap.dedent(
            f"""
            Nhắc lại: {retry_hint.strip()}
            """
        ).strip()

    prompt = textwrap.dedent(
        f"""
        Bạn là trợ lý pháp lý của Công an thành phố Huế. Nhiệm vụ: dựa trên các trích đoạn dưới đây để trả lời câu hỏi của người dân.

        Quy tắc bắt buộc:
        - Không được bịa đặt thông tin ngoài tài liệu.
        - Phải nhắc rõ văn bản (ví dụ: Quyết định 69/QĐ-TW) và mã điều/khoản trong phần trả lời.
        - Cấu trúc trả lời: SUMMARY ngắn gọn -> DETAILS dạng bullet -> CITATIONS chứa thông tin nguồn.
        - Nếu không đủ thông tin, ghi rõ lý do ở phần summary và để danh sách citations rỗng.
        - Tuyệt đối không chép lại schema hay thêm khóa "$defs"; chỉ xuất đối tượng JSON cuối cùng theo mẫu dưới đây.
        - Chỉ in ra CHÍNH XÁC một JSON object, không được thêm chữ 'json', không dùng ``` hoặc văn bản thừa trước/sau.
        - Mỗi bullet DETAILS bắt buộc phải chứa tên văn bản và mã điều/khoản đúng như trong “Bảng tham chiếu” phía dưới.
        - Không được tạo thêm hình thức kỷ luật hoặc điều khoản không xuất hiện trong tài liệu. Nếu không thấy điều/khoản, ghi rõ “(không nêu điều cụ thể)”.
        - Ví dụ định dạng:
          {{
            "summary": "Tóm tắt ...",
            "details": ["- Điều 5 ...", "- Điều 7 ..."],
            "citations": [
              {{
                "document_title": "Quyết định 69/QĐ-TW",
                "section_code": "Điều 5",
                "page_range": "1-2",
                "summary": "Mô tả ngắn gọn",
                "snippet": "Trích dẫn ≤500 ký tự"
              }}
            ]
          }}

        Câu hỏi người dùng: {query}

        Bảng tham chiếu bắt buộc (chỉ sử dụng đúng tên/mã dưới đây):
        {reference_text}

        Các trích đoạn pháp luật:
        {docs_text}

        {prefill_block}

        {retry_hint_block}

        {format_instructions}
        """
    ).strip()

    return prompt


def format_structured_legal_answer(answer: LegalAnswer) -> str:
    """Convert structured answer into human-friendly text with citations."""

    lines: List[str] = []
    if answer.summary:
        lines.append(answer.summary.strip())

    if answer.details:
        lines.append("")
        lines.append("Chi tiết chính:")
        for bullet in answer.details:
            lines.append(f"- {bullet.strip()}")

    if answer.citations:
        lines.append("")
        lines.append("Trích dẫn chi tiết:")
        for idx, citation in enumerate(answer.citations, 1):
            page_text = f" (Trang: {citation.page_range})" if citation.page_range else ""
            lines.append(
                f"{idx}. {citation.document_title} – {citation.section_code}{page_text}"
            )
            lines.append(f"   Tóm tắt: {citation.summary.strip()}")
            lines.append(f"   Trích đoạn: {citation.snippet.strip()}")

    return "\n".join(lines).strip()


def _format_page_range(doc: object) -> Optional[str]:
    start = getattr(doc, "page_start", None)
    end = getattr(doc, "page_end", None)
    if start and end:
        if start == end:
            return str(start)
        return f"{start}-{end}"
    if start:
        return str(start)
    if end:
        return str(end)
    return None


def parse_structured_output(
    parser: PydanticOutputParser, raw_output: str
) -> Optional[LegalAnswer]:
    """Parse raw LLM output to LegalAnswer if possible."""

    if not raw_output:
        return None
    try:
        return parser.parse(raw_output)
    except OutputParserException:
        snippet = raw_output.strip().replace("\n", " ")
        logger.warning(
            "[LLM] Structured parse failed. Preview: %s",
            snippet[:400],
        )
        json_candidate = _extract_json_block(raw_output)
        if json_candidate:
            try:
                return parser.parse(json_candidate)
            except OutputParserException:
                logger.warning("[LLM] JSON reparse also failed.")
                return None
        return None


def _extract_json_block(text: str) -> Optional[str]:
    """
    Best-effort extraction of the first JSON object within text.
    """
    stripped = text.strip()
    if stripped.startswith("```"):
        stripped = stripped.lstrip("`")
        if stripped.lower().startswith("json"):
            stripped = stripped[4:]
        stripped = stripped.strip("`").strip()

    start = text.find("{")
    if start == -1:
        return None

    stack = 0
    for idx in range(start, len(text)):
        char = text[idx]
        if char == "{":
            stack += 1
        elif char == "}":
            stack -= 1
            if stack == 0:
                payload = text[start : idx + 1]
                # Remove code fences if present
                payload = payload.strip()
                if payload.startswith("```"):
                    payload = payload.strip("`").strip()
                try:
                    json.loads(payload)
                    return payload
                except json.JSONDecodeError:
                    return None
    return None