app/services/rag.py

"""
Lógica principal del pipeline RAG:
  - process_query(payload) → QueryResponse
"""

import re
from typing import Any, Dict, List, Set

import app.core.state as state
from app.core.state import embed_query
from app.models.requests import QueryRequest
from app.models.responses import QueryResponse
from app.services.helpers import (
    build_chatbot_fallback_answer,
    collect_metadata_chunks_for_docs,
    contains_any_trigger,
    detect_target_sections,
    extract_explicit_target_doc_ids,
    extract_quoted_titles,
    get_documents_catalog,
    get_topics_catalog_from_metadata,
    list_documents_from_metadata,
    normalize_match_text,
    resolve_doc_ids_from_titles,
    resolve_topic_from_question,
    supplement_retrieval_for_target_docs,
)
from app.services.llm import call_llm


def process_query(payload: QueryRequest) -> QueryResponse:  # noqa: C901 – lógica RAG densa por diseño
    top_k = payload.top_k or state.CONFIG.get("retrieve", {}).get("top_k", 4)
    temperature = payload.temperature
    mode = payload.mode

    query_text = payload.question
    effective_mode = mode
    quoted_titles = extract_quoted_titles(payload.question)
    target_doc_ids_for_filter: Set[str] = set()
    catalog = get_documents_catalog()
    explicit_target_doc_ids = extract_explicit_target_doc_ids(payload.question, catalog)
    if explicit_target_doc_ids:
        target_doc_ids_for_filter = set(explicit_target_doc_ids)

    pregunta_lower = payload.question.lower().strip()

    chatbot_min_top_k = int(state.CONFIG.get("retrieve", {}).get("chatbot_top_k", 16))
    if mode == "chatbot":
        top_k = max(top_k, chatbot_min_top_k)

    # === 0. Triggers: ayuda, smalltalk, saludo ===
    help_triggers = [str(t).lower() for t in state.TRIGGERS.get("help", [])]
    smalltalk_triggers = [str(t).lower() for t in state.TRIGGERS.get("smalltalk", [])]
    greeting_triggers = [str(t).lower() for t in state.TRIGGERS.get("greeting", [])]

    if contains_any_trigger(pregunta_lower, help_triggers):
        return QueryResponse(answer=state.HELP_SCOPE_TEXT, retrieved=[])
    if contains_any_trigger(pregunta_lower, smalltalk_triggers):
        return QueryResponse(answer=state.ABOUT_BOT_TEXT, retrieved=[])
    if contains_any_trigger(pregunta_lower, greeting_triggers):
        return QueryResponse(answer=state.GREETING_TEXT, retrieved=[])

    # === 0.1. Intenciones de resumen / tabla ===
    multi_summary = False
    summary_keywords = ["resumo", "resumen", "summary", "resuma", "resumir"]
    table_keywords = ["tabela", "tabla", "table", "gerar tabela", "gera uma tabela"]
    summary_modes = {"summary", "summary_multi", "summary_sections", "summary_sections_multi"}
    table_modes = {"table", "table_multi"}
    multi_doc_markers = [
        "esses documentos", "esses 5 documentos", "esses cinco documentos",
        "esses 10 documentos", "esses dez documentos",
        "documentos listados", "documentos acima", "documentos mostrados",
        "os documentos", "los documentos", "all documents",
        "todos", "todas", "all", "ambos", "both",
        "os 3", "los 3", "3 documentos", "3 docs",
    ]

    has_summary_intent = contains_any_trigger(pregunta_lower, summary_keywords)
    has_table_intent = contains_any_trigger(pregunta_lower, table_keywords)
    has_multi_doc_reference = contains_any_trigger(pregunta_lower, multi_doc_markers)
    target_sections = detect_target_sections(pregunta_lower)

    if has_summary_intent and has_multi_doc_reference and not state.LAST_LISTED_DOCS:
        return QueryResponse(
            answer=(
                "Não encontrei uma lista recente de documentos no contexto atual. "
                "Por favor, liste os documentos novamente e depois peça o resumo dos documentos listados."
            ),
            retrieved=[],
        )

    if has_table_intent and has_multi_doc_reference:
        docs_scope = state.LAST_LISTED_DOCS if state.LAST_LISTED_DOCS else catalog
        if docs_scope:
            effective_mode = "table_multi"
            target_doc_ids_for_filter = {str(d.get("id")) for d in docs_scope if d.get("id")}
            titulos = [str(d.get("title") or d.get("id")) for d in docs_scope if (d.get("title") or d.get("id"))]
            titulos_str = "; ".join(titulos)
            query_text = (
                "Gere uma tabela comparativa, usando apenas os trechos fornecidos, para os documentos alvo abaixo. "
                "Use colunas objetivas e inclua citações [n] nas células factuais relevantes. "
                f"Documentos alvo ({len(titulos)}): {titulos_str}"
            )
            top_k = max(top_k, max(12, len(docs_scope) * 8))
        else:
            effective_mode = "table"
            top_k = max(top_k, 12)
    elif has_table_intent:
        effective_mode = "table"
        top_k = max(top_k, 12)

    if state.LAST_LISTED_DOCS and has_summary_intent and has_multi_doc_reference and target_sections:
        multi_summary = True
        effective_mode = "summary_sections_multi"
        docs_scope = state.LAST_LISTED_DOCS
        target_doc_ids_for_filter = {str(d.get("id")) for d in docs_scope if d.get("id")}
        titulos = [str(d.get("title") or d.get("id")) for d in docs_scope if (d.get("title") or d.get("id"))]
        titulos_str = "; ".join(titulos)
        sections_text = ", ".join(target_sections)
        query_text = (
            f"Resuma SOMENTE as seções '{sections_text}' de CADA documento listado abaixo, sem misturar conteúdos entre documentos. "
            "Formato obrigatório: um bloco por documento com título. "
            "Para cada seção solicitada, escreva minimo 150 palavras, com mais detalhe técnico. "
            "Use citações [n] ao longo do texto. "
            f"Documentos alvo ({len(titulos)}): {titulos_str}"
        )
        top_k = max(top_k, max(12, len(docs_scope) * 8))

    elif state.LAST_LISTED_DOCS and has_summary_intent and has_multi_doc_reference:
        multi_summary = True
        effective_mode = "summary_multi"
        docs_scope = state.LAST_LISTED_DOCS
        target_doc_ids_for_filter = {str(d.get("id")) for d in docs_scope if d.get("id")}
        titulos_list: List[str] = [
            str(d.get("title") or d.get("id"))
            for d in docs_scope
            if (d.get("title") or d.get("id"))
        ]
        titulos_str = "; ".join(titulos_list)
        query_text = (
            "Com base exclusivamente nos trechos dos documentos abaixo, "
            "elabore um resumo integrado e sintético, destacando temas centrais, "
            "pontos em comum e diferenças relevantes entre eles."
            "IMPORTANTE: use citações no formato [n] ao longo do texto. Documentos: "
            + titulos_str
        )
        top_k = max(top_k, max(12, len(docs_scope) * 6))

    # === Detección de pedido de listado de documentos ===
    keywords_lista = [
        "list", "lista", "listar", "mostrar", "exibir", "quais são",
        "quais sao", "que documentos", "documentos disponíveis",
        "documentos disponiveis", "documentos indexados",
    ]
    doc_words = ["documento", "documentos", "documents"]
    topic_words = ["topic", "topico", "tópico", "tema"]
    topics_catalog = get_topics_catalog_from_metadata()
    requested_topic = resolve_topic_from_question(payload.question, topics_catalog)
    has_topic_reference = contains_any_trigger(pregunta_lower, topic_words) or bool(requested_topic)

    is_list_request = (
        (not multi_summary)
        and (not has_table_intent)
        and contains_any_trigger(pregunta_lower, keywords_lista)
        and contains_any_trigger(pregunta_lower, doc_words)
    )

    requested_n: int | None = None
    num_match = re.search(r"\b(\d{1,3})\b", pregunta_lower)
    if num_match:
        try:
            requested_n = int(num_match.group(1))
        except ValueError:
            requested_n = None

    if is_list_request:
        if has_topic_reference and requested_topic:
            documentos = list_documents_from_metadata(topic=requested_topic)
        else:
            if hasattr(state.RETRIEVER, "list_documents"):
                documentos = state.RETRIEVER.list_documents()
            else:
                documentos = list_documents_from_metadata()

        total = len(documentos)
        limite = max(1, min(requested_n, total)) if requested_n is not None else total
        documentos_mostrados = documentos[:limite]

        if has_topic_reference and requested_topic:
            lista_formatada = "\n".join([
                f"{i+1}. {doc.get('title', doc.get('id', 'Desconhecido'))}"
                for i, doc in enumerate(documentos_mostrados)
            ])
            resposta = (
                f"📚 **Documentos do tópico '{requested_topic}' ({limite} de {total} no total):**\n\n"
                + lista_formatada
            )
        else:
            lista_formatada = "\n".join([
                f"{i+1}. {doc.get('title', doc.get('id', 'Desconhecido'))}"
                for i, doc in enumerate(documentos_mostrados)
            ])
            resposta = (
                f"📚 **Documentos indexados no sistema ({limite} de {total} no total):**\n\n"
                + lista_formatada
            )

        resposta = re.sub(r"\s*</div>\s*$", "", resposta, flags=re.IGNORECASE)
        state.LAST_LISTED_DOCS = documentos_mostrados
        return QueryResponse(answer=resposta, retrieved=[])

    # === Resolución de scope para resúmenes y tablas ===
    is_summary_request = (effective_mode in summary_modes) or has_summary_intent
    is_table_request = (effective_mode in table_modes) or has_table_intent

    if is_summary_request or is_table_request:
        resolved_from_titles = (
            resolve_doc_ids_from_titles(quoted_titles, catalog) if quoted_titles else set()
        )

        if effective_mode in {"summary_sections", "summary_sections_multi"} and target_sections:
            sections_text = ", ".join(target_sections)
            query_text = f"{query_text}\n\nSeções solicitadas: {sections_text}."

        if explicit_target_doc_ids:
            target_doc_ids_for_filter = set(explicit_target_doc_ids)
        elif resolved_from_titles:
            target_doc_ids_for_filter = set(resolved_from_titles)
        elif (
            effective_mode in {"summary_multi", "summary_sections_multi", "table_multi"}
            and state.LAST_LISTED_DOCS
            and not target_doc_ids_for_filter
        ):
            target_doc_ids_for_filter = {
                str(d.get("id")) for d in state.LAST_LISTED_DOCS if d.get("id")
            }

        if is_table_request and quoted_titles and not resolved_from_titles:
            return QueryResponse(
                answer=(
                    "Não há informações suficientes na base. "
                    "Não foi possível localizar exatamente o documento solicitado para gerar a tabela."
                ),
                retrieved=[],
            )

        if target_doc_ids_for_filter:
            top_k = max(top_k, max(12, len(target_doc_ids_for_filter) * 8))
        elif quoted_titles:
            query_text = (
                f"{payload.question}\n\n"
                f"Título alvo (priorizar este documento): {'; '.join(quoted_titles)}"
            )
            top_k = max(top_k, 12)

    # === 1. Embed y búsqueda ===
    index_type = state.CONFIG.get("index", {}).get("type", "faiss").lower()
    normalize = index_type == "faiss"
    q_vec = embed_query(state.EMBED_MODEL, query_text, normalize=normalize)
    retrieved: List[Dict[str, Any]] = state.RETRIEVER.retrieve(q_vec, top_k)

    if len(target_doc_ids_for_filter) == 1 and is_summary_request:
        direct_doc_chunks = collect_metadata_chunks_for_docs(
            metadata=state.METADATA,
            target_doc_ids=target_doc_ids_for_filter,
            max_chunks_per_doc=20,
        )
        if direct_doc_chunks:
            retrieved = direct_doc_chunks

    if target_doc_ids_for_filter:
        scoped = [
            item for item in retrieved
            if str(item.get("document_id", "")) in target_doc_ids_for_filter
        ]
        retrieved = scoped
        chunks_per_doc = 6
        retrieved = supplement_retrieval_for_target_docs(
            retrieved=retrieved,
            target_doc_ids=target_doc_ids_for_filter,
            metadata=state.METADATA,
            chunks_per_doc=chunks_per_doc,
        )

    if quoted_titles and retrieved and not target_doc_ids_for_filter:
        target_titles_norm = [
            normalize_match_text(t) for t in quoted_titles if normalize_match_text(t)
        ]

        def _title_score(item: Dict[str, Any]) -> int:
            title_norm = normalize_match_text(item.get("document_title") or "")
            if not title_norm:
                return 0
            score = 0
            for target in target_titles_norm:
                if title_norm == target:
                    score = max(score, 3)
                elif target in title_norm:
                    score = max(score, 2)
                elif title_norm in target and len(title_norm) >= 20:
                    score = max(score, 1)
            return score

        scored_items = [(item, _title_score(item)) for item in retrieved]
        best_score = max(score for _, score in scored_items)

        if best_score > 0:
            best_item = next(item for item, score in scored_items if score == best_score)
            best_doc_id = best_item.get("document_id")
            if best_doc_id:
                retrieved = [item for item, _ in scored_items if item.get("document_id") == best_doc_id]
            else:
                retrieved = [item for item, score in scored_items if score == best_score]

    # === 2. Asignar citation_id temporal ===
    citation_ids: Dict[str, int] = {}
    next_id = 1
    for m in retrieved:
        doc_id = m.get("document_id")
        if not doc_id:
            continue
        if doc_id not in citation_ids:
            citation_ids[doc_id] = next_id
            next_id += 1
        m["citation_id"] = citation_ids[doc_id]

    # === DEBUG ===
    print("\n[DEBUG] ===== CHUNKS ENVIADOS AL MODELO =====")
    print(f"[DEBUG] question: {payload.question}")
    print(f"[DEBUG] mode: {effective_mode}")
    print(f"[DEBUG] total_chunks: {len(retrieved)}")
    for i, ch in enumerate(retrieved, start=1):
        from app.services.helpers import normalizar_texto
        doc_id = ch.get("document_id", "N/A")
        doc_title = ch.get("document_title", "N/A")
        frag_id = ch.get("fragment_id", ch.get("idx", "N/A"))
        content = normalizar_texto(ch.get("content", ""))
        preview = content[:220] + ("..." if len(content) > 220 else "")
        print(
            f"[DEBUG] #{i} | doc_id={doc_id} | frag={frag_id} | title={doc_title}\n"
            f"        {preview}"
        )
    print("[DEBUG] =====================================\n")

    # === 4. Llamar al LLM ===
    answer = call_llm(query_text, retrieved, temperature, effective_mode)

    # === 5. Detectar tag especial <<LIST_DOCUMENTS>> ===
    if answer.strip() == "<<LIST_DOCUMENTS>>":
        if is_list_request:
            if hasattr(state.RETRIEVER, "list_documents"):
                documentos = state.RETRIEVER.list_documents()
            else:
                documentos = list_documents_from_metadata()
            total = len(documentos)
            limite = max(1, min(requested_n, total)) if requested_n is not None else total
            documentos_mostrados = documentos[:limite]
            lista_formatada = "\n".join([
                f"{i+1}. {doc.get('title', doc.get('id', 'Desconhecido'))}"
                for i, doc in enumerate(documentos_mostrados)
            ])
            resposta = (
                f"📚 **Documentos indexados no sistema ({limite} de {total} no total):**\n\n"
                + lista_formatada
            )
            resposta = re.sub(r"\s*</div>\s*$", "", resposta, flags=re.IGNORECASE)
            state.LAST_LISTED_DOCS = documentos_mostrados
            return QueryResponse(answer=resposta, retrieved=[])
        else:
            answer = state.NO_INFO_PREFIX

    # === 6. Sin información ===
    if answer.strip().startswith(state.NO_INFO_PREFIX) or answer.strip().startswith(
        "Desculpe, mas não tenho informações"
    ):
        if effective_mode == "chatbot" and retrieved:
            stripped = answer.strip()
            if stripped.startswith(state.NO_INFO_PREFIX):
                suffix = stripped[len(state.NO_INFO_PREFIX):].lstrip(" .:-\n\t")
                answer = (suffix[0].upper() + suffix[1:]) if suffix else build_chatbot_fallback_answer(retrieved)
            else:
                answer = build_chatbot_fallback_answer(retrieved)
        else:
            return QueryResponse(answer=answer, retrieved=[])

    if effective_mode == "chatbot" and retrieved and not answer.strip():
        answer = build_chatbot_fallback_answer(retrieved)

    if answer.strip().startswith(state.NO_INFO_PREFIX) or answer.strip().startswith(
        "Desculpe, mas não tenho informações"
    ):
        return QueryResponse(answer=answer, retrieved=[])

    # === 7. Filtro de saludos genéricos ===
    saludo_regex = r"^(ol[áa]|hola|hello|hi)[!,.]?"
    generic_prefixes = [
        "eu sou um assistente",
        "sou um assistente",
        "sou um assistente virtual",
        "eu sou um assistente virtual",
    ]
    if re.match(saludo_regex, answer.strip().lower()) or any(
        answer.strip().lower().startswith(p) for p in generic_prefixes
    ):
        return QueryResponse(answer=answer, retrieved=[])

    # === 8. Renumeración de citas ===
    raw_citations = [int(m.group(1)) for m in re.finditer(r"\[(\d+)\]", answer)]
    if raw_citations:
        unique_old_numbers = list(dict.fromkeys(raw_citations))
        renumber_map = {old: new for new, old in enumerate(unique_old_numbers, start=1)}

        for old, new in renumber_map.items():
            answer = answer.replace(f"[{old}]", f"[TEMP_{new}]")
        answer = answer.replace("TEMP_", "")
        answer = re.sub(r"\[(\d+)\](\s*\[\1\])+", r"[\1]", answer)

        for m in retrieved:
            old_id = m.get("citation_id")
            if old_id in renumber_map:
                m["citation_id"] = renumber_map[old_id]
            else:
                m["citation_id"] = None

        retrieved = [m for m in retrieved if m.get("citation_id") is not None]
    else:
        retrieved = []

    answer = re.sub(r"\s*</div>\s*$", "", answer, flags=re.IGNORECASE)
    return QueryResponse(answer=answer, retrieved=retrieved)