from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_core.output_parsers import StrOutputParser

# Define documents for indexing
urls = [
    "https://lilianweng.github.io/posts/2023-06-23-agent/",
    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
    "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/"
]

# Load and split documents
docs = [WebBaseLoader(url).load() for url in urls]
docs_list = [item for sublist in docs for item in sublist]
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500)
doc_splits = text_splitter.split_documents(docs_list)

# Store documents in a vector database (Chroma)
vectorstore = Chroma.from_documents(doc_splits, collection_name="adaptive-rag", embedding=OpenAIEmbeddings())
retriever = vectorstore.as_retriever()

from typing import List, Literal, TypedDict
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field


# Define routing model
class RouteQuery(BaseModel):
    datasource: Literal["vectorstore", "web_search"]


route_prompt = ChatPromptTemplate.from_messages([
    ("system", "You are an expert at routing a user question to vectorstore or web search."),
    ("human", "{question}")
])

question_router = route_prompt | ChatOpenAI(model="gpt-3.5-turbo", temperature=0).with_structured_output(RouteQuery)


class GradeDocuments(BaseModel):
    binary_score: str = Field(description="Documents are relevant to the question, 'yes' or 'no'")


grade_prompt = ChatPromptTemplate.from_messages([
    ("system", "Evaluate if the document is relevant to the question. Answer 'yes' or 'no'."),
    ("human", "Document: {document}\nQuestion: {question}")
])

retrieval_grader = grade_prompt | ChatOpenAI(model="gpt-3.5-turbo", temperature=0).with_structured_output(GradeDocuments)

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain.schema import Document

web_search_tool = TavilySearchResults(k=3)


def web_search(state):
    search_results = web_search_tool.invoke({"query": state["question"]})
    web_documents = [Document(page_content=result["content"]) for result in search_results if "content" in result]
    return {"documents": web_documents, "question": state["question"]}


from langgraph.graph import StateGraph, START, END


class GraphState(TypedDict):
    question: str
    generation: str
    documents: List[str]


# Define nodes for query handling
def retrieve(state):
    documents = retriever.invoke(state["question"])
    return {"documents": documents, "question": state["question"]}


def grade_documents(state):
    question = state["question"]
    documents = state["documents"]
    filtered_docs = []
    web_search_needed = "No"

    for doc in documents:
        grade = retrieval_grader.invoke({"question": question, "document": doc.page_content}).binary_score
        if grade == "yes":
            filtered_docs.append(doc)
        else:
            web_search_needed = "Yes"

    return {"documents": filtered_docs, "question": question, "web_search": web_search_needed}


def generate(state):
    prompt_template = """
    Use the following context to answer the question concisely and accurately:
    Question: {question} 
    Context: {context} 
    Answer:
    """

    # Define ChatPromptTemplate using the above template
    rag_prompt = ChatPromptTemplate.from_template(prompt_template)

    # Create the RAG generation chain with LLM and output parsing
    rag_chain = (
            rag_prompt |
            ChatOpenAI(model="gpt-4o-mini", temperature=0) |
            StrOutputParser()
    )
    generation = rag_chain.invoke({"context": state["documents"], "question": state["question"]})
    return {"generation": generation}


# Route question based on source
def route_question(state):
    source = question_router.invoke({"question": state["question"]}).datasource
    return "web_search" if source == "web_search" else "retrieve"


# Compile and Run the Graph
workflow = StateGraph(GraphState)
workflow.add_node("web_search", web_search)
workflow.add_node("retrieve", retrieve)
workflow.add_node("grade_documents", grade_documents)
workflow.add_node("generate", generate)

workflow.add_conditional_edges(START, route_question, {"web_search": "web_search", "retrieve": "retrieve"})
workflow.add_edge("web_search", "generate")
workflow.add_edge("retrieve", "grade_documents")
workflow.add_edge("grade_documents", "generate")
workflow.add_edge("generate", END)

app = workflow.compile()

# Run with example inputs
inputs = {"question": "What are the types of agent memory?"}
for output in app.stream(inputs):
    print(output)
