core/citations.py

from __future__ import annotations

from dataclasses import dataclass
import re
import unicodedata
from pathlib import PurePath
from typing import Any, Iterable

MAX_FONTES_RESPOSTA = 6

CABECALHO_FONTE_RE = re.compile(
    r"^\[CABE(?:\u00c7|C|\?)ALHO FONTE:\s*(?P<filename>[^:\]\n]+?\.pdf)\s*:\s*(?P<page>\d+)\s*\]",
    re.IGNORECASE,
)
CITACAO_PDF_RE = re.compile(
    r"\[(?P<filename>[^\]\n]+?\.pdf)\s*:\s*(?P<page>\d+)\]",
    re.IGNORECASE,
)
SMALLTALK_CHARS_RE = re.compile(r"[^a-z0-9\s]+")

SMALLTALK_FRASES = {
    "oi",
    "ola",
    "ol\u00e1",
    "bom dia",
    "boa tarde",
    "boa noite",
    "boas",
    "hey",
    "hello",
    "hi",
    "tudo bem",
    "td bem",
    "obrigado",
    "obrigada",
    "valeu",
}

SMALLTALK_TOKENS = {
    "oi",
    "ola",
    "ol\u00e1",
    "boas",
    "hey",
    "hello",
    "hi",
    "obrigado",
    "obrigada",
    "valeu",
}


@dataclass(frozen=True)
class Citation:
    filename: str
    page: int | None = None
    material_id: str | None = None
    source: str | None = None
    chunk_id: str | None = None

    @property
    def token(self) -> str | None:
        if not self.filename or self.page is None:
            return None
        return f"[{self.filename}:{self.page}]"


def _safe_int(value: Any) -> int | None:
    if value is None:
        return None

    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _clean_filename(value: Any) -> str | None:
    if value is None:
        return None

    filename = str(value).strip()
    if not filename:
        return None

    filename = PurePath(filename).name
    if not filename.lower().endswith(".pdf"):
        return None

    return filename


def _chunk_text(chunk: Any) -> str:
    if isinstance(chunk, str):
        return chunk

    page_content = getattr(chunk, "page_content", None)
    if isinstance(page_content, str):
        return page_content

    if isinstance(chunk, dict):
        content = chunk.get("page_content") or chunk.get("content") or chunk.get("text")
        if isinstance(content, str):
            return content

    return ""


def _chunk_metadata(chunk: Any) -> dict[str, Any]:
    metadata = getattr(chunk, "metadata", None)
    if isinstance(metadata, dict):
        return metadata

    if isinstance(chunk, dict):
        raw_metadata = chunk.get("metadata")
        if isinstance(raw_metadata, dict):
            return raw_metadata

    return {}


def _citation_from_text(text: str) -> Citation | None:
    for line in text.splitlines()[:4]:
        match = CABECALHO_FONTE_RE.search(line.strip())
        if match:
            return Citation(
                filename=match.group("filename").strip(),
                page=_safe_int(match.group("page")),
            )

    match = CITACAO_PDF_RE.search(text)
    if not match:
        return None

    return Citation(
        filename=match.group("filename").strip(),
        page=_safe_int(match.group("page")),
    )


def _citation_from_metadata(metadata: dict[str, Any]) -> Citation | None:
    filename = None
    for key in ("filename", "file_name", "source", "pdf", "material_filename"):
        filename = _clean_filename(metadata.get(key))
        if filename:
            break

    if not filename:
        return None

    page = None
    for key in ("page_number", "pagina", "page"):
        page = _safe_int(metadata.get(key))
        if page is not None:
            if key == "page" and page >= 0:
                page += 1
            break

    return Citation(
        filename=filename,
        page=page,
        material_id=str(metadata.get("material_id")) if metadata.get("material_id") else None,
        source=str(metadata.get("source")) if metadata.get("source") else None,
        chunk_id=str(metadata.get("chunk_id")) if metadata.get("chunk_id") else None,
    )


def extract_citations_from_context(
    chunks: Iterable[Any],
    limit: int = MAX_FONTES_RESPOSTA,
) -> list[Citation]:
    citations: list[Citation] = []
    seen: set[tuple[str, int | None]] = set()

    for chunk in chunks:
        citation = _citation_from_text(_chunk_text(chunk))
        if citation is None:
            citation = _citation_from_metadata(_chunk_metadata(chunk))

        if citation is None or citation.page is None:
            continue

        key = (citation.filename.lower(), citation.page)
        if key in seen:
            continue

        seen.add(key)
        citations.append(citation)

        if len(citations) >= limit:
            break

    return citations


def format_citations(citations: Iterable[Citation]) -> str:
    tokens = [citation.token for citation in citations]
    tokens = [token for token in tokens if token]

    if not tokens:
        return ""

    return "\n\nFontes:\n" + "\n".join(tokens)


def build_missing_citations_block(response_text: str, citations: Iterable[Citation]) -> str:
    existing = {
        match.group(0).replace(" ", "").lower()
        for match in CITACAO_PDF_RE.finditer(response_text or "")
    }

    missing: list[Citation] = []
    for citation in citations:
        token = citation.token
        if token and token.replace(" ", "").lower() not in existing:
            missing.append(citation)

    return format_citations(missing)


def resposta_indica_sem_informacao(response_text: str) -> bool:
    if not response_text:
        return True

    text_lower = response_text.lower()

    def remover_acentos(txt: str) -> str:
        normalized = unicodedata.normalize("NFD", txt)
        return "".join(ch for ch in normalized if unicodedata.category(ch) != "Mn")

    clean_text = remover_acentos(text_lower)

    sem_info_phrases = [
        "nao ha informacoes disponiveis",
        "nao ha informacao disponivel",
        "nao encontrei",
        "nao foi encontrado",
        "nao existe informacao",
        "nao existe informacao suficiente",
        "os materiais nao",
        "nao consta nos materiais",
        "nao consigo encontrar nos materiais",
        "nao foram encontradas informacoes"
    ]

    for phrase in sem_info_phrases:
        if phrase in clean_text:
            return True

    return False


def _normalize_smalltalk(text: str) -> str:
    normalized = unicodedata.normalize("NFD", text or "")
    normalized = "".join(ch for ch in normalized if unicodedata.category(ch) != "Mn")
    normalized = normalized.lower().strip()
    normalized = SMALLTALK_CHARS_RE.sub(" ", normalized)
    return re.sub(r"\s+", " ", normalized).strip()


def is_smalltalk_message(text: str) -> bool:
    normalized = _normalize_smalltalk(text)
    if not normalized or len(normalized) > 40:
        return False

    if normalized in {_normalize_smalltalk(frase) for frase in SMALLTALK_FRASES}:
        return True

    tokens = normalized.split()
    if len(tokens) > 3:
        return False

    normalized_tokens = {_normalize_smalltalk(token) for token in SMALLTALK_TOKENS}
    return any(token in normalized_tokens for token in tokens)