"""
Node-based Document RAG Architecture

This module mirrors the structure of `askme_open_agent/askme_base/sql_nodes.py`,
but for the file/document RAG workflow.

Pipeline Flow (for @file style queries):
1. RetrievalNode: dense (Chroma) + sparse (BM25) retrieval + fusion
2. AnswerNode: produce an answer/summary from retrieved contexts

This file intentionally centralizes the document-RAG pipeline and supporting
utilities (timer, BM25, fusion, and reingest) so the system can evolve in one place.
"""

from __future__ import annotations

import json
import logging
import os
import pickle
import re
import threading
import time
import argparse
from contextlib import AbstractContextManager
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional, Sequence, Tuple

# -----------------------------
# Phase 0: PERF instrumentation
# -----------------------------

_perf_logger = logging.getLogger("askme.perf")


class timer(AbstractContextManager):
    """
    Context manager that logs elapsed time as:
        [PERF] <label>: <X>ms
    """

    def __init__(self, label: str, logger: Optional[logging.Logger] = None):
        self.label = str(label)
        self._logger = logger or _perf_logger
        self.elapsed_ms: Optional[float] = None
        self._start: Optional[float] = None

    def __enter__(self) -> "timer":
        self._start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        start = self._start or time.perf_counter()
        self.elapsed_ms = (time.perf_counter() - start) * 1000.0
        msg = f"[PERF] {self.label}: {self.elapsed_ms:.1f}ms"
        self._logger.info(msg)
        # Print to stdout so timings are always visible in terminal logs.
        print(msg)
        return None


# -----------------------------
# Phase 2: BM25 sparse retrieval
# -----------------------------

_WORD_RE = re.compile(r"[A-Za-z0-9_]+")


def _tokenize(text: str) -> List[str]:
    return [m.group(0).lower() for m in _WORD_RE.finditer(text or "")]


def _checkpoint_enabled() -> bool:
    from app.utils.bm25_checkpoint import checkpoint_enabled

    return checkpoint_enabled()


def _chroma_path_from_pkl(pkl_path: str) -> str:
    from app.utils.bm25_checkpoint import chroma_path_from_pkl

    return chroma_path_from_pkl(pkl_path)


@dataclass(frozen=True)
class BM25Hit:
    chunk_id: str
    score: float
    document: str
    metadata: Dict


class BM25Index:
    """
    Persistent BM25 index backed by a pickle file.
    Append-only; reset by deleting the pickle. Checkpoint layer in bm25_checkpoint.
    """

    def __init__(
        self,
        path: str,
        *,
        chroma_path: Optional[str] = None,
        openai_api_key: Optional[str] = None,
        skip_load: bool = False,
    ):
        self.path = path
        self._chroma_path = chroma_path or _chroma_path_from_pkl(path)
        self._openai_api_key = openai_api_key
        self._lock = threading.RLock()
        self._chunk_ids: List[str] = []
        self._documents: List[str] = []
        self._metadatas: List[Dict] = []
        self._tokenized_corpus: List[List[str]] = []
        self._bm25 = None
        self._disk_mtime: float = 0.0
        if not skip_load:
            self._load()

    @staticmethod
    def for_chroma_path(chroma_path: str, openai_api_key: Optional[str] = None) -> "BM25Index":
        return BM25Index(
            path=f"{chroma_path}_bm25.pkl",
            chroma_path=chroma_path,
            openai_api_key=openai_api_key,
        )

    def _reset_to_empty(self) -> None:
        self._chunk_ids = []
        self._documents = []
        self._metadatas = []
        self._tokenized_corpus = []
        self._bm25 = None
        self._disk_mtime = 0.0

    def _apply_loaded_data(self, data: Dict[str, Any]) -> None:
        self._chunk_ids = list(data.get("chunk_ids", []))
        self._documents = list(data.get("documents", []))
        self._metadatas = list(data.get("metadatas", []))
        self._tokenized_corpus = list(data.get("tokenized_corpus", []))
        self._rebuild_model()
        self._refresh_disk_mtime()

    def _payload_dict(self) -> Dict[str, Any]:
        return {
            "chunk_ids": self._chunk_ids,
            "documents": self._documents,
            "metadatas": self._metadatas,
            "tokenized_corpus": self._tokenized_corpus,
        }

    def _refresh_disk_mtime(self) -> None:
        if self.path and os.path.exists(self.path):
            try:
                self._disk_mtime = os.path.getmtime(self.path)
            except OSError:
                self._disk_mtime = 0.0
        else:
            self._disk_mtime = 0.0

    def _reload_if_stale(self) -> None:
        if not _checkpoint_enabled() or not self.path:
            return
        try:
            mtime = os.path.getmtime(self.path) if os.path.exists(self.path) else 0.0
        except OSError:
            return
        if mtime > self._disk_mtime:
            self._load()

    def _load(self) -> None:
        from app.utils.bm25_checkpoint import (
            is_pkl_corrupted,
            list_backup_paths,
            load_pkl_dict,
            quarantine_pkl,
            schedule_background_rebuild,
            should_skip_restore,
            atomic_save_pkl,
            pkl_lock,
        )
        from app.core.config import settings

        with self._lock:
            if not self.path:
                return
            if not _checkpoint_enabled():
                if not os.path.exists(self.path):
                    return
                with timer("bm25.load"):
                    with open(self.path, "rb") as f:
                        data = pickle.load(f)
                    self._apply_loaded_data(data)
                return

            if not os.path.exists(self.path):
                self._reset_to_empty()
                return

            corrupt, reason = is_pkl_corrupted(self.path)
            if not corrupt:
                try:
                    with timer("bm25.load"):
                        data = load_pkl_dict(self.path)
                    self._apply_loaded_data(data)
                    print("[BM25] loaded main pickle")
                    return
                except Exception as exc:
                    corrupt = True
                    reason = str(exc)

            if corrupt:
                print(f"[BM25] corrupt main pickle: {reason}")

            if not should_skip_restore():
                for bak in list_backup_paths(self.path):
                    bcorrupt, _ = is_pkl_corrupted(bak)
                    if bcorrupt:
                        quarantine_pkl(bak)
                        continue
                    try:
                        with timer("bm25.load.backup"):
                            data = load_pkl_dict(bak)
                        self._apply_loaded_data(data)
                        with pkl_lock(self.path):
                            atomic_save_pkl(self.path, data)
                        self._refresh_disk_mtime()
                        print(f"[BM25] recovered from backup: {bak}")
                        return
                    except Exception as exc:
                        print(f"[BM25] backup load failed {bak}: {exc}")
                        quarantine_pkl(bak)

            if os.path.exists(self.path):
                quarantine_pkl(self.path)
            self._reset_to_empty()

            api_key = (self._openai_api_key or getattr(settings, "OPENAI_API_KEY", "") or "").strip()
            if getattr(settings, "BM25_AUTO_REBUILD", False) and api_key:
                schedule_background_rebuild(
                    self._chroma_path,
                    api_key,
                    pkl_path=self.path,
                )

    def _save(self) -> None:
        from app.utils.bm25_checkpoint import (
            atomic_save_pkl,
            backup_pkl,
            checkpoint_enabled,
            pkl_lock,
        )

        with self._lock:
            if not self.path:
                return
            payload = self._payload_dict()
            if not _checkpoint_enabled():
                os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True)
                with timer("bm25.save"):
                    with open(self.path, "wb") as f:
                        pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
                self._refresh_disk_mtime()
                return
            with timer("bm25.save"):
                with pkl_lock(self.path):
                    backup_pkl(self.path)
                    atomic_save_pkl(self.path, payload)
            self._refresh_disk_mtime()

    def _rebuild_model(self) -> None:
        with timer(f"bm25.rebuild(n={len(self._tokenized_corpus)})"):
            from rank_bm25 import BM25Okapi

            self._bm25 = BM25Okapi(self._tokenized_corpus) if self._tokenized_corpus else None

    def rebuild_from_chroma(
        self,
        vector_db: Any,
        *,
        batch_size: int = 500,
        model_every: int = 10,
        progress_callback: Optional[Any] = None,
    ) -> int:
        """Bulk rebuild from Chroma; single final save. Returns chunks loaded."""
        with self._lock:
            self._reset_to_empty()
            existing: set = set()
            batch_count = 0
            processed = 0
            total = vector_db.count() if vector_db else 0

            for chunk_id, document, metadata in vector_db.iter_all_chunks(batch_size):
                if chunk_id in existing:
                    continue
                self._chunk_ids.append(str(chunk_id))
                self._documents.append(document or "")
                self._metadatas.append(dict(metadata) if isinstance(metadata, dict) else {})
                self._tokenized_corpus.append(_tokenize(document or ""))
                existing.add(chunk_id)
                processed += 1
                batch_count += 1
                if batch_count >= model_every:
                    self._rebuild_model()
                    batch_count = 0
                if progress_callback:
                    progress_callback(processed, total)

            self._rebuild_model()
            self._save()
            return processed

    def add_chunks(self, chunk_ids: Sequence[str], documents: Sequence[str], metadatas: Optional[Sequence[Dict]] = None) -> None:
        metadatas = metadatas or [{} for _ in chunk_ids]
        if len(chunk_ids) != len(documents) or len(chunk_ids) != len(metadatas):
            raise ValueError("chunk_ids, documents, metadatas must have same length")

        with self._lock:
            self._reload_if_stale()
            existing = set(self._chunk_ids)
            added = 0
            with timer(f"bm25.add_chunks(n={len(chunk_ids)})"):
                for cid, doc, meta in zip(chunk_ids, documents, metadatas):
                    if cid in existing:
                        continue
                    self._chunk_ids.append(str(cid))
                    self._documents.append(doc or "")
                    self._metadatas.append(dict(meta) if isinstance(meta, dict) else {})
                    self._tokenized_corpus.append(_tokenize(doc or ""))
                    existing.add(cid)
                    added += 1
            if added:
                self._rebuild_model()
                self._save()

    def delete_file_chunks(self, file_id: str) -> None:
        """Remove all chunks associated with a specific file_id from the index."""
        with self._lock:
            self._reload_if_stale()
            indices_to_keep = [
                i for i, meta in enumerate(self._metadatas)
                if (meta or {}).get("file_id") != file_id
            ]

            if len(indices_to_keep) == len(self._chunk_ids):
                return

            print(f"📁 sparse: removing {len(self._chunk_ids) - len(indices_to_keep)} chunks for file_id={file_id}")
            self._chunk_ids = [self._chunk_ids[i] for i in indices_to_keep]
            self._documents = [self._documents[i] for i in indices_to_keep]
            self._metadatas = [self._metadatas[i] for i in indices_to_keep]
            self._tokenized_corpus = [self._tokenized_corpus[i] for i in indices_to_keep]

            self._rebuild_model()
            self._save()

    def clear_all(self) -> None:
        """Clear all chunks from the index and delete the backing file."""
        from app.utils.bm25_checkpoint import (
            prune_all_backups,
            set_skip_restore,
        )

        set_skip_restore(True)
        try:
            with self._lock:
                self._reset_to_empty()
                if os.path.exists(self.path):
                    try:
                        os.remove(self.path)
                    except Exception as e:
                        print(f"Error deleting BM25 pickle: {e}")
                prune_all_backups(self.path)
        finally:
            set_skip_restore(False)

    def search(
        self,
        query: str,
        n_results: int = 20,
        *,
        agent_ids: Optional[Sequence[str]] = None,
        folder_path: Optional[str] = None,
    ) -> List[BM25Hit]:
        with self._lock:
            self._reload_if_stale()
            if not self._bm25:
                return []
            tokens = _tokenize(query)
            if not tokens:
                return []
            with timer(f"bm25.search(n_results={n_results})"):
                scores = self._bm25.get_scores(tokens)
            ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
            # Filter by agent/folder while walking the global ranking so the caller
            # gets its own true top-K. The pickle is a single global index across all
            # agents; truncating to n_results first would starve a single agent at scale.
            agent_id_set = {str(a) for a in agent_ids} if agent_ids else None
            limit = max(0, int(n_results))
            hits: List[BM25Hit] = []
            for i in ranked:
                if len(hits) >= limit:
                    break
                s = float(scores[i])
                if s <= 0:
                    continue
                meta = self._metadatas[i] or {}
                if agent_id_set is not None and str(meta.get("agent_id")) not in agent_id_set:
                    continue
                if folder_path and meta.get("folder_path") != folder_path:
                    continue
                hits.append(BM25Hit(chunk_id=self._chunk_ids[i], score=s, document=self._documents[i], metadata=meta))
            return hits


# -----------------------------
# Phase 3: Fusion (RRF)
# -----------------------------

def reciprocal_rank_fusion(ranked_lists: Sequence[Sequence[str]], k: int = 60) -> List[str]:
    scores: Dict[str, float] = {}
    for ranked in ranked_lists:
        for rank, doc_id in enumerate(ranked or []):
            if not doc_id:
                continue
            doc_id = str(doc_id)
            scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank + 1)
    return sorted(scores.keys(), key=lambda d: scores[d], reverse=True)


# -----------------------------
# Phase 3: Cross-encoder reranker
# -----------------------------

@dataclass(frozen=True)
class RerankItem:
    chunk_id: str
    text: str
    metadata: Dict


class CrossEncoderReranker:
    _lock = threading.Lock()
    _model = None
    _model_name = "BAAI/bge-reranker-v2-m3"

    @classmethod
    def _get_model(cls):
        with cls._lock:
            if cls._model is None:
                with timer("reranker.load_model"):
                    from sentence_transformers import CrossEncoder

                    cls._model = CrossEncoder(cls._model_name)
            return cls._model

    @classmethod
    def rerank(cls, query: str, items: Sequence[RerankItem], top_k: int = 5) -> List[Tuple[RerankItem, float]]:
        if not items:
            return []
        pairs = [(query, it.text or "") for it in items]
        model = cls._get_model()
        with timer(f"reranker.score(n={len(pairs)})"):
            scores = model.predict(pairs)
        scored = list(zip(items, [float(s) for s in scores]))
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[: max(0, int(top_k))]


# -----------------------------
# Phase 5: Context compression
# -----------------------------

_SENT_SPLIT = re.compile(r"(?<=[.!?])\s+|\n+")


def compress_text(query: str, text: str, keep_ratio: float = 0.5) -> str:
    query = (query or "").strip()
    text = (text or "").strip()
    if not query or not text:
        return text
    sentences = [s.strip() for s in _SENT_SPLIT.split(text) if s and s.strip()]
    if len(sentences) <= 2:
        return text
    k = max(1, int(len(sentences) * max(0.0, min(1.0, keep_ratio))))
    pairs = [(query, s) for s in sentences]
    model = CrossEncoderReranker._get_model()
    with timer(f"context.compress.score(n={len(pairs)})"):
        scores = model.predict(pairs)
    ranked = sorted(range(len(sentences)), key=lambda i: float(scores[i]), reverse=True)[:k]
    ranked_sorted = sorted(ranked)
    compressed = "\n".join(sentences[i] for i in ranked_sorted)
    return compressed.strip() if compressed.strip() else text


class NodeStatus(Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    SUCCESS = "success"
    FAILED = "failed"


@dataclass
class DocumentRAGConfig:
    dense_k: int = 20
    sparse_k: int = 20
    fused_k: int = 20
    rerank_k: int = 5
    compress_keep_ratio: float = 0.5


@dataclass
class DocumentChunk:
    chunk_id: str
    text: str
    metadata: Dict[str, Any] = field(default_factory=dict)
    distance: Optional[float] = None
    bm25_score: Optional[float] = None
    rerank_score: Optional[float] = None


@dataclass
class DocumentRAGContext:
    query: str
    folder_path: Optional[str] = None
    chunks: List[DocumentChunk] = field(default_factory=list)


@dataclass
class DocumentRAGResult:
    answer: str = ""
    chunks: List[DocumentChunk] = field(default_factory=list)
    errors: List[str] = field(default_factory=list)


@dataclass
class DocumentRAGState:
    query: str
    folder_path: Optional[str] = None
    config: DocumentRAGConfig = field(default_factory=DocumentRAGConfig)
    context: Optional[DocumentRAGContext] = None
    result: Optional[DocumentRAGResult] = None
    trace: Optional[Any] = None
    errors: List[str] = field(default_factory=list)
    node_statuses: Dict[str, NodeStatus] = field(default_factory=dict)

    def add_error(self, msg: str) -> None:
        self.errors.append(msg)


class BaseNode(ABC):
    def __init__(self, document_service: Any):
        self.svc = document_service

    @property
    @abstractmethod
    def name(self) -> str: ...

    @abstractmethod
    def process(self, state: DocumentRAGState) -> DocumentRAGState: ...


class RetrievalNode(BaseNode):
    @property
    def name(self) -> str:
        return "RetrievalNode"

    def process(self, state: DocumentRAGState) -> DocumentRAGState:
        state.node_statuses[self.name] = NodeStatus.IN_PROGRESS
        retrieval_span = state.trace.span(
            name="retrieval",
            input={
                "query": state.query,
                "folder_path": state.folder_path,
                "dense_k": state.config.dense_k,
                "sparse_k": state.config.sparse_k,
                "fused_k": state.config.fused_k,
                "rerank_k": state.config.rerank_k,
            },
        ) if state.trace else None
        try:
            if not getattr(self.svc, "vector_db", None):
                raise RuntimeError("vector_db not available")

            query = state.query
            folder_path = state.folder_path
            cfg = state.config

            # Dense retrieval
            dense_span = retrieval_span.span(name="chroma_dense_search") if retrieval_span else None
            # Filter dense search by agent_id metadata
            current_agent_id = getattr(self.svc, "agent_id", None)
            where_filter = None
            if current_agent_id:
                if isinstance(current_agent_id, list):
                    where_filter = {"agent_id": {"$in": [str(aid) for aid in current_agent_id]}}
                else:
                    where_filter = {"agent_id": str(current_agent_id)}
            
            with timer("doc_nodes.dense"):
                dense = self.svc.vector_db.search(query, n_results=cfg.dense_k, folder_path=folder_path, where=where_filter)
            if dense_span:
                dense_hits = len((dense.get("ids") or [[]])[0]) if isinstance(dense, dict) else 0
                dense_span.end(output={"hits": dense_hits})

            dense_ranked: List[str] = []
            dense_map: Dict[str, DocumentChunk] = {}
            if dense and dense.get("ids") and dense["ids"] and dense["ids"][0]:
                for doc_id, document, metadata, distance in zip(
                    dense["ids"][0],
                    dense.get("documents", [[]])[0],
                    dense.get("metadatas", [[]])[0],
                    dense.get("distances", [[]])[0],
                ):
                    dense_ranked.append(str(doc_id))
                    dense_map[str(doc_id)] = DocumentChunk(
                        chunk_id=str(doc_id),
                        text=document or "",
                        metadata=metadata or {},
                        distance=float(distance) if isinstance(distance, (int, float)) else None,
                    )

            # Sparse retrieval
            sparse_ranked: List[str] = []
            sparse_map: Dict[str, DocumentChunk] = {}
            bm25 = getattr(self.svc, "bm25_index", None)
            if bm25:
                sparse_span = retrieval_span.span(name="bm25_sparse_search") if retrieval_span else None
                # Push agent/folder filtering into the BM25 search itself so the agent's
                # true top-K survives instead of being post-filtered out of a global top-K.
                current_agent_id = getattr(self.svc, "agent_id", None)
                if isinstance(current_agent_id, list):
                    agent_ids = [str(a) for a in current_agent_id if a]
                elif current_agent_id:
                    agent_ids = [str(current_agent_id)]
                else:
                    agent_ids = None
                with timer("doc_nodes.sparse"):
                    hits = bm25.search(
                        query,
                        n_results=cfg.sparse_k,
                        agent_ids=agent_ids,
                        folder_path=folder_path,
                    )

                if sparse_span:
                    sparse_span.end(output={"hits": len(hits)})
                for h in hits:
                    sparse_ranked.append(str(h.chunk_id))
                    sparse_map[str(h.chunk_id)] = DocumentChunk(
                        chunk_id=str(h.chunk_id),
                        text=h.document or "",
                        metadata=h.metadata or {},
                        bm25_score=float(h.score),
                    )

            fused = reciprocal_rank_fusion([dense_ranked, sparse_ranked], k=60)
            candidate_ids = fused[: cfg.fused_k] if fused else dense_ranked[: cfg.fused_k]

            rerank_span = retrieval_span.span(name="cross_encoder_rerank") if retrieval_span else None
            reranked_items: List[Tuple[RerankItem, float]] = []
            rerank_items = [
                RerankItem(
                    chunk_id=cid,
                    text=(dense_map.get(cid) or sparse_map.get(cid) or DocumentChunk(cid, "")).text,
                    metadata=(dense_map.get(cid) or sparse_map.get(cid) or DocumentChunk(cid, "")).metadata or {},
                )
                for cid in candidate_ids
                if (dense_map.get(cid) or sparse_map.get(cid))
            ]
            if cfg.rerank_k and cfg.rerank_k > 0 and rerank_items:
                try:
                    with timer("doc_nodes.rerank"):
                        reranked_items = CrossEncoderReranker.rerank(query, rerank_items, top_k=cfg.rerank_k)
                except Exception as rerank_err:
                    # Reranker model unavailable/failed: keep fusion order so search
                    # degrades to hybrid ranking instead of failing the whole retrieval.
                    print(f"[FILE][rerank] disabled, falling back to fusion order: {rerank_err}")
                    reranked_items = [(it, 0.0) for it in rerank_items[: cfg.rerank_k]]
            else:
                reranked_items = [(it, 0.0) for it in rerank_items]
            if rerank_span:
                rerank_span.end(output={"reranked_chunks": len(reranked_items)})

            chunks: List[DocumentChunk] = []
            for item, score in reranked_items:
                cid = item.chunk_id
                base = dense_map.get(cid) or sparse_map.get(cid)
                if not base:
                    continue
                chunks.append(
                    DocumentChunk(
                        chunk_id=cid,
                        text=base.text or item.text,
                        metadata=base.metadata or item.metadata or {},
                        distance=base.distance,
                        bm25_score=base.bm25_score,
                        rerank_score=float(score) if cfg.rerank_k and cfg.rerank_k > 0 else None,
                    )
                )

            state.context = DocumentRAGContext(query=query, folder_path=folder_path, chunks=chunks)
            state.node_statuses[self.name] = NodeStatus.SUCCESS
            if retrieval_span:
                retrieval_span.end(output={"final_chunks": len(chunks)})
        except Exception as e:
            state.add_error(f"{self.name} failed: {e}")
            state.node_statuses[self.name] = NodeStatus.FAILED
            if retrieval_span:
                retrieval_span.end(output={"status": "failed", "error": str(e)})
        return state


class AnswerNode(BaseNode):
    @property
    def name(self) -> str:
        return "AnswerNode"

    def process(self, state: DocumentRAGState) -> DocumentRAGState:
        state.node_statuses[self.name] = NodeStatus.IN_PROGRESS
        answer_span = state.trace.span(
            name="answer",
            input={"query": state.query},
        ) if state.trace else None
        try:
            if not getattr(self.svc, "ai_analyzer", None):
                raise RuntimeError("ai_analyzer not available")
            if not state.context:
                raise RuntimeError("context not available")

            # Adapt to existing AIAnalyzer.summarize_search_results signature which expects
            # a list[dict] with "filename" and "content" keys.
            enriched: List[Dict[str, Any]] = []
            for ch in state.context.chunks:
                meta = ch.metadata or {}
                file_id = meta.get("file_id")
                file_info = self.svc.db_service.get_file(file_id) if file_id else None
                filename = file_info[2] if file_info and len(file_info) > 2 else meta.get("filename", "Unknown")
                enriched.append(
                    {
                        "filename": filename,
                        "content": ch.text,
                        "chunk_id": ch.chunk_id,
                        "metadata": meta,
                    }
                )

            with timer("doc_nodes.answer"):
                answer = self.svc.ai_analyzer.summarize_search_results(
                    state.context.query,
                    enriched,
                    state.context.folder_path,
                    langfuse_trace=state.trace,
                )

            state.result = DocumentRAGResult(answer=answer, chunks=state.context.chunks, errors=list(state.errors))
            state.node_statuses[self.name] = NodeStatus.SUCCESS
            if answer_span:
                answer_span.end(output={"has_answer": bool(answer), "chunks": len(state.context.chunks)})
        except Exception as e:
            state.add_error(f"{self.name} failed: {e}")
            state.node_statuses[self.name] = NodeStatus.FAILED
            if answer_span:
                answer_span.end(output={"status": "failed", "error": str(e)})
        return state


class DocumentRAGPipeline:
    """
    Orchestrates document RAG retrieval and answering.
    """

    def __init__(self, document_service: Any, config: Optional[DocumentRAGConfig] = None):
        self.svc = document_service
        self.config = config or DocumentRAGConfig()
        self.retrieval = RetrievalNode(document_service)
        self.answer = AnswerNode(document_service)

    def run(self, query: str, folder_path: Optional[str] = None, trace: Optional[Any] = None) -> DocumentRAGState:
        state = DocumentRAGState(query=query, folder_path=folder_path, config=self.config, trace=trace)
        state = self.retrieval.process(state)
        state = self.answer.process(state)
        return state
