# ingestion.py
"""
Document ingestion pipeline — handles PDF, Markdown, and web pages.
"""
import os
from pathlib import Path
from typing import List, Dict, Any
from dataclasses import dataclass, field
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.document_loaders import (
    PyPDFLoader,
    DirectoryLoader,
    WebBaseLoader,
)
from langchain.schema import Document
import chromadb
from chromadb.config import Settings


@dataclass
class IngestionConfig:
    chunk_size: int = 500
    chunk_overlap: int = 50
    embedding_model: str = "text-embedding-3-small"
    collection_name: str = "agentic_rag"
    persist_directory: str = "./chroma_db"


class DocumentIngestor:
    def __init__(self, config: IngestionConfig = None):
        self.config = config or IngestionConfig()
        self.embeddings = OpenAIEmbeddings(model=self.config.embedding_model)
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.config.chunk_size,
            chunk_overlap=self.config.chunk_overlap,
            length_function=len,
            separators=["\n\n", "\n", ". ", " ", ""],
        )
        # Initialize ChromaDB
        self.chroma_client = chromadb.Client(Settings(
            chroma_db_impl="duckdb+parquet",
            persist_directory=self.config.persist_directory,
            anonymized_telemetry=False,
        ))
        self.collection = self.chroma_client.get_or_create_collection(
            name=self.config.collection_name,
            metadata={"hnsw:space": "cosine"},
        )

    def load_documents(self, source_path: str) -> List[Document]:
        """Load documents from various sources."""
        documents = []
        path = Path(source_path)

        if path.is_file():
            if path.suffix == ".pdf":
                loader = PyPDFLoader(str(path))
                documents.extend(loader.load())
            elif path.suffix == ".md":
                loader = DirectoryLoader(
                    str(path.parent),
                    glob=path.name,
                    loader_cls=langchain_community.document_loaders.TextLoader,
                )
                documents.extend(loader.load())

        elif path.is_dir():
            # Load all PDFs in directory
            loader = DirectoryLoader(
                str(path),
                glob="**/*.pdf",
                loader_cls=PyPDFLoader,
                show_progress=True,
            )
            documents.extend(loader.load())

            # Load all markdown files
            loader = DirectoryLoader(
                str(path),
                glob="**/*.md",
                show_progress=True,
            )
            documents.extend(loader.load())

        print(f"Loaded {len(documents)} documents")
        return documents

    def load_web_pages(self, urls: List[str]) -> List[Document]:
        """Load content from web pages."""
        loader = WebBaseLoader(urls)
        documents = loader.load()
        print(f"Loaded {len(documents)} web pages")
        return documents

    def process_documents(self, documents: List[Document]) -> List[Document]:
        """Split documents into chunks and add metadata."""
        chunks = self.text_splitter.split_documents(documents)

        # Add metadata
        for i, chunk in enumerate(chunks):
            chunk.metadata["chunk_id"] = i
            chunk.metadata["chunk_size"] = len(chunk.page_content)

        print(f"Split into {len(chunks)} chunks")
        return chunks

    def index_documents(self, chunks: List[Document]):
        """Index chunks into ChromaDB."""
        batch_size = 100
        for i in range(0, len(chunks), batch_size):
            batch = chunks[i:i + batch_size]
            
            # Generate embeddings
            embeddings = self.embeddings.embed_documents(
                [chunk.page_content for chunk in batch]
            )
            
            # Add to ChromaDB
            self.collection.add(
                embeddings=embeddings,
                documents=[chunk.page_content for chunk in batch],
                metadatas=[chunk.metadata for chunk in batch],
                ids=[f"chunk_{i + j}" for j in range(len(batch))],
            )
            
            print(f"Indexed batch {i // batch_size + 1}")

        print(f"Total chunks indexed: {self.collection.count()}")

    def ingest(self, source_path: str = None, urls: List[str] = None):
        """Full ingestion pipeline."""
        documents = []
        
        if source_path:
            documents.extend(self.load_documents(source_path))
        if urls:
            documents.extend(self.load_web_pages(urls))
        
        if not documents:
            raise ValueError("No documents to ingest")
        
        chunks = self.process_documents(documents)
        self.index_documents(chunks)
        return self.collection.count()

# tools.py
"""
Tool definitions for the Agentic RAG agent.
"""
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.tools import TavilySearchResults
import chromadb
from chromadb.config import Settings
import json


@dataclass
class ToolResult:
    success: bool
    content: str
    metadata: Dict[str, Any] = None
    sources: List[Dict[str, str]] = None


class VectorSearchTool:
    """Search the ChromaDB vector store."""
    
    def __init__(self, collection_name: str = "agentic_rag", 
                 persist_directory: str = "./chroma_db"):
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        self.chroma_client = chromadb.Client(Settings(
            chroma_db_impl="duckdb+parquet",
            persist_directory=persist_directory,
            anonymized_telemetry=False,
        ))
        self.collection = self.chroma_client.get_collection(collection_name)
    
    def search(self, query: str, k: int = 5, 
               filters: Dict[str, Any] = None) -> ToolResult:
        """Search for similar documents."""
        query_embedding = self.embeddings.embed_query(query)
        
        kwargs = {
            "query_embeddings": [query_embedding],
            "n_results": k,
        }
        if filters:
            kwargs["where"] = filters
        
        results = self.collection.query(**kwargs)
        
        if not results["documents"][0]:
            return ToolResult(
                success=False,
                content="No relevant documents found.",
                sources=[]
            )
        
        sources = []
        content_parts = []
        
        for i, (doc, metadata, distance) in enumerate(zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0]
        )):
            source = {
                "rank": i + 1,
                "score": round(1 - distance, 4),  # Convert distance to similarity
                "source": metadata.get("source", "Unknown"),
                "page": metadata.get("page", "N/A"),
            }
            sources.append(source)
            content_parts.append(
                f"[Source {i+1}] (Score: {source['score']})\n{doc}"
            )
        
        return ToolResult(
            success=True,
            content="\n\n---\n\n".join(content_parts),
            metadata={"total_results": len(results["documents"][0])},
            sources=sources,
        )


class WebSearchTool:
    """Search the web using Tavily."""
    
    def __init__(self):
        self.search = TavilySearchResults(max_results=5)
    
    def search_web(self, query: str) -> ToolResult:
        """Search the web."""
        try:
            results = self.search.invoke(query)
            
            content_parts = []
            sources = []
            
            for i, result in enumerate(results):
                content_parts.append(
                    f"[Web {i+1}] {result['title']}\n"
                    f"URL: {result['url']}\n"
                    f"{result['content']}"
                )
                sources.append({
                    "rank": i + 1,
                    "title": result["title"],
                    "url": result["url"],
                })
            
            return ToolResult(
                success=True,
                content="\n\n---\n\n".join(content_parts),
                metadata={"total_results": len(results)},
                sources=sources,
            )
        except Exception as e:
            return ToolResult(
                success=False,
                content=f"Web search failed: {str(e)}",
                sources=[],
            )


class QueryDecompositionTool:
    """Decompose complex queries into sub-queries."""
    
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    def decompose(self, query: str) -> List[str]:
        """Break complex query into sub-queries."""
        prompt = f"""Break this complex question into simpler sub-questions 
that can each be answered independently.

Question: {query}

Return ONLY the sub-questions, one per line. If the question is simple enough, 
return just the original question.

Sub-questions:"""
        
        response = self.llm.invoke(prompt)
        sub_queries = [
            q.strip() for q in response.content.strip().split('\n')
            if q.strip() and not q.strip().startswith(('1.', '2.', '3.', '-', '*'))
        ]
        
        return sub_queries if sub_queries else [query]

# agent.py
"""
Agentic RAG agent using LangGraph state machine.
"""
from typing import TypedDict, Annotated, List, Dict, Any, Optional
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from dataclasses import dataclass
import json


# ── State Definition ──

class AgentState(TypedDict):
    """State that flows through the agent graph."""
    question: str
    sub_questions: List[str]
    current_sub_question: str
    sub_question_index: int
    
    retrieved_chunks: List[Dict[str, Any]]
    all_retrieved: List[Dict[str, Any]]
    
    answer: str
    sub_answers: List[str]
    citations: List[Dict[str, str]]
    
    iteration: int
    max_iterations: int
    is_sufficient: bool
    needs_web_search: bool
    
    messages: Annotated[list, add_messages]


# ── Agent Nodes ──

class AgenticRAGAgent:
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
        self.vector_search = VectorSearchTool()
        self.web_search = WebSearchTool()
        self.query_decomposer = QueryDecompositionTool()
        
        self.graph = self._build_graph()
    
    def _build_graph(self) -> StateGraph:
        """Build the LangGraph state machine."""
        workflow = StateGraph(AgentState)
        
        # Add nodes
        workflow.add_node("classify_query", self.classify_query)
        workflow.add_node("decompose_query", self.decompose_query)
        workflow.add_node("retrieve", self.retrieve)
        workflow.add_node("evaluate_retrieval", self.evaluate_retrieval)
        workflow.add_node("web_search", self.web_search_node)
        workflow.add_node("generate_answer", self.generate_answer)
        workflow.add_node("reflect", self.reflect)
        workflow.add_node("refine_query", self.refine_query)
        
        # Set entry point
        workflow.set_entry_point("classify_query")
        
        # Add edges
        workflow.add_conditional_edges(
            "classify_query",
            self.should_decompose,
            {
                "decompose": "decompose_query",
                "direct": "retrieve",
            }
        )
        
        workflow.add_edge("decompose_query", "retrieve")
        
        workflow.add_conditional_edges(
            "retrieve",
            self.after_retrieve,
            {
                "evaluate": "evaluate_retrieval",
                "web_search": "web_search",
                "generate": "generate_answer",
            }
        )
        
        workflow.add_conditional_edges(
            "evaluate_retrieval",
            self.after_evaluate,
            {
                "sufficient": "generate_answer",
                "insufficient": "refine_query",
                "web_search": "web_search",
            }
        )
        
        workflow.add_edge("web_search", "generate_answer")
        
        workflow.add_conditional_edges(
            "generate_answer",
            self.after_generate,
            {
                "reflect": "reflect",
                "end": END,
            }
        )
        
        workflow.add_conditional_edges(
            "reflect",
            self.after_reflect,
            {
                "retry": "retrieve",
                "refine": "refine_query",
                "end": END,
            }
        )
        
        workflow.add_edge("refine_query", "retrieve")
        
        return workflow.compile()
    
    # ── Node Implementations ──
    
    def classify_query(self, state: AgentState) -> AgentState:
        """Classify query complexity and determine strategy."""
        question = state["question"]
        
        prompt = f"""Analyze this question and determine its complexity.

Question: {question}

Classify as:
- SIMPLE: Factual question, one retrieval should suffice
- COMPLEX: Requires multiple pieces of information, should be decomposed
- OPINION: Asks for opinion/recommendation, may not need retrieval
- CURRENT: Asks about recent events, may need web search

Respond with ONLY the classification (SIMPLE/COMPLEX/OPINION/CURRENT):"""
        
        response = self.llm.invoke(prompt)
        classification = response.content.strip().upper()
        
        state["needs_web_search"] = classification == "CURRENT"
        
        return state
    
    def decompose_query(self, state: AgentState) -> AgentState:
        """Decompose complex query into sub-questions."""
        question = state["question"]
        sub_questions = self.query_decomposer.decompose(question)
        
        state["sub_questions"] = sub_questions
        state["sub_question_index"] = 0
        state["current_sub_question"] = sub_questions[0]
        
        return state
    
    def retrieve(self, state: AgentState) -> AgentState:
        """Retrieve relevant documents."""
        query = state.get("current_sub_question", state["question"])
        
        # Vector search
        vector_results = self.vector_search.search(query, k=5)
        
        if vector_results.success:
            chunks = []
            for content in vector_results.content.split("\n\n---\n\n"):
                chunks.append({
                    "content": content,
                    "source": "knowledge_base"
                })
            state["retrieved_chunks"] = chunks
            state["all_retrieved"] = state.get("all_retrieved", []) + chunks
        
        state["iteration"] = state.get("iteration", 0) + 1
        
        return state
    
    def evaluate_retrieval(self, state: AgentState) -> AgentState:
        """Evaluate if retrieved content is sufficient."""
        chunks = state.get("retrieved_chunks", [])
        
        if not chunks:
            state["is_sufficient"] = False
            return state
        
        context = "\n\n".join([c["content"] for c in chunks])
        
        prompt = f"""Evaluate if the retrieved content is sufficient to answer the question.

Question: {state.get("current_sub_question", state["question"])}

Retrieved Content:
{context[:2000]}

Is this sufficient to answer the question?
- SUFFICIENT: Content directly addresses the question
- INSUFFICIENT: Content is related but not enough
- NO_MATCH: Content is not relevant

Respond with ONLY one of: SUFFICIENT/INSUFFICIENT/NO_MATCH"""
        
        response = self.llm.invoke(prompt)
        evaluation = response.content.strip().upper()
        
        state["is_sufficient"] = evaluation == "SUFFICIENT"
        state["needs_web_search"] = evaluation == "NO_MATCH"
        
        return state
    
    def web_search_node(self, state: AgentState) -> AgentState:
        """Search the web for information."""
        query = state.get("current_sub_question", state["question"])
        
        web_results = self.web_search.search_web(query)
        
        if web_results.success:
            chunks = []
            for content in web_results.content.split("\n\n---\n\n"):
                chunks.append({
                    "content": content,
                    "source": "web_search"
                })
            state["retrieved_chunks"] = chunks
            state["all_retrieved"] = state.get("all_retrieved", []) + chunks
        
        return state
    
    def generate_answer(self, state: AgentState) -> AgentState:
        """Generate answer from retrieved content."""
        chunks = state.get("all_retrieved", state.get("retrieved_chunks", []))
        
        context = "\n\n".join([
            f"[{i+1}] {c['content']}" for i, c in enumerate(chunks)
        ])
        
        prompt = f"""Answer the question based ONLY on the provided context.
If the context doesn't contain enough information, say so.
Always cite your sources using [1], [2], etc.

Question: {state["question"]}

Context:
{context[:4000]}

Answer:"""
        
        messages = [
            SystemMessage(content="You are a helpful research assistant. Always cite sources."),
            HumanMessage(content=prompt),
        ]
        
        response = self.llm.invoke(messages)
        state["answer"] = response.content
        
        # Extract citations
        state["citations"] = [
            {"source_id": str(i+1), "content": c["content"][:100]}
            for i, c in enumerate(chunks)
        ]
        
        return state
    
    def reflect(self, state: AgentState) -> AgentState:
        """Reflect on the answer quality."""
        prompt = f"""Evaluate this answer for accuracy and completeness.

Question: {state["question"]}

Answer: {state["answer"]}

Evaluate:
1. Is the answer accurate based on the context?
2. Is it complete or does it need more information?
3. Are the citations appropriate?

Respond with JSON:
{{
  "accuracy": "HIGH/MEDIUM/LOW",
  "completeness": "COMPLETE/NEEDS_MORE",
  "needs_retry": true/false,
  "suggestion": "brief suggestion if retry needed"
}}"""
        
        response = self.llm.invoke(prompt)
        
        try:
            evaluation = json.loads(response.content)
            state["is_sufficient"] = not evaluation.get("needs_retry", False)
        except json.JSONDecodeError:
            state["is_sufficient"] = True
        
        return state
    
    def refine_query(self, state: AgentState) -> AgentState:
        """Refine the query based on what was found."""
        prompt = f"""Given the original question and what we've already found, 
suggest a refined search query that might find better information.

Original question: {state["question"]}

What we found so far: {state.get("answer", "Nothing yet")}

Provide ONLY the refined query:"""
        
        response = self.llm.invoke(prompt)
        state["current_sub_question"] = response.content.strip()
        
        return state
    
    # ── Edge Decision Functions ──
    
    def should_decompose(self, state: AgentState) -> str:
        if state.get("needs_web_search"):
            return "direct"
        
        question = state["question"]
        # Simple heuristic: if question has multiple clauses, decompose
        if len(question.split()) > 20 or "?" in question[:-1]:
            return "decompose"
        return "direct"
    
    def after_retrieve(self, state: AgentState) -> str:
        if state.get("needs_web_search"):
            return "web_search"
        if state.get("retrieved_chunks"):
            return "evaluate"
        return "web_search"
    
    def after_evaluate(self, state: AgentState) -> str:
        if state.get("is_sufficient"):
            return "sufficient"
        if state.get("needs_web_search"):
            return "web_search"
        return "insufficient"
    
    def after_generate(self, state: AgentState) -> str:
        if state.get("iteration", 0) < state.get("max_iterations", 3):
            return "reflect"
        return "end"
    
    def after_reflect(self, state: AgentState) -> str:
        if state.get("is_sufficient"):
            return "end"
        if state.get("iteration", 0) >= state.get("max_iterations", 3):
            return "end"
        if state.get("needs_web_search"):
            return "refine"
        return "retry"
    
    # ── Main Entry Point ──
    
    def query(self, question: str, max_iterations: int = 3) -> Dict[str, Any]:
        """Query the Agentic RAG system."""
        initial_state = {
            "question": question,
            "sub_questions": [],
            "current_sub_question": question,
            "sub_question_index": 0,
            "retrieved_chunks": [],
            "all_retrieved": [],
            "answer": "",
            "sub_answers": [],
            "citations": [],
            "iteration": 0,
            "max_iterations": max_iterations,
            "is_sufficient": False,
            "needs_web_search": False,
            "messages": [HumanMessage(content=question)],
        }
        
        result = self.graph.invoke(initial_state)
        
        return {
            "question": result["question"],
            "answer": result["answer"],
            "citations": result["citations"],
            "iterations": result["iteration"],
            "sub_questions": result.get("sub_questions", []),
        }

# api.py
"""
FastAPI endpoint for the Agentic RAG agent.
"""
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Any
import uvicorn

app = FastAPI(
    title="Agentic RAG API",
    description="Intelligent research agent with multi-step reasoning",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize agent
agent = AgenticRAGAgent()


class QueryRequest(BaseModel):
    question: str
    max_iterations: int = 3


class QueryResponse(BaseModel):
    question: str
    answer: str
    citations: List[Dict[str, str]]
    iterations: int
    sub_questions: List[str]


@app.post("/query", response_model=QueryResponse)
async def query_agent(request: QueryRequest):
    """Query the Agentic RAG agent."""
    try:
        result = agent.query(
            question=request.question,
            max_iterations=request.max_iterations,
        )
        return QueryResponse(**result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health_check():
    return {"status": "healthy"}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

# ingest.py
"""Run document ingestion."""
from ingestion import DocumentIngestor, IngestionConfig

config = IngestionConfig(
    chunk_size=500,
    chunk_overlap=50,
    persist_directory="./chroma_db",
)

ingestor = DocumentIngestor(config)

# Ingest local documents
count = ingestor.ingest(source_path="./documents")

# Or ingest web pages
count = ingestor.ingest(urls=[
    "https://docs.python.org/3/tutorial/index.html",
    "https://docs.fast.ai/",
])

print(f"Indexed {count} chunks")

# test_agent.py
"""Test the Agentic RAG agent."""
from agent import AgenticRAGAgent

agent = AgenticRAGAgent()

# Test 1: Simple factual question
result = agent.query("What is backpropagation?")
print(f"Answer: {result['answer']}")
print(f"Citations: {len(result['citations'])}")

# Test 2: Complex multi-hop question
result = agent.query(
    "How does the transformer architecture relate to attention mechanisms, "
    "and how did this influence the development of GPT models?"
)
print(f"Sub-questions: {result['sub_questions']}")
print(f"Answer: {result['answer']}")
print(f"Iterations: {result['iterations']}")

# Test 3: Question requiring web search
result = agent.query("What are the latest developments in AI regulation in 2024?")
print(f"Answer: {result['answer']}")