"""
06_rag_pipeline.py — Retrieval-Augmented Generation (RAG) pipeline

Simulates a real-world RAG system: query_parser fans out to 3 parallel
retrievers → reranker merges them → answer_generator → confidence_scorer.

Pipeline:
  query_parser
    ├─ retriever_dense   (parallel, auto-spawned)
    ├─ retriever_sparse  (parallel, auto-spawned)
    └─ retriever_web     (parallel, auto-spawned)
  → reranker  (merge: llm_input=["retriever_dense","retriever_sparse","retriever_web"])
  → answer_generator
  → confidence_scorer

WHAT THE DASHBOARD SHOULD SHOW:
  Live Graph:
    • query_parser → 3 retrievers (fan-out)
    • 3 retrievers → reranker (fan-in / merge point)
    • reranker → answer_generator → confidence_scorer (sequential tail)
    • "Hourglass" or diamond shape: wide in the middle, narrow at ends

  State View:
    • query_parser: INPUT = query string
    • retrievers: each shows chunk list in OUTPUT
    • reranker: INPUT = merged chunks, OUTPUT = top-5 ranked chunks
    • answer_generator: shows full RAG answer
    • confidence_scorer: shows numeric score + reasoning

  Time Travel:
    • Rewind to reranker to see the raw chunks BEFORE ranking
    • Rewind further to see all 3 retrievers' raw results

  Token Dashboard:
    • answer_generator tallest bar (longest prompt: query + chunks)
    • 3 retrievers roughly equal
    • confidence_scorer smallest

  Run Comparison:
    • retriever_web likely slowest (simulated network latency)
    • reranker fast (no LLM call in real RAG — here we use LLM for demo)
"""
import asyncio
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from agentradar import radar, reset_run

load_dotenv()

llm = ChatOpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL"),
    model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
    temperature=0.2,
)


@radar(tag="query_parser")
async def query_parser(query: str):
    """
    Coordinator: auto-spawns 3 parallel retrievers.
    All retrievers receive the same query string.
    Returns None → library returns gathered [dense, sparse, web] results.
    """
    print(f"[query_parser] Dispatching query to 3 retrievers: {query[:60]}")


@radar(tag="retriever_dense", llm_input=["query_parser"])
async def retriever_dense(query: str) -> list:
    """Parallel retriever A: dense vector search (semantic similarity)."""
    await asyncio.sleep(0.3)
    message = llm.invoke(
        f"Simulate 3 semantically similar document chunks for: {query}\n"
        f"Return as a numbered list, each chunk 1-2 sentences."
    )
    chunks = [f"[Dense] {line.strip()}" for line in message.content.split("\n") if line.strip()]
    return chunks[:3]


@radar(tag="retriever_sparse", llm_input=["query_parser"])
async def retriever_sparse(query: str) -> list:
    """Parallel retriever B: sparse keyword search (BM25-style)."""
    await asyncio.sleep(0.2)
    message = llm.invoke(
        f"Simulate 3 keyword-matched document chunks for: {query}\n"
        f"Return as a numbered list, each chunk 1-2 sentences."
    )
    chunks = [f"[Sparse] {line.strip()}" for line in message.content.split("\n") if line.strip()]
    return chunks[:3]


@radar(tag="retriever_web", llm_input=["query_parser"])
async def retriever_web(query: str) -> list:
    """Parallel retriever C: live web search results."""
    await asyncio.sleep(0.5)  # web is slowest
    message = llm.invoke(
        f"Simulate 3 recent web search results (title + snippet) for: {query}\n"
        f"Return as a numbered list."
    )
    chunks = [f"[Web] {line.strip()}" for line in message.content.split("\n") if line.strip()]
    return chunks[:3]


@radar(tag="reranker", llm_input=["retriever_dense", "retriever_sparse", "retriever_web"])
async def reranker(all_chunks: list) -> list:
    """
    Merge point: combine results from all 3 retrievers and pick top 5 chunks.
    all_chunks is a flat list containing results from all 3 retrievers.
    """
    flat = []
    for chunk_group in all_chunks:
        if isinstance(chunk_group, list):
            flat.extend(chunk_group)
        else:
            flat.append(chunk_group)

    chunks_text = "\n".join(flat)
    message = llm.invoke(
        f"From these retrieved chunks, select and rank the top 5 most relevant ones. "
        f"Return only the top 5, numbered:\n{chunks_text}"
    )
    return [line.strip() for line in message.content.split("\n") if line.strip()][:5]


@radar(tag="answer_generator", llm_input=["reranker"])
async def answer_generator(top_chunks: list) -> str:
    """Generate the final RAG answer from top-ranked chunks."""
    context = "\n".join(top_chunks)
    message = llm.invoke(
        f"Using only the context below, answer the question comprehensively:\n\n"
        f"Context:\n{context}"
    )
    return message.content


@radar(tag="confidence_scorer", llm_input=["answer_generator"])
async def confidence_scorer(answer: str) -> dict:
    """Score the answer quality and confidence."""
    message = llm.invoke(
        f"Rate this answer on a scale of 1-10 for: accuracy, completeness, relevance. "
        f"Then give overall confidence (0.0-1.0). Format: JSON.\n\nAnswer: {answer[:500]}"
    )
    return {
        "answer": answer,
        "confidence_raw": message.content,
        "status": "complete",
    }


async def main():
    reset_run()
    query = "What are the latest breakthroughs in protein folding prediction?"

    print("=== RAG Pipeline (Parallel Retrieval + Merge) ===\n")

    print("[1] Parsing query + parallel retrieval (dense + sparse + web)...")
    # query_parser auto-spawns retrievers in parallel, returns [dense, sparse, web] results
    all_chunks = await query_parser(query)
    total = sum(len(c) for c in all_chunks if isinstance(c, list))
    print(f"  → Retrieved {total} chunks total\n")

    print("[2] Reranking...")
    top = await reranker(all_chunks)
    print(f"  → Top {len(top)} chunks selected\n")

    print("[3] Generating answer...")
    answer = await answer_generator(top)
    print(f"  → Answer: {answer[:150]}...\n")

    print("[4] Scoring confidence...")
    scored = await confidence_scorer(answer)
    print(f"  → Confidence: {scored['confidence_raw'][:100]}...")

    print("\n✓ Dashboard: diamond/hourglass graph, 3-way fan-out, reranker merge point")


if __name__ == "__main__":
    asyncio.run(main())
    input("\nPress Enter to view dashboard (http://localhost:7111) or Ctrl+C to exit...")
