import os
import gradio as gr
from huggingface_hub import InferenceClient
from langchain_openai import ChatOpenAI
from crewai_tools import PDFSearchTool
from langchain_community.tools.tavily_search import TavilySearchResults
from crewai_tools import tool
from crewai import Crew, Task, Agent
from sentence_transformers import SentenceTransformer
from typing import List, Tuple

# === Hardcoded API Keys ===
GROQ_API_KEY = "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"
TAVILY_API_KEY = "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"

# Set environment variables for API keys
os.environ['GROQ_API_KEY'] = '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'
os.environ['TAVILY_API_KEY'] = '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'

# === Model and Tool Initialization ===
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

llm = ChatOpenAI(
    openai_api_base="https://api.groq.com/openai/v1",
    openai_api_key=GROQ_API_KEY,
    model_name="llama3-70b-8192",
    temperature=0.1,
    max_tokens=1000
)

rag_tool = PDFSearchTool(
    pdf='finance.pdf',
    config=dict(
        llm=dict(
            provider="groq",
            config=dict(
                model="llama3-8b-8192",
            ),
        ),
        embedder=dict(
            provider="huggingface",
            config=dict(
                model="BAAI/bge-small-en-v1.5",
            ),
        ),
    )
)

web_search_tool = TavilySearchResults(k=3, api_key=TAVILY_API_KEY)

# === Tool Definitions ===
@tool
def router_tool(question: str) -> str:
    """Router Function: Decides between web search and vectorstore."""
    return 'web_search'

# === Agent Definitions ===
Router_Agent = Agent(
    role='Router',
    goal='Route user question to a vectorstore or web search',
    backstory=(
        "You are an expert at routing a user question to a vectorstore or web search. "
        "Use the vectorstore for questions on concept related to Retrieval-Augmented Generation. "
        "You do not need to be stringent with the keywords in the question related to these topics. Otherwise, use web-search."
    ),
    verbose=True,
    allow_delegation=False,
    llm=llm,
)

Retriever_Agent = Agent(
    role="Retriever",
    goal="Use the information retrieved from the vectorstore to answer the question",
    backstory=(
        "You are an assistant for question-answering tasks. "
        "Use the information present in the retrieved context to answer the question. "
        "You have to provide a clear concise answer."
    ),
    verbose=True,
    allow_delegation=False,
    llm=llm,
)

Grader_agent = Agent(
    role='Answer Grader',
    goal='Filter out erroneous retrievals',
    backstory=(
        "You are a grader assessing relevance of a retrieved document to a user question. "
        "If the document contains keywords related to the user question, grade it as relevant. "
        "It does not need to be a stringent test. You have to make sure that the answer is relevant to the question."
    ),
    verbose=True,
    allow_delegation=False,
    llm=llm,
)

hallucination_grader = Agent(
    role="Hallucination Grader",
    goal="Filter out hallucination",
    backstory=(
        "You are a hallucination grader assessing whether an answer is grounded in / supported by a set of facts. "
        "Make sure you meticulously review the answer and check if the response provided is in alignment with the question asked."
    ),
    verbose=True,
    allow_delegation=False,
    llm=llm,
)

answer_grader = Agent(
    role="Answer Grader",
    goal="Filter out hallucination from the answer.",
    backstory=(
        "You are a grader assessing whether an answer is useful to resolve a question. "
        "Make sure you meticulously review the answer and check if it makes sense for the question asked. "
        "If the answer is relevant generate a clear and concise response. "
        "If the answer generated is not relevant then perform a websearch using 'web_search_tool'."
    ),
    verbose=True,
    allow_delegation=False,
    llm=llm,
)

# === Task Definitions ===
router_task = Task(
    description=(
        "Analyse the keywords in the question {question}. "
        "Based on the keywords decide whether it is eligible for a vectorstore search or a web search. "
        "Return a single word 'vectorstore' if it is eligible for vectorstore search. "
        "Return a single word 'websearch' if it is eligible for web search. "
        "Do not provide any other preamble or explanation."
    ),
    expected_output=(
        "Give a binary choice 'websearch' or 'vectorstore' based on the question. "
        "Do not provide any other preamble or explanation."
    ),
    agent=Router_Agent,
    tools=[router_tool],
)

retriever_task = Task(
    description=(
        "Based on the response from the router task extract information for the question {question} with the help of the respective tool. "
        "Use the web_search_tool to retrieve information from the web in case the router task output is 'websearch'. "
        "Use the rag_tool to retrieve information from the vectorstore in case the router task output is 'vectorstore'."
    ),
    expected_output=(
        "You should analyse the output of the 'router_task'. "
        "If the response is 'websearch' then use the web_search_tool to retrieve information from the web. "
        "If the response is 'vectorstore' then use the rag_tool to retrieve information from the vectorstore. "
        "Return a clear and concise text as response."
    ),
    agent=Retriever_Agent,
    context=[router_task],
)

grader_task = Task(
    description=(
        "Based on the response from the retriever task for the question {question} evaluate whether the retrieved content is relevant to the question."
    ),
    expected_output=(
        "Binary score 'yes' or 'no' score to indicate whether the document is relevant to the question. "
        "You must answer 'yes' if the response from the 'retriever_task' is in alignment with the question asked. "
        "You must answer 'no' if the response from the 'retriever_task' is not in alignment with the question asked. "
        "Do not provide any preamble or explanations except for 'yes' or 'no'."
    ),
    agent=Grader_agent,
    context=[retriever_task],
)

hallucination_task = Task(
    description=(
        "Based on the response from the grader task for the question {question} evaluate whether the answer is grounded in / supported by a set of facts."
    ),
    expected_output=(
        "Binary score 'yes' or 'no' score to indicate whether the answer is sync with the question asked. "
        "Respond 'yes' if the answer is useful and contains fact about the question asked. "
        "Respond 'no' if the answer is not useful and does not contains fact about the question asked. "
        "Do not provide any preamble or explanations except for 'yes' or 'no'."
    ),
    agent=hallucination_grader,
    context=[grader_task],
)

answer_task = Task(
    description=(
        "Based on the response from the hallucination task for the question {question} evaluate whether the answer is useful to resolve the question. "
        "If the answer is 'yes' return a clear and concise answer. "
        "If the answer is 'no' then perform a 'websearch' and return the response."
    ),
    expected_output=(
        "Return a clear and concise response if the response from 'hallucination_task' is 'yes'. "
        "Perform a web search using 'web_search_tool' and return a clear and concise response only if the response from 'hallucination_task' is 'no'. "
        "Otherwise respond as 'Sorry! unable to find a valid response'."
    ),
    context=[hallucination_task],
    agent=answer_grader,
)

# === Crew Definition ===
rag_crew = Crew(
    agents=[Router_Agent, Retriever_Agent, Grader_agent, hallucination_grader, answer_grader],
    tasks=[router_task, retriever_task, grader_task, hallucination_task, answer_task],
    verbose=True,
)

def respond(
    message: str,
    history: List[Tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    """Main response function for Gradio chat interface."""
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""
    inputs = {"question": message}
    result = rag_crew.kickoff(inputs=inputs)

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

# === Gradio Interface ===
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)

if __name__ == "__main__":
    demo.launch()