#%%
### Router

from src.index import *

from typing import Literal

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
# from langchain_openai import AzureChatOpenAI

#%%
# Data model
class RouteQuery(BaseModel):
    """Route a user query to the most relevant datasource."""

    datasource: Literal["vectorstore", "web_search"] = Field(
        ...,
        description="Given a user question choose to route it to web search or a vectorstore.",
    )


# LLM with function call
# llm = AzureChatOpenAI(model="gpt-4o-mini", temperature=0.3)
structured_llm_router = llm.with_structured_output(RouteQuery)

# Prompt
system = """You are an expert at routing a user question to a vectorstore or web search.
The vectorstore contains documents related to Indian Penal Code and The Indian Constitution. 
It can answer questions related to Indian Law, IPC and the Constitution.
Use vectorstore if the question is a legal query within the scope of IPC, Indian Law and the Indian Constitution.
Use web-search if the question is a legal query outside the scope of IPC, Indian Law and the Indian Constitution.
Use web-search and your own knowledge if the question requires general legal help.
Use web-search if the questions is a legal query that requires latest information."""
route_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "{question}"),
        ("human", "{chat_context}"),
    ]
)

question_router = route_prompt | structured_llm_router

# %%
### Retrieval Grader

# Data model
class GradeDocuments(BaseModel):
    """Binary score for relevance check on retrieved documents."""

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

# LLM with function call
# llm = AzureChatOpenAI(model="gpt-4o-mini", temperature=0.3)
structured_llm_grader = llm.with_structured_output(GradeDocuments)

# Prompt
system = """You are a grader assessing relevance of a retrieved document to a user question. \n 
    If the document contains keyword(s) or is relevant to the user question, grade it as relevant. \n
    The goal is to filter out erroneous retrievals. \n
    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.
    Return 'yes' if the document is relevant to the question, otherwise return 'no'.
    Also return 'yes' if the document may be relevant, and might be useful, otherwise return 'no'."""
grade_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Retrieved document: \n\n {document} \n\n User question: {question} \n\n Chat context: {chat_context}"),
    ]
)

retrieval_grader = grade_prompt | structured_llm_grader


# question = "agent memory"
# docs = retriever.invoke(question)
# doc_txt = docs[1].page_content
# print(retrieval_grader.invoke({"question": question, "document": doc_txt}))

#%%

from langchain import hub
from langchain_core.output_parsers import StrOutputParser

# Prompt
prompt = hub.pull("rlm/rag-prompt")

# LLM
# llm = AzureChatOpenAI(model_name="gpt-4o-mini", temperature=0.3)



# Post-processing
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


# Chain
rag_chain = prompt | llm | StrOutputParser()

# # Run
# generation = rag_chain.invoke({"context": docs, "question": question})
# print(generation)

#%%

### Hallucination Grader


# Data model
class GradeHallucinations(BaseModel):
    """Binary score for hallucination present in generation answer."""

    binary_score: str = Field(
        description="Answer is grounded in the facts, 'yes' or 'no'"
    )


# LLM with function call
# llm = AzureChatOpenAI(model="gpt-4o-mini", temperature=0.3)
structured_llm_grader = llm.with_structured_output(GradeHallucinations)

# Prompt
system = """You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \n 
     Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts."""
hallucination_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Set of facts: \n\n {documents} \n\n LLM generation: {generation} \n\n Chat context: {chat_context}"),
    ]
)

hallucination_grader = hallucination_prompt | structured_llm_grader
# hallucination_grader.invoke({"documents": docs, "generation": generation})

#%%
### Answer Grader


# Data model
class GradeAnswer(BaseModel):
    """Binary score to assess answer addresses question."""

    binary_score: str = Field(
        description="Answer addresses the question, 'yes' or 'no'"
    )


# LLM with function call
# llm = AzureChatOpenAI(model="gpt-4o-mini", temperature=0.3)
structured_llm_grader = llm.with_structured_output(GradeAnswer)

# Prompt
system = """You are a grader assessing whether an answer addresses / resolves a question \n 
     Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question."""
answer_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "User question: \n\n {question} \n\n LLM generation: {generation} \n\n Chat context: {chat_context}"),
    ]
)

answer_grader = answer_prompt | structured_llm_grader
# answer_grader.invoke({"question": question, "generation": generation})

#%%
### Question Re-writer

# LLM
# # llm = AzureChatOpenAI(model="gpt-4o-mini", temperature=0.3)

# Prompt
system = """You a question re-writer that converts an input question to a better version that is optimized \n 
     for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning."""
re_write_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        (
            "human",
            "Here is the initial question: \n\n {question} \n\n Here is the chat context: \n\n {chat_context} \n. Use it to form a better question. Formulate an improved question.",
        ),
    ]
)

question_rewriter = re_write_prompt | llm | StrOutputParser()
# question_rewriter.invoke({"question": question})

class IntentClassifier(BaseModel):
    """Classify the intent of the user query."""

    intent: Literal["greeting", "legal_query", "follow_up", "off_topic"] = Field(
        ...,
        description="""Classify the intent of the user query.
        'greeting' if the user is saying greetings,
        'legal_query' if the user is asking for information about law,
        'follow_up' if the user is asking for information related to the previous conversation. If you think the user is referring to a previous conversation, you can classify it as 'follow_up'.
        'off_topic' if the user is asking for information about anything else."""
    )

# LLM with function call
# llm = AzureChatOpenAI(model="gpt-4o-mini", temperature=0.3)
structured_llm_intent_classifier = llm.with_structured_output(IntentClassifier)

# Prompt
system = """You are an intent classifier that classifies the intent of a user query. \n 
    Give the intent as one of the following: 'greeting', 'legal_query', 'follow_up', 'off_topic'."""
intent_classifier_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Here is the user query: \n\n {question} \n\n Here is the chat context: \n\n {chat_context} \n\n Classify the intent of the user query."),
    ]
)

intent_classifier = intent_classifier_prompt | structured_llm_intent_classifier