"""RAG chain: retrieval, context formatting, prompt, LLM, and response with citations."""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal, Optional, Union

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.language_models import BaseChatModel
from langchain_core.retrievers import BaseRetriever
from langsmith import traceable

from pinrag.config import (
    get_collection_name,
    get_multi_query_count,
    get_retrieve_k,
    get_rerank_retrieve_k,
    get_rerank_top_n,
    get_use_multi_query,
    get_use_rerank,
)
from pinrag.rag.formatting import format_docs, format_sources
from pinrag.rag.multiquery import wrap_retriever_with_multiquery
from pinrag.rag.query_preprocess import preprocess_query
from pinrag.rag.prompts import get_rag_prompt
from pinrag.rag.rerank import is_rerank_available, wrap_retriever_with_cohere_rerank
from pinrag.vectorstore.chroma_client import DEFAULT_PERSIST_DIR
from pinrag.vectorstore.retriever import create_retriever


PathLike = Union[str, Path]
logger = logging.getLogger(__name__)


@dataclass
class RAGResult:
    """Result of running the RAG chain: answer text and source citations."""

    answer: str
    sources: list[dict[str, str | int]]
    documents: list[Document] = field(default_factory=list)


def _docs_to_langsmith_output(docs: list[Document]) -> list[dict]:
    """Convert Documents for LangSmith retriever trace rendering."""
    return [
        {"page_content": d.page_content, "type": "Document", "metadata": d.metadata}
        for d in docs
    ]


@traceable(
    name="retrieve",
    run_type="retriever",
    process_outputs=lambda out: _docs_to_langsmith_output(out) if out else [],
)
def _retrieve(retriever: BaseRetriever, query: str) -> list[Document]:
    """Retrieve documents; LangSmith logs output in retriever format for special rendering."""
    return retriever.invoke(query)


@traceable(name="run_rag", run_type="chain")
def run_rag(
    query: str,
    llm: BaseChatModel,
    *,
    retriever: Optional[BaseRetriever] = None,
    k: Optional[int] = None,
    use_rerank: Optional[bool] = None,
    persist_directory: PathLike = DEFAULT_PERSIST_DIR,
    collection_name: Optional[str] = None,
    embedding: Optional[Embeddings] = None,
    document_id: Optional[str] = None,
    page_min: Optional[int] = None,
    page_max: Optional[int] = None,
    tag: Optional[str] = None,
    document_type: Optional[str] = None,
    file_path: Optional[str] = None,
    response_style: Literal["thorough", "concise"] = "thorough",
) -> RAGResult:
    """Run a 2-step RAG pipeline: retrieve chunks, format context, prompt LLM, return answer with citations.

    Uses LangChain Retriever (store.as_retriever()). Pass a retriever directly for maximum
    flexibility (e.g. wrapped with reranking), or use legacy params to build one from Chroma.

    When PINRAG_USE_MULTI_QUERY=true, generates query variants via LLM, retrieves per
    variant, merges (unique union), then optionally reranks. Improves recall for terse queries.

    Args:
        query: Natural language question to answer.
        llm: Chat model for generation (e.g. from get_chat_model()).
        retriever: Optional BaseRetriever. If provided, used directly; else built from legacy params.
        k: Number of chunks to retrieve. If None, uses PINRAG_RETRIEVE_K (default 10). Ignored when retriever is provided.
        use_rerank: Override config to enable/disable Cohere re-ranking. If None, uses PINRAG_USE_RERANK.
        persist_directory: Chroma persistence directory (used when retriever is None).
        collection_name: Chroma collection name (used when retriever is None). If None, uses provider-based name.
        embedding: Optional embedding model for retrieval (used when retriever is None).
        document_id: Optional document ID to filter retrieval (e.g. PDF file name).
        page_min: Optional start of page range (inclusive). Use with page_max.
        page_max: Optional end of page range (inclusive). Single page: page_min=64, page_max=64.
        tag: Optional tag to filter retrieval (e.g. "PI_PICO").
        document_type: Optional type to filter: "pdf", "youtube", "discord", "github", or "plaintext".
        file_path: Optional file path within a document (GitHub: e.g. src/ria/api/atr.c). Use list_documents to see files.
        response_style: Answer style for prompt instructions ("thorough" or "concise").

    Returns:
        RAGResult with answer (str) and sources (list of {document_id, page}).
    """
    query_for_retrieval = preprocess_query(query)
    built_with_multi_query_no_rerank = False
    truncate_k: Optional[int] = None
    if retriever is None:
        if collection_name is None:
            collection_name = get_collection_name()
        use_rerank = (
            use_rerank if use_rerank is not None else get_use_rerank()
        )
        use_multi_query = get_use_multi_query()

        if use_rerank:
            available, err = is_rerank_available()
            if available:
                base_k = k if k is not None else get_rerank_retrieve_k()
                top_n = get_rerank_top_n()
                base_retriever = create_retriever(
                    k=base_k,
                    persist_directory=persist_directory,
                    collection_name=collection_name,
                    embedding=embedding,
                    document_id=document_id,
                    page_min=page_min,
                    page_max=page_max,
                    tag=tag,
                    document_type=document_type,
                    file_path=file_path,
                )
                if use_multi_query:
                    base_retriever = wrap_retriever_with_multiquery(
                        base_retriever,
                        llm,
                        num_queries=get_multi_query_count(),
                    )
                retriever = wrap_retriever_with_cohere_rerank(
                    base_retriever, top_n=top_n
                )
            else:
                logger.warning(
                    "Re-ranking disabled: %s. Using standard retrieval.", err
                )
                effective_k = k if k is not None else get_retrieve_k()
                base_retriever = create_retriever(
                    k=effective_k,
                    persist_directory=persist_directory,
                    collection_name=collection_name,
                    embedding=embedding,
                    document_id=document_id,
                    page_min=page_min,
                    page_max=page_max,
                    tag=tag,
                    document_type=document_type,
                    file_path=file_path,
                )
                if use_multi_query:
                    retriever = wrap_retriever_with_multiquery(
                        base_retriever,
                        llm,
                        num_queries=get_multi_query_count(),
                    )
                    built_with_multi_query_no_rerank = True
                    truncate_k = effective_k
                else:
                    retriever = base_retriever
        else:
            effective_k = k if k is not None else get_retrieve_k()
            base_retriever = create_retriever(
                k=effective_k,
                persist_directory=persist_directory,
                collection_name=collection_name,
                embedding=embedding,
                document_id=document_id,
                page_min=page_min,
                page_max=page_max,
                tag=tag,
                document_type=document_type,
                file_path=file_path,
            )
            if use_multi_query:
                retriever = wrap_retriever_with_multiquery(
                    base_retriever,
                    llm,
                    num_queries=get_multi_query_count(),
                )
                built_with_multi_query_no_rerank = True
                truncate_k = effective_k
            else:
                retriever = base_retriever
    docs = _retrieve(retriever, query_for_retrieval)

    if built_with_multi_query_no_rerank and truncate_k is not None and len(docs) > truncate_k:
        docs = docs[:truncate_k]

    if not docs:
        return RAGResult(
            answer="No relevant passages found; try a different query or add more documents.",
            sources=[],
            documents=[],
        )

    prompt = get_rag_prompt(response_style)
    messages = prompt.invoke(
        {"context": format_docs(docs), "question": query}
    ).messages
    try:
        response = llm.invoke(messages)
        content = response.content if hasattr(response, "content") else str(response)
        answer = content if isinstance(content, str) else str(content)
        return RAGResult(answer=answer, sources=format_sources(docs), documents=docs)
    except Exception as e:
        logger.exception("LLM invocation failed in run_rag")
        err = str(e).lower()
        if "rate" in err or "limit" in err:
            msg = "Answer generation failed: rate limit exceeded. Please try again later."
        elif "timeout" in err:
            msg = "Answer generation failed: request timed out. Please try again."
        else:
            msg = "Answer generation failed. Please try again."
        return RAGResult(answer=msg, sources=[], documents=docs)
