src/services/citations.py

import hashlib

from src.models.schemas import Citation, EffectiveStatus, LegalChunk, LegalDocument

MISSING_METADATA = "metadata-missing"


class CitationServiceError(Exception):
    code = "citation_error"

    def __init__(self, message: str, details: dict[str, object] | None = None) -> None:
        super().__init__(message)
        self.message = message
        self.details = details or {}


class CitationReferenceError(CitationServiceError):
    code = "citation_reference_error"


class CitationService:
    def create_citation(
        self,
        *,
        document: LegalDocument,
        chunk: LegalChunk,
        allowed_chunk_ids: set[str] | None = None,
        text_span: str | None = None,
        retrieval_rank: int | None = None,
        confidence: float | None = None,
    ) -> Citation:
        self.validate_reference(document=document, chunk=chunk, allowed_chunk_ids=allowed_chunk_ids)
        resolved_span = self._resolve_text_span(chunk=chunk, text_span=text_span)
        label = self.build_label(document=document, chunk=chunk)
        citation_id = self._stable_id(
            "cit",
            document.document_id,
            chunk.chunk_id,
            label,
            resolved_span,
            str(retrieval_rank or ""),
        )
        page = chunk.page_start
        return Citation(
            citation_id=citation_id,
            chunk_id=chunk.chunk_id,
            document_id=document.document_id,
            label=label,
            text_span=resolved_span,
            article=chunk.article,
            clause=chunk.clause,
            point=chunk.point,
            page=page,
            source_uri=document.source_uri if document.source_uri else None,
            effective_status=document.effective_status,
            confidence=confidence,
            retrieval_rank=retrieval_rank,
            metadata={
                "document_number": document.document_number or MISSING_METADATA,
                "authority": document.authority or MISSING_METADATA,
                "effective_date": (
                    document.effective_date.isoformat()
                    if document.effective_date is not None
                    else MISSING_METADATA
                ),
                "parse_flags": [flag.value for flag in chunk.parse_flags],
            },
        )

    def build_label(self, *, document: LegalDocument, chunk: LegalChunk) -> str:
        base = document.document_number or document.title or document.document_id
        path_parts = [value for value in [chunk.article, chunk.clause, chunk.point] if value]
        if path_parts:
            return ", ".join([base, *path_parts])

        fallback_parts = [base]
        if document.authority:
            fallback_parts.append(document.authority)
        display_date = document.issued_date or document.effective_date
        if display_date:
            fallback_parts.append(f"ngày {display_date.strftime('%d/%m/%Y')}")
        return ", ".join(fallback_parts)

    def build_debug_label(self, *, document: LegalDocument, chunk: LegalChunk) -> str:
        parts = [
            self._value(document.title, "title"),
            self._value(document.document_number, "document_number"),
            self._value(chunk.article, "article"),
            self._value(chunk.clause, "clause"),
            self._value(chunk.point, "point"),
            f"status:{document.effective_status.value or EffectiveStatus.UNKNOWN.value}",
            self._date_value(document.effective_date, "effective_date"),
            self._value(document.source_uri, "source_uri"),
            f"chunk:{chunk.chunk_id}",
        ]
        return ", ".join(parts)

    def validate_reference(
        self,
        *,
        document: LegalDocument,
        chunk: LegalChunk,
        allowed_chunk_ids: set[str] | None = None,
    ) -> None:
        if chunk.document_id != document.document_id:
            raise CitationReferenceError(
                "Citation chunk does not belong to the supplied document.",
                details={
                    "document_id": document.document_id,
                    "chunk_document_id": chunk.document_id,
                    "chunk_id": chunk.chunk_id,
                },
            )
        if allowed_chunk_ids is not None and chunk.chunk_id not in allowed_chunk_ids:
            raise CitationReferenceError(
                "Citation chunk was not available in the current retrieved source set.",
                details={"chunk_id": chunk.chunk_id},
            )

    @staticmethod
    def _resolve_text_span(*, chunk: LegalChunk, text_span: str | None) -> str:
        if text_span is None:
            return chunk.content
        if text_span not in chunk.content:
            raise CitationReferenceError(
                "Citation text span is not present in the source chunk.",
                details={"chunk_id": chunk.chunk_id},
            )
        return text_span

    @staticmethod
    def _value(value: str | None, field: str) -> str:
        if value:
            return value
        return f"{field}:{MISSING_METADATA}"

    @staticmethod
    def _date_value(value: object | None, field: str) -> str:
        if value is None:
            return f"{field}:{MISSING_METADATA}"
        return f"{field}:{value}"

    @staticmethod
    def _stable_id(prefix: str, *parts: str) -> str:
        digest = hashlib.sha256("\n".join(parts).encode("utf-8")).hexdigest()[:24]
        return f"{prefix}_{digest}"